def rf_lw152(num_classes, pretrained=True, **kwargs):
model = ResNetLW(Bottleneck, [3, 8, 36, 3], num_classes=num_classes, **kwargs)
if pretrained:
dataset = data_info.get(num_classes, None)
if dataset:
bname = '152_' + dataset.lower()
key = 'rf_lw' + bname
url = models_urls[bname]
model.load_state_dict(maybe_download(key, url), strict=False)
return model
def mbv2(num_classes, pretrained=True, **kwargs):
"""Constructs the network.
Args:
num_classes (int): the number of classes for the segmentation head to output.
"""
model = MBv2(21, **kwargs)
if pretrained:
bname = 'mbv2_voc'
key = 'rf_lw' + bname
url = models_urls[bname]
state = maybe_download(key, url)
model.load_state_dict(state, strict=False)
print('model loaded')
model.segm = conv3x3(256, num_classes, bias=True)
return model
def mbv2(num_classes=2, pretrained=False, **kwargs):
"""Constructs the network.
Args:
num_classes (int): the number of classes for the segmentation head to output.
"""
model = MBv2(num_classes, **kwargs)
if pretrained:
dataset = data_info.get(num_classes, None)
if dataset == 'MHP':
bname = 'mbv2_voc'
key = 'rf_lw' + bname
url = models_urls[bname]
state = maybe_download(key, url)
segm_bias = state.pop('segm.bias')
state['segm.bias'] = torch.tensor([segm_bias[0], segm_bias[15]])
segm_weight = state.pop('segm.weight')
state['segm.weight'] = torch.tensor(
[segm_weight[0].numpy(), segm_weight[15].numpy()])
model.load_state_dict(state, strict=False)
print('model loaded')
return model