Add Facebook Research Semi-Supervised and Semi-Weakly Supervised ResNet model weights.

Author: rwightman

sotabench.py

@@ -167,6 +167,65 @@ def _entry(model_name, paper_model_name, paper_arxiv_id, batch_size=BATCH_SIZE,
     _entry('ig_resnext101_32x48d', 'ResNeXt-101 32x48d (288x288 Mean-Max Pooling)', '1805.00932',
            ttp=True, args=dict(img_size=288), batch_size=BATCH_SIZE // 8),
 
+    ## Facebook SSL weights
+    _entry('ssl_resnet18', 'ResNet-18', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnet50', 'ResNet-50', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext50_32x4d', 'ResNeXt-50 32x4d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext101_32x4d', 'ResNeXt-101 32x4d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext101_32x8d', 'ResNeXt-101 32x8d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext101_32x16d', 'ResNeXt-101 32x16d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+
+    _entry('ssl_resnet50', 'ResNet-50 (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext50_32x4d', 'ResNeXt-50 32x4d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext101_32x4d', 'ResNeXt-101 32x4d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext101_32x8d', 'ResNeXt-101 32x8d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('ssl_resnext101_32x16d', 'ResNeXt-101 32x16d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288), batch_size=BATCH_SIZE // 2,
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+
+    ## Facebook SWSL weights
+    _entry('swsl_resnet18', 'ResNet-18', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnet50', 'ResNet-50', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext50_32x4d', 'ResNeXt-50 32x4d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext101_32x4d', 'ResNeXt-101 32x4d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext101_32x8d', 'ResNeXt-101 32x8d', '1905.00546',
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext101_32x16d', 'ResNeXt-101 32x16d', '1905.00546'),
+
+    _entry('swsl_resnet50', 'ResNet-50 (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext50_32x4d', 'ResNeXt-50 32x4d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext101_32x4d', 'ResNeXt-101 32x4d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext101_32x8d', 'ResNeXt-101 32x8d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288),
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+    _entry('swsl_resnext101_32x16d', 'ResNeXt-101 32x16d (288x288 Mean-Max Pooling)', '1905.00546',
+           ttp=True, args=dict(img_size=288), batch_size=BATCH_SIZE // 2,
+           model_desc='Weights from https://github.com/facebookresearch/semi-supervised-ImageNet1K-models/'),
+
     ## DLA official impl weights (to remove if sotabench added to source)
     _entry('dla34', 'DLA-34', '1707.06484'),
     _entry('dla46_c', 'DLA-46-C', '1707.06484'),

timm/models/helpers.py

@@ -57,30 +57,32 @@ def resume_checkpoint(model, checkpoint_path):
         raise FileNotFoundError()
 
 
-def load_pretrained(model, default_cfg, num_classes=1000, in_chans=3, filter_fn=None):
-    if 'url' not in default_cfg or not default_cfg['url']:
+def load_pretrained(model, cfg=None, num_classes=1000, in_chans=3, filter_fn=None):
+    if cfg is None:
+        cfg = getattr(model, 'default_cfg')
+    if cfg is None or 'url' not in cfg or not cfg['url']:
         logging.warning("Pretrained model URL is invalid, using random initialization.")
         return
 
-    state_dict = model_zoo.load_url(default_cfg['url'], progress=False)
+    state_dict = model_zoo.load_url(cfg['url'], progress=False)
 
     if in_chans == 1:
-        conv1_name = default_cfg['first_conv']
+        conv1_name = cfg['first_conv']
         logging.info('Converting first conv (%s) from 3 to 1 channel' % conv1_name)
         conv1_weight = state_dict[conv1_name + '.weight']
         state_dict[conv1_name + '.weight'] = conv1_weight.sum(dim=1, keepdim=True)
     elif in_chans != 3:
         assert False, "Invalid in_chans for pretrained weights"
 
     strict = True
-    classifier_name = default_cfg['classifier']
-    if num_classes == 1000 and default_cfg['num_classes'] == 1001:
+    classifier_name = cfg['classifier']
+    if num_classes == 1000 and cfg['num_classes'] == 1001:
         # special case for imagenet trained models with extra background class in pretrained weights
         classifier_weight = state_dict[classifier_name + '.weight']
         state_dict[classifier_name + '.weight'] = classifier_weight[1:]
         classifier_bias = state_dict[classifier_name + '.bias']
         state_dict[classifier_name + '.bias'] = classifier_bias[1:]
-    elif num_classes != default_cfg['num_classes']:
+    elif num_classes != cfg['num_classes']:
         # completely discard fully connected for all other differences between pretrained and created model
         del state_dict[classifier_name + '.weight']
         del state_dict[classifier_name + '.bias']