直接加载整个模型
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_i