1 自定义块和层

1.1 自定义块

复杂的网络，通常会有很多复用的模块，为了方便模型的表示，引入了块的概念。
它可以描述单个层或是由多个层组成的组件。
一个块通常只需要提供构造函数和前向传播函数，除非我们实现一个新的运算符， 否则我们不必担心反向传播函数或参数初始化， 系统将自动生成这些。

# 自定义块
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)
    
    def forward(self, X):
        return self.output(F.relu(self.hidden(X)))
  • 继承自nn.Module
  • 要先调用父类的初始化函数super().__init__()
  • 在每次调用前向传播函数时调用实例化的层
  • F.relu是激活函数的函数形式

之前使用的Sequential类其实也可以看作是一个块，里面的各层是按顺序堆叠的。
其实里面就实现了两个关键函数：

  • 将块/层逐个添加到列表
  • 将输入按追加块/层的顺序传递
class MySequential(nn.Module):
    def __init__(self, *args):
        super().__init__()
        for idx, module in enumerate(args):
            self._modules[str(idx)] = module
    
    def forward(self, X):
        for module in self._modules.values():
            X = module(X)
        return X
  • *arg接受任意个位置参数，并把它们打包成一个元组
  • enumerate返回带下标的迭代器
  • nn.Module内部有一个特殊的字典：self._modules，用来存放子模块

并不是所有的架构都是简单的顺序架构，我们可能希望在前向传播函数中执行Python的控制流
或者是在某一层加入一个常量之类的操作
自定义块都是允许的

class FancyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.rand_weight = torch.rand((20, 20), requires_grad=False)
        self.linear = nn.Linear(20, 20)
    
    def forward(self, x):
        x = self.linear(x)
        x = F.relu(torch.mm(x, self.rand_weight) + 1)
        x = self.linear(x)
        while x.abs().sum() > 1:
            x /= 2
        if x.norm().sum() < 1:
            x *= 10
        return x.sum()

1.2 自定义层

自定义层和自定义块其实没有严格的界限，主要是语义上的区别。
他们都是继承自nn.Module
要构建自定义层，我们只需继承基础层类并实现前向传播功能。

class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        return x - x.mean()

2 参数管理

2.1 参数访问

2.1.1 访问目标参数

当我们通过Sequential来定义模型的时候，我们可以通过索引来访问模型的任意层，就像一个列表一样
比如我们想访问第三层的参数：print(net[2].state_dict())
在PyTorch中，模型的科学系参数和缓冲区都保存在一个字典state_dict中
调用state_dict()会返回一个有序字典，保存了所有参数（权重、偏置）和缓冲（如 BatchNorm 的均值方差）。
注意：state_dict 只保存参数和缓冲区，不保存模型的结构或 forward 函数。常用来保存或者加载模型参数。

参数是复合的torch.nn.Parameter对象，包含值、梯度和额外信息。

print(net[1].weight.data)           # 访问权重的值
print(net[1].bias.grad)             # 访问偏置的梯度
print(net[0].hidden.weight.data)    # 对于块内参数访问

2.1.2 访问全部参数

当我们需要对所有参数执行操作时，逐个访问它们可能会很麻烦。 当我们处理更复杂的块（例如，嵌套块）时，情况可能会变得特别复杂， 因为我们需要递归整个树来提取每个子块的参数。
model.parameters()会返回一个迭代器，包含模型中所有需要梯度的参数（nn.Parameter），一般是把这个丢给优化器。
对于我们访问参数更常用的是model.named_parameters()，会返回"参数名字+参数"更方便的定位。

# 访问第一个全连接层的参数
print(*[(name, param.shape) for name, param in net[0].named_parameters()])
# 访问整个网络的参数
print(*[(name, param.shape) for name, param in net.named_parameters()])
  • model.named_parameters()返回一个迭代器
  • (name, param.shape) for name, param in net.named_parameters()是一个生成器，需要用列表接收它
  • *[...]解包列表里的东西

2.1.3 访问嵌套块里的参数

先把模型结构打印出来

Sequential(
  (0): MLP(
    (hidden): Linear(in_features=20, out_features=256, bias=True)
    (output): Linear(in_features=256, out_features=10, bias=True)
  )
  (1): Linear(in_features=10, out_features=20, bias=True)
  (2): Sequential(
    (0): MLP(
      (hidden): Linear(in_features=20, out_features=256, bias=True)
      (output): Linear(in_features=256, out_features=10, bias=True)
    )
    (1): Linear(in_features=10, out_features=20, bias=True)
  )
  (3): FancyMLP(
    (linear): Linear(in_features=20, out_features=20, bias=True)
  )
  (4): CenteredLayer()
)

第三层，即(2): Sequential，其实是个块，里面还嵌套了(0): MLP
所以可以用索引访问子模块

print(net[2][0].hidden.weight.data)  # 访问嵌套块的第一个块

2.2 参数初始化

深度学习框架提供默认随机初始化， 也允许我们创建自定义初始化方法， 满足我们通过其他规则实现初始化权重。
PyTorch的nn.init模块提供了多种预置初始化方法。
将权重初始化为标准差为0.01的高斯随机变量

def init_normal(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.zeros_(m.bias)

初始化为给定的常数

def init_constant(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 1)
        nn.init.zeros_(m.bias)

Xavier初始化方法

def init_xavier(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)

我们也可以自定义任意的初始化方法

def my_init(m):
    if type(m) == nn.Linear:
        print("Init", *[(name, param.shape)
                        for name, param in m.named_parameters()][0])
        nn.init.uniform_(m.weight, -10, 10)
        m.weight.data *= m.weight.data.abs() >= 5

2.3 共享参数

其实就是让多个层的名字一样这样他们就会使用一样的参数

shared = nn.Linear(8, 8)
net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                    shared, nn.ReLU(),
                    shared, nn.ReLU(),
                    nn.Linear(8, 1))

如果我们改变其中一个参数，另一个参数也会改变。
注意：在反向传播的时候两个层的梯度会累加，然后在调用优化器更新参数的时候，会根据已累加的梯度对这个共享层的参数进行更新。

3 延后初始化

延迟初始化很方便，允许框架自动推断参数形状，从而轻松修改架构。

net = nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10))

4 读写文件

4.1 张量读写

torch.save和torch.load就可以读写张量。

x = torch.arange(4)
print(x)
torch.save(x, 'x-file')
x2 = torch.load('x-file')
print(x2)

4.2 模型参数读写

由于net.state_dict()只会保存参数的字典和值，不会保存结构
因此我们需要实例化一个内部结构一样的模型，才能加载

net = MLP()
torch.save(net.state_dict(), 'mlp.params')

net2 = MLP()
net2.load_state_dict(torch.load('mlp.params'))
net2.eval()

5 GPU加速

GPU可以极大加速模型的运行速度。
我们可以通过print(torch.cuda.is_available())查看有无可用的GPU

通过以下函数我们可以同时兼顾GPU和CPU环境（但是对于稍大一点的模型通常在CPU计算的速度非常感人）

def try_gpu(i=0):
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

为了加速计算，我们需要把模型和数据放在同一个gpu上，
可以在定义的时候就指定设备，也可以在定义完后使用to()将张量/模型转移到指定设备上。

X = torch.ones(2, 3, device=try_gpu())
net = net.to(device=try_gpu())