mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Add missing dropout for head reset in ConvNeXt default head

Browse Source
This commit is contained in:
Ross Wightman 2022-01-13 22:41:23 -08:00 committed by Ross Wightman
parent b093dcb46d
commit edd3d73695
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
timm/models/convnext.py
View File

@ -1,4 +1,4 @@
""" ConvNext """ ConvNeXt
Paper: `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf Paper: `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf
@ -229,6 +229,7 @@ class ConvNeXt(nn.Module):
('global_pool', SelectAdaptivePool2d(pool_type=global_pool)), ('global_pool', SelectAdaptivePool2d(pool_type=global_pool)),
('norm', norm_layer(self.num_features)), ('norm', norm_layer(self.num_features)),
('flatten', nn.Flatten(1) if global_pool else nn.Identity()), ('flatten', nn.Flatten(1) if global_pool else nn.Identity()),
('drop', nn.Dropout(self.drop_rate)),
('fc', nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()) ('fc', nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity())
])) ]))