File size: 1,244 Bytes
92f0e98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import collections, torch, torchvision, numpy
def load_places_vgg16(weight_file):
model = torchvision.models.vgg16(num_classes=365)
model.features = torch.nn.Sequential(collections.OrderedDict(zip([
'conv1_1', 'relu1_1',
'conv1_2', 'relu1_2',
'pool1',
'conv2_1', 'relu2_1',
'conv2_2', 'relu2_2',
'pool2',
'conv3_1', 'relu3_1',
'conv3_2', 'relu3_2',
'conv3_3', 'relu3_3',
'pool3',
'conv4_1', 'relu4_1',
'conv4_2', 'relu4_2',
'conv4_3', 'relu4_3',
'pool4',
'conv5_1', 'relu5_1',
'conv5_2', 'relu5_2',
'conv5_3', 'relu5_3',
'pool5'],
model.features)))
model.classifier = torch.nn.Sequential(collections.OrderedDict(zip([
'fc6', 'relu6',
'drop6',
'fc7', 'relu7',
'drop7',
'fc8a'],
model.classifier)))
state_dict = torch.load(weight_file)
converted_state_dict = ({
l: torch.from_numpy(numpy.array(v)).view_as(p)
for k, v in state_dict.items()
for l, p in model.named_parameters() if k in l})
model.load_state_dict(converted_state_dict)
# TODO: figure out normalizations etc.
return model
|