Add ResNet-101 pre-trained model

graingert · Jan 12, 2017 · 11821d6 · 11821d6
1 parent 7d150e8
commit 11821d6
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/torchvision/models/resnet.py b/torchvision/models/resnet.py
@@ -11,6 +11,7 @@
     'resnet18': 'https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth',
     'resnet34': 'https://s3.amazonaws.com/pytorch/models/resnet34-333f7ec4.pth',
     'resnet50': 'https://s3.amazonaws.com/pytorch/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://s3.amazonaws.com/pytorch/models/resnet101-5d3b4d8f.pth',
 }
 
 
@@ -171,8 +172,11 @@ def resnet50(pretrained=False):
     return model
 
 
-def resnet101():
-    return ResNet(Bottleneck, [3, 4, 23, 3])
+def resnet101(pretrained=False):
+    model = ResNet(Bottleneck, [3, 4, 23, 3])
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model
 
 
 def resnet152():