PYTORCH保存训练好的模型

1.保存模型的dict

#...
XXXmodel = model(*args, **kwargs)
torch.save(XXXmodel.state_dict(),'onlyDict.pkl'):

读取dict

pre_model = model(*args, **kwargs)
pre_model.load_state_dict(torch.load('onlyDict.pkl'))

2.保存整个模型（费时间、内存）

#...
XXXmodel = model(*args, **kwargs)
torch.save(XXXmodel,'fullmodel.pkl'):

读取model

pre_model = torch.load('fullmodel.pkl')