PyTorch——读取与存储

参考链接

1. https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter04_DL_computation/4.5_read-write

读写Tensor

存储

save使用Python的pickle实用程序将对象进行序列化，然后将序列化的对象保存到disk，使用save可以保存各种对象,包括模型、张量和字典等。

torch.save(x, 'x.pt')

读取

load使用pickle unpickle工具将pickle的对象文件反序列化为内存。

x2 = torch.load('x.pt')

读写模型

保存和加载模型参数(state_dict)

在PyTorch中，Module的可学习参数(即权重和偏差)，模块模型包含在参数中(通过model.parameters()访问)。state_dict是一个从参数名称隐射到参数Tesnor的字典对象。

注意，只有具有可学习参数的层(卷积层、线性层等)才有state_dict中的条目。优化器(optim)也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。

torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth

model.load_state_dict(torch.load(PATH))

保存和加载整个模型

torch.save(model, PATH)

model = torch.load(PATH)