tensor和模型 保存与加载 PyTorch

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_ *= 100

print(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_ *= 100

print(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 F

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()

        # convolution layers
        self.conv1 = nn.Conv2d(1,6,3)
        self.conv2 = nn.Conv2d(6,16,3)

        # fully-connection layers
        self.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 layers
        x = 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 functions
        x = x.view(-1,16*6*6)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))

        # final fully-connected without activation functon
        x = self.fc3(x)

        return x

net = 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)