PyTorch教程-7：PyTorch中保存與加載tensor和模型詳解
保存和讀取Tensor
PyTorch中的tensor可以保存成 .pt 或者 .pth 格式的文件，使用torch.save()方法保存張量，使用torch.load()來讀取張量：

x = torch.rand(4,5)
torch.save(x, "./myTensor.pt")y = torch.load("./myTensor.pt")
print(y)

tensor([[0.9363, 0.2292, 0.1612, 0.9558, 0.9414],[0.3649, 0.9622, 0.3547, 0.5772, 0.7575],[0.7005, 0.8115, 0.6132, 0.6640, 0.1173],[0.6999, 0.1023, 0.8544, 0.7708, 0.1254]])

當然，save和load方法也適用于其他數據類型，比如list、tuple、dict等：

a = {'a':torch.rand(2,2), 'b':torch.rand(3,4)}
torch.save(a, "./myDict.pth")b = torch.load("./myDict.pth")
print(b)

{'a': tensor([[0.9356, 0.0240],[0.6004, 0.3923]]), 'b': tensor([[0.0222, 0.1799, 0.9172, 0.8159],[0.3749, 0.6689, 0.4796, 0.5772],[0.5016, 0.5279, 0.5109, 0.0592]])}

保存Tensor的純數據
PyTorch中，使用 torch.save 保存的不僅有其中的數據，還包括一些它的信息，包括它與其它數據（可能存在）的關系，這一點是很有趣的。
詳細的原文可以參考：https://pytorch.org/docs/stable/notes/serialization.html#saving-and-loading-tensors-preserves-views

這里結合例子給出一個簡單的解釋。

x = torch.arange(20)
y = x[:5]torch.save([x,y], "./myTensor.pth")
x_, y_ = torch.load("././myTensor.pth")y_ *= 100print(x_)

tensor([  0, 100, 200, 300, 400,   5,   6,   7,   8,   9,  10,  11,  12,  13, 14,  15,  16,  17,  18,  19])

比如在上邊的例子中，我們看到y是x的一個前五位的切片，當我們同時保存x和y后，它們的切片關系也被保存了下來，再將他們加載出來，它們之間依然保留著這個關系，因此可以看到，我們將加載出來的 y_ 乘以100后，x_ 也跟著變化了。

如果不想保留他們的關系，其實也很簡單，再保存y之前使用 clone 方法保存一個只有數據的“克隆體”，這樣就能只保存數據而不保留關系：

x = torch.arange(20)
y = x[:5]torch.save([x,y.clone()], "./myTensor.pth")
x_, y_ = torch.load("././myTensor.pth")y_ *= 100print(x_)

tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19])

當我們只保存y而不同時保存x會怎樣呢？這樣的話確實可以避免如上的情況，即不會再在讀取數據后保留他們的關系，但是實際上有一個不容易被看到的影響存在，那就是保存的數據所占用的空間會和其“父親”級別的數據一樣大：

x = torch.arange(1000)
y = x[:5]torch.save(y, "./myTensor1.pth")
torch.save(y.clone(), "./myTensor2.pth")y1_ = torch.load("./myTensor1.pth")
y2_ = torch.load("./myTensor2.pth")print(y1_.storage().size())
print(y2_.storage().size())

1000
5

如果你去觀察他們保存的文件，會發現占用的空間確實存在很大的差距：

myTensor1.pth      9KB
myTensor2.pth      1KB

綜上所述，對于一些“被關系”的數據來說，如果不想保留他們的關系，最好使用 clone 來保存其“純數據”

保存與加載模型
保存與加載state_dict
這是一種較為推薦的保存方法，即只保存模型的參數，保存的模型文件會較小，而且比較靈活。但是當加載時，需要先實例化一個模型，然后通過加載將參數賦給這個模型的實例，也就是說加載的時候也需要直到模型的結構。

保存：

torch.save(model.state_dict(), PATH)

加載：

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

比較重要的點是：

保存模型時調用 state_dict() 獲取模型的參數，而不保存結構
加載模型時需要預先實例化一個對應的結構
加載模型使用 load_state_dict 方法，其參數不是文件路徑，而是 torch.load(PATH)
如果加載出來的模型用于驗證，不要忘了使用 model.eval() 方法，它會丟棄 dropout、normalization 等層，因為這些層不能在inference的時候使用，否則得到的推斷結果不一致。
一個例子：

import torch
import torch.nn as nn
import torch.nn.functional as Fclass Net(nn.Module):def __init__(self):super(Net,self).__init__()# convolution layersself.conv1 = nn.Conv2d(1,6,3)self.conv2 = nn.Conv2d(6,16,3)# fully-connection layersself.fc1 = nn.Linear(16*6*6,120)self.fc2 = nn.Linear(120,84)self.fc3 = nn.Linear(84,10)def forward(self,x):# max pooling over convolution layersx = F.max_pool2d(F.relu(self.conv1(x)),2)x = F.max_pool2d(F.relu(self.conv2(x)),2)# fully-connected layers followed by activation functionsx = x.view(-1,16*6*6)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))# final fully-connected without activation functonx = self.fc3(x)return xnet = Net()torch.save(net.state_dict(), "./myModel.pth")loaded_net = Net()
loaded_net.load_state_dict(torch.load("./myModel.pth"))
loaded_net.eval()

Net((conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))(conv2): Conv2d(6, 16, kernel_size=(3, 3), stride=(1, 1))(fc1): Linear(in_features=576, out_features=120, bias=True)(fc2): Linear(in_features=120, out_features=84, bias=True)(fc3): Linear(in_features=84, out_features=10, bias=True)
)

保存與加載整個模型
這種方式不僅保存、加載模型的數據，也包括模型的結構一并存儲，存儲的文件會較大，好處是加載時不需要提前知道模型的結構，解來即用。實際上這與上文提到的保存Tensor是一致的。

保存：

torch.save(model, PATH)

加載：

model = torch.load(PATH)
model.eval()

同樣的，如果加載的模型用于inference，則需要使用 model.eval()

保存與加載模型與其他信息
有時我們不僅要保存模型，還要連帶保存一些其他的信息。比如在訓練過程中保存一些 checkpoint，往往除了模型，還要保存它的epoch、loss、optimizer等信息，以便于加載后對這些 checkpoint 繼續訓練等操作；或者再比如，有時候需要將多個模型一起打包保存等。這些其實也很簡單，正如我們上文提到的，torch.save 可以保存dict、list、tuple等多種數據結構，所以一個字典可以很完美的解決這個問題，比如一個簡單的例子：

# saving
torch.save({'epoch': epoch,'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss': loss,...}, PATH)# loading
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']model.eval()
# - or -
model.train()

跨設備存儲與加載
跨設備的情況指對于一些數據的保存、加載在不同的設備上，比如一個在CPU上，一個在GPU上的情況，大致可以分為如下幾種情況：

從CPU保存，加載到CPU
實際上，這就是默認的情況，我們上文提到的所有內容都沒有關心設備的問題，因此也就適應于這種情況。

從CPU保存，加載到GPU
保存：依舊使用默認的方法
加載：有兩種可選的方式
使用 torch.load() 函數的 map_location 參數指定加載后的數據保存的設備
對于加載后的模型使用 to() 函數發送到設備

torch.save(net.state_dict(), PATH)device = torch.device("cuda")loaded_net = Net()
loaded_net.load_state_dict(torch.load(PATH, map_location=device))
# or
loaded_net.to(device)

從GPU保存，加載到CPU
保存：依舊使用默認的方法
加載：只能使用 torch.load() 函數的 map_location 參數指定加載后的數據保存的設備

torch.save(net.state_dict(), PATH)device = torch.device("cuda")loaded_net = Net()
loaded_net.load_state_dict(torch.load(PATH, map_location=device))

從GPU保存，加載到GPU
保存：依舊使用默認的方法
加載：只能使用 對于加載后的模型進行 to() 函數發送到設備

torch.save(net.state_dict(), PATH)device = torch.device("cuda")loaded_net = Net()
loaded_net.to(device)

總結

以上是生活随笔為你收集整理的tensor和模型 保存与加载 PyTorch的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。