1、保存文件的后缀
有.pt,.pth,.pkl,在格式上没有任何区别,都能用,随便哪种都行。
2、保存和加载整个模型
torch.save(net, 'net_model.pkl')
加载模型
net_loaded = torch.load('net_model.pkl')
3、只保存模型参数
torch.save(net.state_dict(), 'net_parameter.pkl')
加载模型参数
net1.load_state_dict(torch.load('net_parameter.pkl'))