深度学习+pytorch自学笔记（七）——模型构造、模型参数的访问、初始化和共享、自定义层、读取和存储、GPU计算

参考书籍《动手学深度学习(pytorch版），参考网址为：

https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter04_DL_computation/4.1_model-construction

https://zh-v2.d2l.ai/chapter_deep-learning-computation/model-construction.html

请大家也多多支持这两个很好用的平台~

大部分内容为书中内容，也有部分自己实验和添加的内容，如涉及侵权，会进行删除。

一、模型构造

回顾一下在多层感知机的简洁实现中含单隐藏层的多层感知机的实现方法。我们首先构造Sequential实例，然后依次添加两个全连接层。其中第一层的输出大小为256，即隐藏层单元个数是256；第二层的输出大小为10，即输出层单元个数是10。我们在上一章的其他节中也使用了Sequential类构造模型。这里我们介绍另外一种基于Module类的模型构造方法：它让模型构造更加灵活。

1.1 继承Module类来构造模型

Module类是nn模块里提供的一个模型构造类，是所有神经网络模块的基类，我们可以继承它来定义我们想要的模型。下面继承Module类构造本节开头提到的多层感知机。这里定义的MLP类重载了Module类的__init__函数和forward函数。它们分别用于创建模型参数和定义前向计算。前向计算也即正向传播。

demo1:

import torch
from torch import nn

class MLP(nn.Module):
    # 声明带有模型参数的层，这里声明了两个全连接层
    def __init__(self, **kwargs):
        # 调用MLP父类Module的构造函数来进行必要的初始化。这样在构造实例时还可以指定其他函数
        # 参数，如“模型参数的访问、初始化和共享”一节将介绍的模型参数params
        super(MLP, self).__init__(**kwargs)
        self.hidden = nn.Linear(784, 256) # 隐藏层
        self.act = nn.ReLU()
        self.output = nn.Linear(256, 10)  # 输出层


    # 定义模型的前向计算，即如何根据输入x计算返回所需要的模型输出
    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)

X = torch.rand(2, 784)
net = MLP()  # 创建一个实例
print(net)
print(net(X))

out1:

MLP(
  (hidden): Linear(in_features=784, out_features=256, bias=True)
  (act): ReLU()
  (output): Linear(in_features=256, out_features=10, bias=True)
)
tensor([[-0.1412,  0.3369, -0.3535, -0.1426, -0.0124,  0.2102,  0.0440, -0.0997,
          0.0469,  0.0598],
        [-0.1235,  0.3672, -0.3056, -0.1914,  0.0456,  0.1594,  0.0753, -0.0011,
          0.1199,  0.0562]], grad_fn=<AddmmBackward>)

以上的MLP类中无须定义反向传播函数。系统将通过自动求梯度而自动生成反向传播所需的backward函数。

我们可以实例化MLP类得到模型变量net。初始化net并传入输入数据X做一次前向计算。其中，net(X)会调用MLP继承自Module类的__call__函数，这个函数将调用MLP类定义的forward函数来完成前向计算。

注意，这里并没有将Module类命名为Layer（层）或者Model（模型）之类的名字，这是因为该类是一个可供自由组建的部件。它的子类既可以是一个层（如PyTorch提供的Linear类），又可以是一个模型（如这里定义的MLP类），或者是模型的一个部分。我们下面通过两个例子来展示它的灵活性。

super( test, self).__init__()

首先找到test的父类（比如是类A），然后把类test的对象self转换为类A的对象，然后“被转换”的类A对象调用自己的__init__函数。

其中，super()用法参考链接：https://blog.csdn.net/a__int__/article/details/104600972

1.2 Module的子类

我们刚刚提到，Module类是一个通用的部件。事实上，PyTorch还实现了继承自Module的可以方便构建模型的类: 如Sequential、ModuleList和ModuleDict等等。

1.2.1 Sequential类

当模型的前向计算为简单串联各个层的计算时，Sequential类可以通过更加简单的方式定义模型。这正是Sequential类的目的：它可以接收一个子模块的有序字典（OrderedDict）或者一系列子模块作为参数来逐一添加Module的实例，而模型的前向计算就是将这些实例按添加的顺序逐一计算。

下面我们实现一个与Sequential类有相同功能的MySequential类。这或许可以帮助读者更加清晰地理解Sequential类的工作机制。

demo2:

# Sequential类
from collections import OrderedDict  # 保持Key的顺序
class MySequential(nn.Module):
    def __init__(self, *args):
        super(MySequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict): # 如果传入的是一个OrderedDict
            for key, module in args[0].items():
                self.add_module(key, module)  # add_module方法会将module添加进self._modules(一个OrderedDict)
        else:  # 传入的是一些Module
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)
    def forward(self, input):
        # self._modules返回一个 OrderedDict，保证会按照成员添加时的顺序遍历成员
        for module in self._modules.values():
            input = module(input)
        return input

net = MySequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
        )
print(net)
print(net(X))

out2:

MySequential(
  (0): Linear(in_features=784, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)
tensor([[ 0.1418,  0.0071, -0.0182,  0.0574, -0.2706, -0.1830,  0.0574,  0.1369,
          0.2271, -0.0421],
        [-0.0328, -0.0201,  0.1502,  0.1099, -0.0528, -0.2136,  0.0483,  0.0781,
          0.3694, -0.1114]], grad_fn=<AddmmBackward>)

其中，collections模块中OrderedDict的使用方法参考链接：

https://blog.csdn.net/brucewong0516/article/details/83992642

1.2.2 ModuleList类

ModuleList接收一个子模块的列表作为输入，然后也可以类似List那样进行append和extend操作:

demo3:

# ModuleList类
net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # # 类似List的append操作
print(net[-1])  # 类似List的索引访问
print(net)
# net(torch.zeros(1, 784)) # 会报NotImplementedError

out3:

Linear(in_features=256, out_features=10, bias=True)
ModuleList(
  (0): Linear(in_features=784, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)

既然Sequential和ModuleList都可以进行列表化构造网络，那二者区别是什么呢。

ModuleList仅仅是一个储存各种模块的列表，这些模块之间没有联系也没有顺序（所以不用保证相邻层的输入输出维度匹配），而且没有实现forward功能需要自己实现，所以上面执行net(torch.zeros(1, 784))会报NotImplementedError；而Sequential内的模块需要按照顺序排列，要保证相邻层的输入输出大小相匹配，内部forward功能已经实现。

1.2.3 ModuleDict类

ModuleDict接收一个子模块的字典作为输入, 然后也可以类似字典那样进行添加访问操作:

demo4:

# ModuleDict类
net = nn.ModuleDict({
    'linear': nn.Linear(784, 256),
    'act': nn.ReLU(),
})
net['output'] = nn.Linear(256, 10) # 添加
print(net['linear']) # 访问
print(net.output)
print(net)
# net(torch.zeros(1, 784)) # 会报NotImplementedError

out4:

Linear(in_features=784, out_features=256, bias=True)
Linear(in_features=256, out_features=10, bias=True)
ModuleDict(
  (linear): Linear(in_features=784, out_features=256, bias=True)
  (act): ReLU()
  (output): Linear(in_features=256, out_features=10, bias=True)
)

和ModuleList一样，ModuleDict实例仅仅是存放了一些模块的字典，并没有定义forward函数需要自己定义。同样，ModuleDict也与Python的Dict有所不同，ModuleDict里的所有模块的参数会被自动添加到整个网络中。

1.3 构造复杂的模型

虽然上面介绍的这些类可以使模型构造更加简单，且不需要定义forward函数，但直接继承Module类可以极大地拓展模型构造的灵活性。下面我们构造一个稍微复杂点的网络FancyMLP。在这个网络中，我们通过get_constant函数创建训练中不被迭代的参数，即常数参数。在前向计算中，除了使用创建的常数参数外，我们还使用Tensor的函数和Python的控制流，并多次调用相同的层。

在这个FancyMLP模型中，我们使用了常数权重rand_weight（注意它不是可训练模型参数）、做了矩阵乘法操作（torch.mm）并重复使用了相同的Linear层。

demo5:

# 构建复杂的模型
class FancyMLP(nn.Module):
    def __init__(self, **kwargs): # 关键字参数允许你传入0个或任意个含参数名的参数，这些关键字参数在函数内部自动组装为一个dict
        super(FancyMLP, self).__init__(**kwargs)

        self.rand_weight = torch.rand((20, 20), requires_grad=False) # 不可训练参数（常数参数）
        self.linear = nn.Linear(20, 20)

    def forward(self, x):
        x = self.linear(x)
        # 使用创建的常数参数，以及nn.functional中的relu函数和mm函数
        x = nn.functional.relu(torch.mm(x, self.rand_weight.data) + 1)

        # 复用全连接层。等价于两个全连接层共享参数
        x = self.linear(x)
        # 控制流，这里我们需要调用item函数来返回标量进行比较
        while x.norm().item() > 1:  # norm()计算Tensor中所有元素的2范数
            x /= 2
        if x.norm().item() < 0.8:
            x *= 10
        return x.sum()
#测试该模型的前向计算
X = torch.rand(2, 20)
net = FancyMLP()
print(net)
print(net(X))

#因为FancyMLP和Sequential类都是Module类的子类，所以可以嵌套调用它们
class NestMLP(nn.Module):
    def __init__(self, **kwargs):
        super(NestMLP, self).__init__(**kwargs)
        self.net = nn.Sequential(nn.Linear(40, 30), nn.ReLU())

    def forward(self, x):
        return self.net(x)

net = nn.Sequential(NestMLP(), nn.Linear(30, 20), FancyMLP())

X = torch.rand(2, 40)
print(net)
print(net(X))

out5:

FancyMLP(
  (linear): Linear(in_features=20, out_features=20, bias=True)
)
tensor(0.4308, grad_fn=<SumBackward0>)
Sequential(
  (0): NestMLP(
    (net): Sequential(
      (0): Linear(in_features=40, out_features=30, bias=True)
      (1): ReLU()
    )
  )
  (1): Linear(in_features=30, out_features=20, bias=True)
  (2): FancyMLP(
    (linear): Linear(in_features=20, out_features=20, bias=True)
  )
)
tensor(4.3046, grad_fn=<SumBackward0>)

小结

1.可以通过继承Module类来构造模型。
2.Sequential、ModuleList、ModuleDict类都继承自Module类。
3.与Sequential不同，ModuleList和ModuleDict并没有定义一个完整的网络，它们只是将不同的模块存放在一起，需要自己定义forward函数。
4.虽然Sequential等类可以使模型构造更加简单，但直接继承Module类可以极大地拓展模型构造的灵活性。

二、模型参数的访问、初始化和共享

我们先定义一个与上一节中相同的含单隐藏层的多层感知机。我们依然使用默认方式初始化它的参数，并做一次前向计算。与之前不同的是，在这里我们从nn中导入了init模块，它包含了多种模型初始化方法。

2.1 访问模型参数

demo6:

import torch
from torch import nn

net = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))  # pytorch已进行默认初始化

print(net)
X = torch.rand(2, 4)
Y = net(X).sum()


# 访问模型参数
# 对于Sequential实例中含模型参数的层，我们可以通过Module类的parameters()或者named_parameters方法来访问所有参数（以迭代器的形式返回），
# 后者除了返回参数Tensor外还会返回其名字。访问多层感知机net的所有参数：
print(type(net.named_parameters()))
for name, param in net.named_parameters():
    print(name, param.size())

# 返回的名字自动加上了层数的索引作为前缀。 我们再来访问net中单层的参数。
# 对于使用Sequential类构造的神经网络，我们可以通过方括号[]来访问网络的任一层。索引0表示隐藏层为Sequential实例最先添加的层
for name, param in net[0].named_parameters():
    print(name, param.size(), type(param))

# 因为这里是单层的所以没有了层数索引的前缀。另外返回的param的类型为torch.nn.parameter.Parameter，
# 其实这是Tensor的子类，和Tensor不同的是如果一个Tensor是Parameter，那么它会自动被添加到模型的参数列表里
class MyModel(nn.Module):
    def __init__(self, **kwargs):
        super(MyModel, self).__init__(**kwargs)
        self.weight1 = nn.Parameter(torch.rand(20, 20))
        self.weight2 = torch.rand(20, 20)
    def forward(self, x):
        pass

n = MyModel()
for name, param in n.named_parameters():
    print(name)

# 上面的代码中weight1在参数列表中但是weight2却没在参数列表中
# 因为Parameter是Tensor，即Tensor拥有的属性它都有，比如可以根据data来访问参数数值，用grad来访问参数梯度
weight_0 = list(net[0].parameters())[0]
print(weight_0.data)
print(weight_0.grad) # 反向传播前梯度为None
Y.backward()
print(weight_0.grad)

out6:

Sequential(
  (0): Linear(in_features=4, out_features=3, bias=True)
  (1): ReLU()
  (2): Linear(in_features=3, out_features=1, bias=True)
)
<class 'generator'>
0.weight torch.Size([3, 4])
0.bias torch.Size([3])
2.weight torch.Size([1, 3])
2.bias torch.Size([1])
weight torch.Size([3, 4]) <class 'torch.nn.parameter.Parameter'>
bias torch.Size([3]) <class 'torch.nn.parameter.Parameter'>
weight1
tensor([[-0.1439,  0.0765, -0.0037, -0.3121],
        [ 0.3675,  0.0878,  0.3473,  0.4148],
        [-0.0551,  0.2627, -0.2301, -0.4903]])
None
tensor([[-0.3761, -0.5268, -0.5040, -0.3642],
        [-0.5463, -0.7652, -0.7320, -0.5290],
        [ 0.0000,  0.0000,  0.0000,  0.0000]])

2.2 初始化模型参数

（数值稳定性和模型初始化）中提到了PyTorch中nn.Module的模块参数都采取了较为合理的初始化策略（不同类型的layer具体采样的哪一种初始化方法的可参考源代码）。但我们经常需要使用其他方法来初始化权重。PyTorch的init模块里提供了多种预设的初始化方法。在下面的例子中，我们将权重参数初始化成均值为0、标准差为0.01的正态分布随机数，并依然将偏差参数清零。

自定义初始化方法

有时候我们需要的初始化方法并没有在init模块中提供。这时，可以实现一个初始化方法，从而能够像使用其他初始化方法那样使用它。在这之前我们先来看看PyTorch是怎么实现这些初始化方法的，例如torch.nn.init.normal_：

def normal_(tensor, mean=0, std=1):
    with torch.no_grad():
        return tensor.normal_(mean, std)

可以看到这就是一个inplace改变Tensor值的函数，而且这个过程是不记录梯度的。 类似的我们来实现一个自定义的初始化方法。在例子里，我们令权重有一半概率初始化为0，有另一半概率初始化为[-10,-5]和[5,10]两个区间里均匀分布的随机数。

demo7:

import torch
from torch import nn
from torch.nn import init

net = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))  # pytorch已进行默认初始化

# 将权重参数初始化成均值为0、标准差为0.01的正态分布随机数，并依然将偏差参数清零
for name, param in net.named_parameters():
    if 'weight' in name:
        init.normal_(param, mean=0, std=0.01)
        print(name, param.data)

# 使用常数来初始化权重参数
for name, param in net.named_parameters():
    if 'bias' in name:
        init.constant_(param, val=0)
        print(name, param.data)


# 自定义初始化方法
# 令权重有一半概率初始化为0，有另一半概率初始化为[-10,-5]和[5,10]两个区间里均匀分布的随机数
def init_weight_(tensor):
    with torch.no_grad():
        tensor.uniform_(-10, 10)  # 生成与数据tensor同形的均匀分布的数据
        tensor *= (tensor.abs() >= 5).float()  # 绝对值小于5的都变为0

for name, param in net.named_parameters():
    if 'weight' in name:
        init_weight_(param)
        print(name, param.data)

# 改变这些参数的data来改写模型参数值同时不会影响梯度
for name, param in net.named_parameters():
    if 'bias' in name:
        param.data += 1
        print(name, param.data)

out7:

0.weight tensor([[ 0.0004, -0.0088, -0.0044, -0.0100],
        [-0.0065,  0.0003, -0.0122, -0.0008],
        [-0.0064,  0.0069, -0.0028, -0.0080]])
2.weight tensor([[ 0.0184, -0.0301,  0.0021]])
0.bias tensor([0., 0., 0.])
2.bias tensor([0.])
0.weight tensor([[-9.7925,  0.0000,  0.0000, -7.5638],
        [-7.7713, -0.0000,  7.1479,  0.0000],
        [-9.4375,  0.0000, -0.0000,  7.5921]])
2.weight tensor([[ 0.0000, -5.0063,  7.0720]])
0.bias tensor([1., 1., 1.])
2.bias tensor([1.])

2.4 共享模型参数

在有些情况下，我们希望在多个层之间共享模型参数。1.3提到了如何共享模型参数: Module类的forward函数里多次调用同一个层。此外，如果我们传入Sequential的模块是同一个Module实例的话参数也是共享的。

demo8:

import torch
from torch import nn
from torch.nn import init

linear = nn.Linear(1, 1, bias=False)
net = nn.Sequential(linear, linear)  # 传入Sequential的模块是同一个Module实例的话参数也是共享的
print(net)
for name, param in net.named_parameters():
    init.constant_(param, val=3)
    print(name, param.data)

# 在内存中，这两个线性层其实一个对象
print(id(net[0]) == id(net[1]))
print(id(net[0].weight) == id(net[1].weight))

# 因为模型参数里包含了梯度，所以在反向传播计算时，这些共享的参数的梯度是累加的
x = torch.ones(1, 1)
y = net(x).sum()
print(y)
y.backward()
print(net[0].weight.grad) # 单次梯度是3，两次所以就是6

out8:

Sequential(
  (0): Linear(in_features=1, out_features=1, bias=False)
  (1): Linear(in_features=1, out_features=1, bias=False)
)
0.weight tensor([[3.]])
True
True
tensor(9., grad_fn=<SumBackward0>)
tensor([[6.]])

三、自定义层

深度学习的一个魅力在于神经网络中各式各样的层，例如全连接层和后面章节中将要介绍的卷积层、池化层与循环层。虽然PyTorch提供了大量常用的层，但有时候我们依然希望自定义层。本节将介绍如何使用Module来自定义层，从而可以被重复调用。

3.1 不含模型参数的自定义层

我们先介绍如何定义一个不含模型参数的自定义层。事实上，这和模型构造中介绍的使用Module类构造模型类似。下面的CenteredLayer类通过继承Module类自定义了一个将输入减掉均值后输出的层，并将层的计算定义在了forward函数里。这个层里不含模型参数。

demo9:

import torch
from torch import nn

# CenteredLayer类通过继承Module类自定义了一个将输入减掉均值后输出的层，并将层的计算定义在了forward函数里。这个层里不含模型参数
class CenteredLayer(nn.Module):
    def __init__(self, **kwargs):
        super(CenteredLayer, self).__init__(**kwargs)
    def forward(self, x):
        return x - x.mean()

# 实例化这个层，然后做前向计算
layer = CenteredLayer()
print(layer(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float)))

#构造更复杂的模型
net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
y = net(torch.rand(4, 8))
print(y.shape)
print(y.mean().item())  # 均值是浮点数，所以它的值是一个很接近0的数

out9:

tensor([-2., -1.,  0.,  1.,  2.])
torch.Size([4, 128])
6.984919309616089e-09

3.2 含模型参数的自定义层

我们还可以自定义含模型参数的自定义层。其中的模型参数可以通过训练学出。

在模型参数的访问、初始化和共享中介绍了Parameter类其实是Tensor的子类，如果一个Tensor是Parameter，那么它会自动被添加到模型的参数列表里。所以在自定义含模型参数的层时，我们应该将参数定义成Parameter，除了直接定义成Parameter类外，还可以使用ParameterList和ParameterDict分别定义参数的列表和字典。

ParameterList接收一个Parameter实例的列表作为输入然后得到一个参数列表，使用的时候可以用索引来访问某个参数，另外也可以使用append和extend在列表后面新增参数。

而ParameterDict接收一个Parameter实例的字典作为输入然后得到一个参数字典，然后可以按照字典的规则使用了。例如使用update()新增参数，使用keys()返回所有键值，使用items()返回所有键值对等等，可参考官方文档。

demo10:

import torch
from torch import nn

# 含模型参数的自定义层
class MyListDense(nn.Module):
    def __init__(self):
        super(MyListDense, self).__init__()
        self.params = nn.ParameterList([nn.Parameter(torch.randn(4, 4)) for i in range(3)])
        self.params.append(nn.Parameter(torch.randn(4, 1)))

    def forward(self, x):
        for i in range(len(self.params)):
            x = torch.mm(x, self.params[i])
        return x
net = MyListDense()
print(net)

class MyDictDense(nn.Module):
    def __init__(self):
        super(MyDictDense, self).__init__()
        self.params = nn.ParameterDict({
                'linear1': nn.Parameter(torch.randn(4, 4)),
                'linear2': nn.Parameter(torch.randn(4, 1))
        })
        self.params.update({'linear3': nn.Parameter(torch.randn(4, 2))}) # 新增

    def forward(self, x, choice='linear1'):
        return torch.mm(x, self.params[choice])

net = MyDictDense()
print(net)

# 根据传入的键值来进行不同的前向传播
x = torch.ones(1, 4)
print(net(x, 'linear1'))
print(net(x, 'linear2'))
print(net(x, 'linear3'))

# 使用自定义层构造模型，它和PyTorch的其他层在使用上很类似。
net = nn.Sequential(
    MyDictDense(),
    MyListDense(),
)
print(net)
print(net(x))

out10:

MyListDense(
  (params): ParameterList(
      (0): Parameter containing: [torch.FloatTensor of size 4x4]
      (1): Parameter containing: [torch.FloatTensor of size 4x4]
      (2): Parameter containing: [torch.FloatTensor of size 4x4]
      (3): Parameter containing: [torch.FloatTensor of size 4x1]
  )
)
MyDictDense(
  (params): ParameterDict(
      (linear1): Parameter containing: [torch.FloatTensor of size 4x4]
      (linear2): Parameter containing: [torch.FloatTensor of size 4x1]
      (linear3): Parameter containing: [torch.FloatTensor of size 4x2]
  )
)
tensor([[ 1.8806,  2.2471,  1.1397, -3.8681]], grad_fn=<MmBackward>)
tensor([[3.0825]], grad_fn=<MmBackward>)
tensor([[ 1.1917, -0.8614]], grad_fn=<MmBackward>)
Sequential(
  (0): MyDictDense(
    (params): ParameterDict(
        (linear1): Parameter containing: [torch.FloatTensor of size 4x4]
        (linear2): Parameter containing: [torch.FloatTensor of size 4x1]
        (linear3): Parameter containing: [torch.FloatTensor of size 4x2]
    )
  )
  (1): MyListDense(
    (params): ParameterList(
        (0): Parameter containing: [torch.FloatTensor of size 4x4]
        (1): Parameter containing: [torch.FloatTensor of size 4x4]
        (2): Parameter containing: [torch.FloatTensor of size 4x4]
        (3): Parameter containing: [torch.FloatTensor of size 4x1]
    )
  )
)
tensor([[10.8090]], grad_fn=<MmBackward>)

四、读取和存储

到目前为止，学习了如何处理数据以及如何构建、训练和测试深度学习模型。然而在实际中，有时需要把训练好的模型部署到很多不同的设备。在这种情况下，我们可以把内存中训练好的模型参数存储在硬盘上供后续读取使用。

4.1 读写Tensor

我们可以直接使用save函数和load函数分别存储和读取Tensor。save使用Python的pickle实用程序将对象进行序列化，然后将序列化的对象保存到disk，使用save可以保存各种对象,包括模型、张量和字典等。而load使用pickle unpickle工具将pickle的对象文件反序列化为内存。

demo11:

import torch

# 创建了Tensor变量x，并将其存在文件名同为x.pt的文件里
x = torch.ones(3)
torch.save(x, 'x.pt')

# 将数据从存储的文件读回内存
x2 = torch.load('x.pt')
print(x2)

# 存储一个Tensor列表并读回内存
y = torch.zeros(4)
torch.save([x, y], 'xy.pt')
xy_list = torch.load('xy.pt')
print(xy_list)

# 存储并读取一个从字符串映射到Tensor的字典
torch.save({'x': x, 'y': y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
print(xy)

out11:

tensor([1., 1., 1.])
[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]
{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}

4.2 读写模型

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

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

PyTorch中保存和加载训练模型有两种常见的方法:

1.仅保存和加载模型参数(state_dict)；
2.保存和加载整个模型。

1.保存和加载state_dict(推荐方式)

保存：

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

加载：

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

2.保存和加载整个模型

保存：

torch.save(model, PATH)

加载：

model = torch.load(PATH)

demo12:

import torch
from torch import nn

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()
print(net.state_dict())

# 优化器(optim)的state_dict
optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
print(optimizer.state_dict())


# 保存和加载state_dict(推荐方式)
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)
print(Y2 == Y) # net和net2都有同样的模型参数，对同一个输入X的计算结果是一样的

out12:

OrderedDict([('hidden.weight', tensor([[-0.1988,  0.2219, -0.2090],
        [-0.3193, -0.3393,  0.1803]])), ('hidden.bias', tensor([-0.4834,  0.4835])), ('output.weight', tensor([[0.5424, 0.1767]])), ('output.bias', tensor([0.1720]))])
{'state': {}, 'param_groups': [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3]}]}
tensor([[True],
        [True]])

五、GPU计算

到目前为止，我们一直在使用CPU计算。对复杂的神经网络和大规模的数据来说，使用CPU来计算可能不够高效。在本节中，将介绍如何使用单块NVIDIA GPU来计算。所以需要确保已经安装好了PyTorch GPU版本。

5.1 计算设备

PyTorch可以指定用来存储和计算的设备，如使用内存的CPU或者使用显存的GPU。默认情况下，PyTorch会将数据创建在内存，然后利用CPU来计算。

demo13:

import torch

print(torch.cuda.is_available()) # 查看GPU是否可用,True

# 查看GPU数量
print(torch.cuda.device_count()) # 输出 1

# 查看当前GPU索引号，索引号从0开始：
print(torch.cuda.current_device()) # 输出 0

# 根据索引号查看GPU名字:
print(torch.cuda.get_device_name(0)) # 输出 'NVIDIA GeForce RTX 3050 Laptop GPU'

out13:

True
1
0
NVIDIA GeForce RTX 3050 Laptop GPU

5.2 Tensor的GPU计算

默认情况下，Tensor会被存在内存上。因此，之前我们每次打印Tensor的时候看不到GPU相关标识。

1.使用.cuda()可以将CPU上的Tensor转换（复制）到GPU上。如果有多块GPU，我们用.cuda(i)来表示第 i 块GPU及相应的显存（i从0开始）且cuda(0)和cuda()等价。
2.我们可以通过Tensor的device属性来查看该Tensor所在的设备。
3.可以直接在创建的时候就指定设备。
4.如果对在GPU上的数据进行运算，那么结果还是存放在GPU上。
5.需要注意的是，存储在不同位置中的数据是不可以直接进行计算的。即存放在CPU上的数据不可以直接与存放在GPU上的数据进行运算，位于不同GPU上的数据也是不能直接进行计算的。
z = y + x.cpu()会报错:RuntimeError: Expected object of type torch.cuda.LongTensor but found type torch.LongTensor for argument #3 'other'

demo14:

import torch

# Tensor的GPU计算
x = torch.tensor([1, 2, 3])
print(x)
# 使用.cuda()可以将CPU上的Tensor转换（复制）到GPU上
x = x.cuda(0)
print(x)

# 通过Tensor的device属性来查看该Tensor所在的设备
print(x.device)

# 可以直接在创建的时候就指定设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = torch.tensor([1, 2, 3], device=device)
# x = torch.tensor([1, 2, 3]).to(device)  # 这种创建方式也可以
print(x)

# 对在GPU上的数据进行运算，结果还是存放在GPU上
y = x**2
print(y)

out14:

tensor([1, 2, 3])
tensor([1, 2, 3], device='cuda:0')
cuda:0
tensor([1, 2, 3], device='cuda:0')
tensor([1, 4, 9], device='cuda:0')

5.3 模型的GPU计算

同Tensor类似，PyTorch模型也可以通过.cuda转换到GPU上。我们可以通过检查模型的参数的device属性来查看存放模型的设备。

同样的，我么需要保证模型输入的Tensor和模型都在同一设备上，否则会报错。

demo15:

import torch
from torch import nn

#  模型的GPU计算
net = nn.Linear(3, 1)
print(list(net.parameters())[0].device)  # 模型在CPU上

# 转换到GPU上
net.cuda()
print(list(net.parameters())[0].device)

# 需要保证模型输入的Tensor和模型都在同一设备上，否则会报错
x = torch.rand(2,3).cuda()
print(net(x))

out15:

cpu
cuda:0
tensor([[0.1535],
        [0.4055]], device='cuda:0', grad_fn=<AddmmBackward>)