Questions about load pretrained model

vision
2

Thanks to Alpha’s Topic, I got an answer.

vgg_torch = torchvision.models.vgg16(pretrained = True)

vgg_my = models.vgg16(pretrained = False)
vgg_my_dict = vgg_my.state_dict()

for k,v in vgg_torch.state_dict().items()[:26]:
    if 0<=int(k.split('.')[1])<=2:
        my_key = 'stage1.'+str(k.split('.')[1])+'.'+str(k.split('.')[2])
        vgg_my_dict[my_key] = v
    if 5<=int(k.split('.')[1])<=7:
        my_key = 'stage2.'+str(int(k.split('.')[1])-5)+'.'+str(k.split('.')[2])
        vgg_my_dict[my_key] = v
    if 10<=int(k.split('.')[1])<=14:
        my_key = 'stage3.'+str(int(k.split('.')[1])-10)+'.'+str(k.split('.')[2])
        vgg_my_dict[my_key] = v
    if 17<=int(k.split('.')[1])<=21:
        my_key = 'stage4.'+str(int(k.split('.')[1])-17)+'.'+str(k.split('.')[2])
        vgg_my_dict[my_key] = v
    if 24<=int(k.split('.')[1])<=28:
        my_key = 'stage5.'+str(int(k.split('.')[1])-24)+'.'+str(k.split('.')[2])
        vgg_my_dict[my_key] = v
    else :
        continue

vgg_my.load_state_dict(vgg_my_dict)