模型加载预训练权重

将预训练好的模型加载到新的网络上。

import torch
from vit_pytorch import ViT
model = ViT(
    image_size = 256,
    patch_size = 8,
    num_classes = 33,
    dim = 256,
    depth = 6,
    heads = 4,
    mlp_dim = 256,
    dropout = 0.1,
    emb_dropout = 0.1
)
#params=model.state_dict() #获得模型的原始状态以及参数
#for k,v in params.items():
    #print(k) #只打印key值，不打印具体参数。
# 随机化一个图像输入
#img = torch.randn(1, 3, 256, 256)
# 获取输出
#preds = model(img) # (1, 1000)
#print(preds.shape)
dir = '/home/best.pth'
model.load_state_dict(torch.load(dir))
for k,v in model.items():
    print(k) #只打印key值，不打印具体参数。