pytorch 加载模型 模型大小测试速度

直接加载整个模型

Pytorch保存和加载整个模型：

save_net=model
if hasattr(model, 'module'):
     save_net=model.module
torch.save(save_net, 'model.pth')
model = torch.load('model.pth')

Pytorch保存和加载预训练模型参数：

save_net=model
if hasattr(model, 'module'):
     save_net=model.module
torch.save(save_net.state_dict(), 'params.pth')
model.load_state_dict(torch.load('params.pth'))

加载部分参数：

import torch
from src import convmlp_s
if __name__ == '__main__':
    model = convmlp_s(num_classes=100)
    state_dict = torch.load("convmlp_s_imagenet.pth")
    new_state_dict = {}
    for k, v in state_dict.items():
        if 'module.' in k:
            new_state_dict[k[7:]] = v
        else:
            if "head