4.5 讀取和存儲 · 動手學深度學習--pytorch

# 4.5 讀取和存儲到目前為止，我們介紹了如何處理數據以及如何構建、訓練和測試深度學習模型。然而在實際中，我們有時需要把訓練好的模型部署到很多不同的設備。在這種情況下，我們可以把內存中訓練好的模型參數存儲在硬盤上供后續讀取使用。 ## 4.5.1 讀寫`Tensor` 我們可以直接使用`save`函數和`load`函數分別存儲和讀取`Tensor`。`save`使用Python的pickle實用程序將對象進行序列化，然后將序列化的對象保存到disk，使用`save`可以保存各種對象,包括模型、張量和字典等。而`laod`使用pickle unpickle工具將pickle的對象文件反序列化為內存。下面的例子創建了`Tensor`變量`x`，并將其存在文件名同為`x.pt`的文件里。 ``` python import torch from torch import nn x = torch.ones(3) torch.save(x, 'x.pt') ``` 然后我們將數據從存儲的文件讀回內存。 ``` python x2 = torch.load('x.pt') x2 ``` 輸出： ``` tensor([1., 1., 1.]) ``` 我們還可以存儲一個`Tensor`列表并讀回內存。 ``` python y = torch.zeros(4) torch.save([x, y], 'xy.pt') xy_list = torch.load('xy.pt') xy_list ``` 輸出： ``` [tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])] ``` 存儲并讀取一個從字符串映射到`Tensor`的字典。 ``` python torch.save({'x': x, 'y': y}, 'xy_dict.pt') xy = torch.load('xy_dict.pt') xy ``` 輸出： ``` {'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])} ``` ## 4.5.2 讀寫模型 ### 4.5.2.1 `state_dict` 在PyTorch中，`Module`的可學習參數(即權重和偏差)，模塊模型包含在參數中(通過`model.parameters()`訪問)。`state_dict`是一個從參數名稱隱射到參數`Tesnor`的字典對象。 ``` python 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() ``` 輸出： ``` OrderedDict([('hidden.weight', tensor([[ 0.2448, 0.1856, -0.5678], [ 0.2030, -0.2073, -0.0104]])), ('hidden.bias', tensor([-0.3117, -0.4232])), ('output.weight', tensor([[-0.4556, 0.4084]])), ('output.bias', tensor([-0.3573]))]) ``` 注意，只有具有可學習參數的層(卷積層、線性層等)才有`state_dict`中的條目。優化器(`optim`)也有一個`state_dict`，其中包含關于優化器狀態以及所使用的超參數的信息。 ``` python optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9) optimizer.state_dict() ``` 輸出： ``` {'param_groups': [{'dampening': 0, 'lr': 0.001, 'momentum': 0.9, 'nesterov': False, 'params': [4736167728, 4736166648, 4736167368, 4736165352], 'weight_decay': 0}], 'state': {}} ``` ### 4.5.2.2 保存和加載模型 PyTorch中保存和加載訓練模型有兩種常見的方法: 1. 僅保存和加載模型參數(`state_dict`)； 2. 保存和加載整個模型。 #### 1. 保存和加載`state_dict`(推薦方式) 保存： ``` python torch.save(model.state_dict(), PATH) # 推薦的文件后綴名是pt或pth ``` 加載： ``` python model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH)) ``` #### 2. 保存和加載整個模型保存： ``` python torch.save(model, PATH) ``` 加載： ``` python model = torch.load(PATH) ``` 我們采用推薦的方法一來實驗一下: ``` python 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 ``` 輸出： ``` tensor([[1], [1]], dtype=torch.uint8) ``` 因為這`net`和`net2`都有同樣的模型參數，那么對同一個輸入`X`的計算結果將會是一樣的。上面的輸出也驗證了這一點。此外，還有一些其他使用場景，例如GPU與CPU之間的模型保存與讀取、使用多塊GPU的模型的存儲等等，使用的時候可以參考[官方文檔](https://pytorch.org/tutorials/beginner/saving_loading_models.html)。 ## 小結 * 通過`save`函數和`load`函數可以很方便地讀寫`Tensor`。 * 通過`save`函數和`load_state_dict`函數可以很方便地讀寫模型的參數。 ----------- > 注：本節與原書此節有一些不同，[原書傳送門](https://zh.d2l.ai/chapter_deep-learning-computation/read-write.html)