Skip to content

Commit 610f13f

Browse files
committed
Add pretrained model weights.
1 parent 359d941 commit 610f13f

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

torchvision/models/segmentation/segmentation.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,12 @@
1414
model_urls = {
1515
'fcn_resnet50_coco': 'https://download.pytorch.org/models/fcn_resnet50_coco-1167a1af.pth',
1616
'fcn_resnet101_coco': 'https://download.pytorch.org/models/fcn_resnet101_coco-7ecb50ca.pth',
17-
'fcn_mobilenet_v3_large_coco': None,
17+
'fcn_mobilenet_v3_large_coco': 'https://download.pytorch.org/models/fcn_mobilenet_v3_large-23a17f82.pth',
1818
'deeplabv3_resnet50_coco': 'https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth',
1919
'deeplabv3_resnet101_coco': 'https://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pth',
20-
'deeplabv3_mobilenet_v3_large_coco': None,
21-
'lraspp_mobilenet_v3_large_coco': None,
20+
'deeplabv3_mobilenet_v3_large_coco':
21+
'https://download.pytorch.org/models/deeplabv3_mobilenet_v3_large-924f767e.pth',
22+
'lraspp_mobilenet_v3_large_coco': 'https://download.pytorch.org/models/lraspp_mobilenet_v3_large-d234d4ea.pth',
2223
}
2324

2425

@@ -77,7 +78,7 @@ def _load_model(arch_type, backbone, pretrained, progress, num_classes, aux_loss
7778

7879
def _load_weights(model, arch_type, backbone, progress):
7980
arch = arch_type + '_' + backbone + '_coco'
80-
model_url = model_urls[arch]
81+
model_url = model_urls.get(arch, None)
8182
if model_url is None:
8283
raise NotImplementedError('pretrained {} is not supported as of now'.format(arch))
8384
else:

0 commit comments

Comments
 (0)