PyTorch笔记27--模型保存与加载

序列化与反序列化

#序列化
torch.save
#主要参数：
#obj：对象
#f：输出路径

#反序列化
torch.load
#主要参数
#f：文件路径
#map_location：指定存放位置, cpu or gpu

模型保存与加载的两种方式

法1：保存整个Module

torch.save(net, path)

法2：保存模型参数

state_dict = net.state_dict()
torch.save(state_dict , path)

模型断点续训练

checkpoint = {
"model_state_dict": net.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"epoch": epoch
}