Added ability to load pretrained model with custom number of classes

Author: lukemelas

efficientnet_pytorch/model.py

@@ -192,7 +192,7 @@ def from_name(cls, model_name, override_params=None):
 
     @classmethod
     def from_pretrained(cls, model_name, num_classes=1000):
-        model = EfficientNet.from_name(model_name, override_params={'num_classes': 1000})
+        model = EfficientNet.from_name(model_name, override_params={'num_classes': num_classes})
         load_pretrained_weights(model, model_name, load_fc=(num_classes == 1000))
         return model
 

efficientnet_pytorch/utils.py

@@ -279,5 +279,11 @@ def get_model_params(model_name, override_params):
 def load_pretrained_weights(model, model_name, load_fc=True):
     """ Loads pretrained weights, and downloads if loading for the first time. """
     state_dict = model_zoo.load_url(url_map[model_name])
-    model.load_state_dict(state_dict)
+    if load_fc:
+        model.load_state_dict(state_dict)
+    else:
+        state_dict.pop('_fc.weight')
+        state_dict.pop('_fc.bias')
+        res = model.load_state_dict(state_dict, strict=False)
+        assert str(res.missing_keys) == str(['_fc.weight', '_fc.bias']), 'issue loading pretrained weights'
     print('Loaded pretrained weights for {}'.format(model_name))