參考

4.5 讀取和存儲(chǔ)

到目前為止,我們介紹了如何處理數(shù)據(jù)以及如何構(gòu)建、訓(xùn)練和測試深度學(xué)習(xí)模型。然而在實(shí)際中,我們有時(shí)需要把訓(xùn)練好的模型部署到很多不同的設(shè)備。在這種情況下,我們可以把內(nèi)存中訓(xùn)練好的模型參數(shù)存儲(chǔ)在硬盤上供后續(xù)讀取使用。

4.5.1 讀寫tensor

我們可以直接使用save函數(shù)和load函數(shù)分別存儲(chǔ)和讀取Tensor。

下面的例子創(chuàng)建了Tensor變量x,并將其存儲(chǔ)在文件名為x.pt的文件里.

import torch import torch.nn as nnx = torch.ones(3) torch.save(x, 'x.pt')

然后我們將數(shù)據(jù)從存儲(chǔ)的文件讀回內(nèi)存

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

存儲(chǔ)一個(gè)Tensor列表并返回

y =torch.zeros(4) torch.save([x, y], 'xy.pt') xy_list = torch.load('xy.pt') xy_list

存儲(chǔ)并讀取一個(gè)從字符串映射到Tensor的字典

torch.save({'x': x,'y': y }, 'xy_dict.pt') xy = torch.load('xy_dict.pt') xy

4.5.2 讀寫模型

4.5.2.1 state_dict

static_dict是一個(gè)從參數(shù)名稱映射到參數(shù)Tensor的字典對(duì)象

class MLP(nn.Module):def __init__(self):super(MLP, self).__init__()self.hidden = nn.Linear(3, 2)self.act = nn.ReLU()self.output = nn.Linear(2, 1)def forward(self, x):a = self.act(self.hidden(x))return self.output(a)net = MLP() net.state_dict()

注意,只有具有可學(xué)習(xí)參數(shù)的層(卷積層、線性層)才有 state_dict中的條目

optimizer = torch.optim.SGD(net.parameters(), lr= 0.001, momentum=0.9) optimizer.state_dict()

4.5.2.2 保存和加載模型

PyTorch中保存和加載訓(xùn)練模型有兩種常見的方法:

僅保存和加載模型參數(shù)(state_dict)

保存和加載整個(gè)模型。

1. 保存加載static_dict(推薦方式)
torch.save(model.state_dict(), PATH)

# 保存 torch.save(model.state_dict(), PATH)# 加載 model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.laod(PATH))

2. 保存和加載整個(gè)模型

# 保存 torch.save(model, PATH)# 加載 model = torch.load(PATH)

采用第一種方法來試驗(yàn)一下:

X = torch.randn(2, 3) Y = net(X)PATH = "./net.pt" torch.save(net.state_dict(), PATH)net2 = MLP() net2.load_state_dict(torch.load(PATH)) Y2 = net2(X)Y2 ==Y

總結(jié)

以上是生活随笔為你收集整理的[pytorch、学习] - 4.5 读取和存储的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。