PyTorch笔记 #2 神经网络

本文介绍使用 PyTorch 搭建神经网络模型的方法，由于涉及的内容众多，本文将以完成一个神经网络实验的顺序展开介绍。分别是：数据加载、模型搭建、训练阶段、评估阶段、模型保存与加载。

本文内容基于 官方 API 文档 及 PyTorch 中文文档、网上的博客等，实战部分推荐以下教程：

• pytorch-tutorial：逐步搭建网络，包括 FFNN、CNN、RNN、ResNet、BiLSTM 等经典模型，代码简练。
• PyTorchZeroToAll：HKUST 的课程，有视频讲解，也是经典模型的搭建过程。
• 深度学习 PyTorch 训练代码模板：完整的深度学习模板，在处理轻量级任务时可以参考。在快速开发时建议不要重复造轮子，优先使用模板（比赛的公开 baseline、要对标的开源模型等）。

数据加载

PyTorch 中的 torch.utils.data 模块封装了所有关于数据集的操作，可以帮助我们使用一些预加载数据集或者自定义的数据，其中 Dataset 可以储存样本以及对应标签，而 Dataloader 则用于将 Dataset 封装成可迭代对象，使得我们能够更加容易地获取样本数据。

预加载 or 自定义

PyTorch 中提供了大量预加载的数据集（如 MNIST、FashionMNIST），存放在不同的领域库中，分别是：torchvision、torchtext、torchaudio 库。

下面以 FashionMNIST 为例介绍加载流程，其包含 60000 个训练样本以及 10000 个测试样本，每个样本由一张灰度图、一个类别标签组成。

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor

training_data = datasets.FashionMNIST(
    root="data",			# root 是训练/测试数据所存储的路径，相对于当前位置
    train=True,				# train 指定训练或测试数据集
    download=True,			# 如果数据集在 root 中不存在则从互联网上下载
    transform=ToTensor()	# 特征与标签的变换，这里转成张量
)
test_data = datasets.FashionMNIST(
    root="data",
    train=False,			# False 表示指定测试集
    download=True,
    transform=ToTensor()
)

自定义数据集需要通过面向对象操作来完成，案例如下：

import os
import pandas as pd
from torchvision.io import read_image
from torch.utils.data import Dataset, DataLoader

# 定义 MyDataset() 类, 继承自 Dataset 父类，重写 init、getitem、len 方法
class MyDataset(Dataset):
    
    # 实例化数据集对象，从文件目录中获取数据集，数据变换也在这里定义
    # 如果数据特征过大，无法一次读入，这里就要记录文件目录，按需读取即可
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file) # (file_name, label) 列表
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform
    
    # 返回数据集中样本的数量
	def __len__(self):
		return len(self.img_labels)
    
    # 加载并返回数据集中给定索引 idx 的样本特征与标签
    # 如果数据集不在内存中，要从指定目录中读取，并根据实际情况组织数据
    # 通常在这里对数据进行变换和增强，并转为 Tensor 返回
	def __getitem__(self, idx):
		img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0]) # 获取文件路径
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1] # 获取对应图像的标签
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

数据集迭代访问

获取的数据会以列表的形式存储，特征和标签构成一对元组。如果这样通过遍历列表的形式来访问数据集，则效率太低，我们希望以小批次（Mini-batch）为单位读取样本，并且在每一个 epoch 中打乱数据的顺序以防止模型过拟合，同时利用 Python 的并行处理（Multi-Processing）功能加速数据处理。

使用 DataLoader 模块就可以实现上述功能：

from torch.utils.data import DataLoader

# 声明 DataLoader 模块，得到可迭代对象
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True, num_workers=4)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True, num_workers=4)

# 用 iter() 转迭代器，再用 next() 访问，每次获取一个批次，分包含 64 个特征与标签
train_features, train_labels = next(iter(train_dataloader))

# 可迭代对象可以用 for 循环遍历，训练时每个 epoch 中的顺序会变化
num_epoches = 100
for epoch in range(num_epoches):
    for img, label in train_dataloader:
        # 训练代码，这里每一组 img, label 都是 64 个样本

其中 num_workers 参数代表在声明时一次性创建的工作进程数量，此时 batch_sampler 会将指定 batch 分配给每个 worker，worker 负责将 batch 载入显存。当 DataLoader 在某轮迭代需要用到该 batch 时，可以直接使用。分配的进程数量越多，获取下一 batch 的速度越快，因为它可能早就放在显存中了，但这也加重了 CPU 负担。一般将该值设置为 CPU 核心数量，根据性能变化逐步调整。

除了以上四个常用参数，有时还会用到一个 collate_fn 参数，顾名思义，该函数用于对样本进行核对、校勘。在 Dataset 类中调用 __getitem__ 后，每次返回的都是一个样本，最后构成一个 Batch 的样本。如果再这个过程中出现出现什么样本错误（如存在空值、错误值），或者需要对批次进行特殊处理（如对齐整个 Batch 样本的长度），则需要在该函数中完成。

collate_fn 默认是等于 default_collate，这是 PyTorch 自带的处理函数，也可以自己定义：

# 输入的 batch 是 list of tuple (img, label)，本例子为过滤 None 数据
def collate_fn(batch):
    batch = list(filter(lambda x: x[0] is not None, batch))
    if len(batch) == 0: 
        return torch.Tensor()
    return default_collate(batch)

# 假设 label 是文本，本例子为对齐文本的长度
def collate_fn(batch):
    batch.sort(key=lambda x: len(x[1]), reverse=True)
    img, label = zip(*batch)
    pad_label = []
    lens = []
    max_len = len(label[0])
    for i in range(len(label)):
        temp_label = [0] * max_len
        temp_label[:len(label[i])] = label[i]
        pad_label.append(temp_label)
        lens.append(len(label[i]))
    return img, pad_label, lens

数据预处理

前文中已经出现过 torchvision.transforms 模块，该模块的功能是对图像进行预处理，内置各种数据增强操作，包括随机裁剪、翻转和旋转、变形和填充、修改属性等。每一个单独的变换都是一个函数，输入数据即可得到结果。

preprocess = transforms.Resize([256, 256])
img_transformed = preprocess(img)

如果要连续进行多种操作，则可以用 transforms.Compose() 组合，参数为操作列表。例如：

import torchvision.transforms as transforms

myTransforms = transforms.Compose([
    transforms.RandomResizedCrop(224),	# 将给定图像随机裁剪并缩放到 224x224
    transforms.RandomHorizontalFlip(),	# 给定的概率随机水平旋转，默认 0.5
    transforms.ToTensor(),				# 将给定图像转为 Tensor
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 归一化处理
])

data = MyDataset(
	annotations_file="annotations.csv",
    img_dir="data/img",
    transform=myTransforms	# 调用自定义的预处理组合
)

该案例中归一化处理传入的两个参数分别代表均值和标准差，这里取的是 ImageNet 数据集的值，这是 CV 领域一种常见的做法，通常只要数据集是自然场景的普通照片就能使用。如果是特殊场景、特殊风格的数据，则需要重新计算。

模型搭建

网络的定义也是通过面向对象操作来完成，torch.nn 模块封装了神经网络的大部分操作，案例如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义 MyNet() 类, 继承自 nn.Module 父类
# 父类中封装了 parameters() train() eval() zero_grad() 等常用方法
class MyNet(nn.Module):
    
    # 重写构造函数 __init__(), 在实例化 MyNet() 时自动调用
    # 输入参数除了必须的 self, 还包含其他实例化对象时必备的参数
    def __init__(self, num_classes=10):
        # 由于子类定义了构造函数，父类的构造函数不再自动执行，需要显式调用
        super().__init__()
        # 定义对象变量, 通过实例化 nn.Conv2d() 等类获得
        self.conv = nn.Conv2d(3, 6, 3)
        self.fc = nn.Linear(6 * 5 * 5, num_classes)

    # 重写前向传播方法 forward(), 对输入 x 进行操作
    def forward(self, x):
        # 可以调用 self 的对象变量，也可以直接调用 F 函数
        x = F.max_pool2d(F.relu(self.conv(x)), 2)
        x = torch.flatten(x, 1) # 将除了 batch 的维度展开到 1 维
        x = F.relu(self.fc(x))
        return x

net = MyNet()
input = torch.randn(1, 3, 12, 12) # (batch_size, channel, hight, width)
out = net(input)

如果这时候 print(net)，则会打印出类的所有属性，即每个构造函数中的网络部件：

MyNet(
  (conv): Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1))
  (fc): Linear(in_features=150, out_features=10, bias=True)
)

由于继承了父类的方法，调用 net.paprameters() 可以得到参数列表，各个层的参数依次排列：

>>> params = list(net.parameters())
>>> print(len(params)) # 参数的长度，这里卷积层和全连接层各有 weight 和 bias
# 4,
>>> print(params[0].size()) # 按索引访问指定参数
# torch.Size([6, 3, 3, 3])
>>> print([(name, param.shape) for name, param in net.named_parameters()]) # 一次性访问全部
# [('conv.weight', torch.Size([6, 3, 3, 3])),
#  ('conv.bias', torch.Size([6])),
#  ('fc.weight', torch.Size([10, 150])),
#  ('fc.bias', torch.Size([10]))]
>>> print(sum(p.numel() for p in net.parameters() if p.requires_grad)) # 计算总参数量

预加载的网络

PyTorch 中自带了搭建好的经典模型，可以直接调用，包括 AlexNet、VGG 等，详见 官方文档。可以选择是否使用预训练好的权重、原有的转换。下面的例子展示了使用预加载网络推理一个样本的过程：

from torchvision.io import read_image
from torchvision.models import resnet50, ResNet50_Weights

img = read_iamge("test.jpg")

# 使用预训练好的权重参数调用模型，也可以不使用预训练，即随机初始化权重
weights = ResNet50_Weights.DEFAULT
model = resnet50(weights=weights) 	# model = resnet50(weights=None)
model.eval() # 进入验证模式

# 使用模型默认的预处理方法（与权重对应）
preprocess = weights.transforms()	# 使用 preprocess(img) 调用变换

# 此处只有一张照片，因此我们在 dim=0 处补 1 表示 batch_size
batch = preprocess(img).unsqueeze(0)

# 使用预加载的网络预测出类别，并打印
prediction = model(batch).squeeze(0).softmax(0)
class_id = prediction.argmax().item()
score = prediction[class_id].item()
category_name = weights.meta["categories"][class_id]
print(f"{category_name}: {100 * score:.1f}%")

此外，我们还可以对已经定义好的模型进行添加、修改网络层：

vgg16 = torchvision.models.vgg16(pretrained=True)
# 添加网络层，第一个参数是网络层的命名
vgg16.classifier.add_module("add_linear", nn.Linear(1000,10)) # 在 classifier 里加一层
# 修改网络层，直接指定层数修改
vgg16.classifier[6] = nn.Linear(4096, 10) # 将 classifier 最后一层的输出由 1000 改为 10
# print(vgg16)

nn 网络模块

这里介绍一下在构建网络时常见的一些网络模块，即在 __init__ 中定义的部分，当然也可以在 nn.Sequential() 中组合，稍后会介绍。

• 全连接层：

'''
	in_features: 输入特征张量的大小，可以理解为该层神经元个数，即 (batch_size, in_features)
	out_features: 输出特征张量的大小，可以理解为下一层神经元个数
	bias: 是否启用偏置，除了 weight 数组，还会初始化一个 bias 数组
'''
nn.Flatten()	# 展平层，通常放在全连接层之前，如 (64, 3, 5, 5) -> (64, 75)
nn.Linear(in_features, out_features, bias=True) # 如 (64, 75) -> (64, 10)

• 常用的卷积层，最终输出的张量维度为 $( , ) $：

'''
	in_channels/out_channels: 输入/输出通道数，四维张量 [N, C, H, W] 中的 C
	kernel_size: 卷积核大小，可输入一个值或元组
	stride: 步幅，默认为 1，可输入一个值或元组
	padding: 填充，默认不填充，可输入一个值或元组
	padding_mode: 可选 zeros(零) reflect(镜像) replicate(复制) circular(循环)
	dilation: 是否采用空洞卷积，默认为 1 不采用
	卷积核参数量: [out_channels, in_channels, kernel_height, kernel_width]
'''
nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=True, padding_mode='zeros')

• 常用的池化层（无参数），最终输出的张量维度同上：

'''
	kernel_size: 池化窗口大小，可输入一个值或元组
	stride: 默认为窗口的大小，即每个窗口不重叠，与卷积不同
	return_indices: 是否返回最大值位置索引
	ceil_mode: 输出的形状取整方向，默认为 False 向下取整
'''
nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)
'''
	count_include_pad: 计算平均值时是否包括零填充，默认包括
	divisor_override: 除数大小，默认就是池化窗口大小，但可以自己指定
'''
nn.AvgPool2d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)
'''
	output_size: 固定输出尺寸，即可自适应窗口大小，可输入一个值或元组
'''
nn.MaxPool2d(output_size)
nn.AdaptiveAvgPool2d(output_size)

• 常用激活函数：也作为层来定义，但都没有参数，默认 inplace=Flase。

nn.Sigmoid() # S 型激活函数，[0, 1]，存在梯度消失/梯度爆炸/不好收敛/效率低的问题
nn.Tanh() # 双曲正切函数，[-1, 1]，相对更好收敛，但还是容易梯度消失/梯度爆炸/效率低
nn.ReLU() # 线性修正单元，计算快，单侧饱和，不存在梯度消失问题，但会有 dead 问题
nn.LeakyReLU(negative_slope=0.01) # 负区间也采用线性，梯度不再为 0，解决 dead 问题，但不再单侧饱和
nn.PReLU(num_parameters=1) # 可学习 alpha 的 LeakyReLU，需要有较多数据时才能使用
nn.ELU(alpha=1.0) # 负区间采用指数函数，既单侧饱和，又解决了 dead 问题
nn.Softplus(beta=1, threshold=20) # ReLU 的近似，保证输出为正数
nn.Softmax(dim=None) # 最后对输出值的处理，[0, 1] 且和为 1

• 特殊功能层：注意 \(\textrm{Dropout}\) 只在训练阶段启用，并且不带参数。而 \(\textrm{BN}\) 层实际是有参数存储的，包括仿射参数和测试阶段的移动平均。同样 \(\textrm{LN}\) 层也是有参数存储的，即仿射参数，但不会跟踪全局的移动平均，所有阶段都一致。

'''
	p: 张量元素被置为 0 的概率，置为 0 的神经元整个不激活
	inplace: 是否进行原地操作，默认为否
'''
# 主要用于一维数据，一般在线性层、激活函数之后
nn.Dropout(p=0.5, inplace=False)
# 主要用于二维数据，一般在卷积层、激活函数之后
nn.Dropout2d(p=0.5, inplace=False)
'''
	num_features: 通道数，1d 中代表 [N, L] 中的 L, 2d 中代表 [N, C, H, W] 中的 C
	eps: 归一化时加到分母（方差）上的小量，防止除零溢出
	momentum: 预测模型下进行移动平均的动量
	affine: 是否具有可学习的仿射参数 gamma 和 beta，用于拉伸和偏移分布
	track_running_stats: 是否跟踪均值和方差，进行移动平均
'''
# 一维 BN 是取 axis=0 进行归一化的，一般在线性层和激活函数之间
nn.BatchNorm1d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
# 二维 BN 是取 axis=(0,2,3) 进行归一化的，一般在卷积层和激活函数之间
nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'''
	normalized_shape: 传入一个整数表示最后一维的长度，传入列表表示最后两维的长度
	eps: 归一化时加到分母（方差）上的小量，防止除零溢出
	elementwise_affine: 是否具有可学习的仿射参数 gamma 和 beta，用于拉伸和偏移分布
'''
# 取 axis=-1 或 axis=(-2, -1) 进行归一化
nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True)

• 词嵌入层：在 NLP 任务中经常使用，用来代替 One-Hot 实现降维，并学习词之间的联系。

'''
	num_embeddings: 词典的大小，如果有 5000 词，id 就是 0-4999
	embedding_dim: 嵌入向量的维度，即用多少维的空间来表示一个词向量
	padding_idx: 输入句子长度不一致时，末尾的 PAD 的索引值
'''
nn.Embedding(num_embeddings, embedding_dim, padding_idx=None)
# 使用 Embedding 时要在输入数据前就预处理好格式，具体步骤如下：
# 原始句子 			['I am a boy.','How are you?','I am very lucky.']
# 分词映射到索引 	  [[3,6,5,6,7],[6,4,7,9,5]，[4,5,8,7]]
# 填充[EOS][PAD] 	 [[3,6,5,6,7,1],[6,4,7,9,5,1],[4,5,8,7,1,2]]
# 切片用于序列输入   [[3,6,4],[6,4,5],[5,7,8],[6,9,7],[7,5,1],[1,1,2]]
# 转为张量			batch = torch.LongTensor(batch)
embed = torch.nn.Embedding(num_embeddings, embedding_dim)
embed_batch = embed(batch) # [seq_len, batch_size] -> [seq_len, batch_size, embedding_dim]

如果想观察网络中每层的通道数、数据规模是否符合预期，可以构造一个单独的样本来前向传播，并实时获取其形状：

X = torch.randn(1, 3, 224, 224)
for layer in net:
    X=layer(X)
    print(layer.__class__.__name__, '\t\ output shape:\t', X.shape)

需要注意的是，如果要设计全新的网络模块，除了实现一个运算函数，还需要定义可训练的参数。nn.Parameter() 将一个不可训练的 Tensor 转换成可训练的类型 Parameter，并将这个 Parameter 绑定到模型的参数列表里面。比起参数 requires_grad=True 还多了绑定的过程：

class Conv2D(nn.Module):
    def __init__(self, kernel_size):
        super().__init__()
        self.weight = nn.Parameter(torch.rand(kernel_size))
        self.bias = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return corr2d(x, self.weight) + self.bias

F 函数式编程

在 PyTorch 的 nn 模块中，我们不需要手动定义网络层的权重和偏置，这就是 PyTorch 比 Tensorflow 简洁的地方。当然，PyTorch 也提供了 nn.Functional 函数式编程的方法，其中的 F.conv2d() 就和 Tensorflow 一样，要预先定义好卷积核的权重和偏置 nn.Parameter()，作为形参之一输入。

当然，为了不多此一举，网络中具有可学习参数的层（如全连接层、卷积层）通常会放在 __init__ 中显式定义，而不具有参数的层（如 ReLU、Dropout、BN 层）一般放在 forward 中使用 F 函数调用。

• 无参数的池化操作：

F.max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)
F.avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)

• 无参数的激活函数：以下函数都支持 F.relu_ 的原地版本。

F.sigmoid(input)
F.tanh(input)
F.relu(input)
F.leaky_relu(input, negative_slope=0.01)
F.prelu(input, weight) 	# 特例：这里 weight 是个可学习的参数，一般不会使用函数式编程
F.elu(input, alpha=1.0)
F.softplus(input, beta=1, threshold=20)
F.softmax(input)

• 无参数的特殊功能函数：包括 Dropout 和一些计算向量距离的函数。

F.dropout(input, p=0.5, inplace=False)
F.dropout2d(input, p=0.5, inplace=False)
F.cosine_similarity(x1, x2)  		# 计算两个向量的余弦相似度
F.pairwise_distance(x1, x2, p=2.0) 	# 计算两个向量的 p 范数距离
F.one_hot(x, num_classes)			# 生成长度为 num_classes 的 one_hot 张量组

注意在 F 函数式编程中，函数名全都为小写，与 nn 网络模块中不同！

使用块搭建深层网络

块（Block）比单个层（Layer）大，又比整个模型（Model）小，是实现一个深层网络的必备操作，通常用于封装以重复模式排列的若干层，最早在 \(\textrm{VGG}\) 中得到应用。

从编程角度来看，每个块也需要用类定义，并实现初始化、前向传播功能。实际中 PyTorch 提供了一个简单的方法 nn.Sequential，可以将不同神经网络层进行顺序拼接：

self.net = nn.Sequential(
    nn.Linear(20, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)
output = net(torch.rand(2, 20)) # output.shape = (2, 10)

如果希望多个层共享参数，只需先实例化一个层，并将其作为网络的多个部分，此时所有层的权重张量共享一个内存空间。这就相当于 __init__ 中的同一层在 forward 中被调用了两次，因此每次调用的 grad 也会相加。如下：

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

如果要封装更复杂的块，例如 VGG 中的 VGG_Block，则可以将其单独用类定义，并设置好参数：

def VGG_Block(num_convs, in_channels, out_channels):
    layers = [] # 以列表的形式拼接网络模块
    for _ in range(num_convs):
        layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))
        layers.append(nn.ReLU())
        in_channels = out_channels
    layers.append(nn.MaxPool2d(kernel_size=2,stride=2))
    return nn.Sequential(*layers) # 这里 * 表示解包，将列表拆分

网络初始化

默认情况下，PyTorch 会根据一个范围均匀地初始化权重和偏置矩阵， 这个范围是根据输入和输出维度计算出的。其中nn.init 模块提供了多种预置初始化方法。包括：

nn.init.zeros_(weight)				# 初始化为全零
nn.init.ones_(weight)				# 初始化为全一
nn.init.eye_(weight)				# 对二维矩阵，初始化为单位矩阵
nn.init.orthogonal_(weight) 		# 对二维矩阵，初始化为正交矩阵
nn.init.constant_(weight, val) 		# 初始化为指定常数
nn.init.uniform_(weight, a=0, b=1) 	# 初始化为均匀分布 U(a,b)
nn.init.normal_(weight, mean=0, std=1) # 从给定均值和标准差，初始化为正态分布；
nn.init.xavier_uniform_(weight) 	# 使用 Xavier 初始化生成均匀分布，范围由扇入、扇出确定
nn.init.xavier_normal_(weight)		# 使用 Xavier 初始化生成正态分布，方差由扇入、扇出确定
nn.init.kaiming_uniform_(weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
nn.init.kaiming_normal_(weight, a=0, mode='fan_in', nonlinearity='leaky_relu') 
# 使用 He 初始化，范围由 mode 决定，可选扇入或扇出，a 为这层后 Leaky ReLU 的斜率系数，rulu 则不需要

调用上述内置的初始化器，我们可以实现：

def init_weights(m):
    if type(m) == nn.Linear: # 仅当网络层的类型匹配时采用
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.zeros_(m.bias)
net.apply(init_weights) # 对所有层应用初始化方法

注意到上面的方法要求网络类型完全匹配，而 PyTorch 中自带有很多相似的层，如 Conv2d、ConvTranspose2d 等，此时可以获取类名进行字符串匹配：

def init_weights(m):
    classname = m.__class__.__name__    	# 返回 m 的类名
    if classname.find('Conv') != -1:		# 如果类名包含 Conv
    	nn.init.kaiming_uniform_(m.weight)
        nn.init.uniform_(weight, a=-1, b=1)
    elif classname.find('BatchNorm') != -1:           
        nn.init.normal_(m.weight, 1.0, 0.02)
        nn.init.zeros_(m.bias)
net.apply(init_weights)

训练阶段

假设我们已经定义好 MyDataset 类和 MyModel 类，现在正式进入训练阶段。除了实例化数据集和模型，还需要额外定义损失函数和优化器。下面给出一个完整的训练过程：

# 实例化数据集和模型
dataset = MyDataset(file)
train_set = DataLoader(dataset, 16, shuffle=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device) 				# 将模型放到指定的设备上
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), 0.1)
# 训练过程
model.train()								# 将模型切换到「训练模式」，也可以放到循环内
for epoch in range(num_epochs):
	for X, y in train_set:
		optimizer.zero_grad()				# 每个 batch 前都要清空梯度
        X, y = X.to(device), y.to(device)	# 将数据放到指定的设备上
        pred = model(X)
        loss = criterion(pred, y)			# 计算损失
        loss.backward()						# 反向传播
        optimizer.step()					# 每个 batch 更新一次参数

所谓的「训练模式」，是相对另一个「评估模式」下的模型传播路径。实现上，模型通过 model.train() 将变量 self.training 置为 True，通过 model.eval() 将 self.training 置为 False，所以即便训练与测试共用一个模型，也能通过该变量来区分现在属于训练还是测试。

当模型中使用了 Dropout、BatchNorm 等模块，在训练阶段和测试阶段的操作不同，此时区分模式就很有必要。此外，如果网络中有一部分层只在训练阶段启用，则可以用 if self.training: 来执行。

损失函数

PyTorch 在 torch.nn 模块中已经定义了大部分常用的 Loss，具体的计算过程详见 手册：

损失函数	名称	适用场景
nn.L1Loss()	绝对值误差损失	回归
nn.MSELoss()	均方误差损失	回归
nn.SmoothL1Loss()	Huber 损失（结合 L1 和 L2）	回归
nn.BCELoss()	二分类交叉熵损失	二分类
nn.BCEWithLogitsLoss()	自带 Sigmoid 的 BCELoss	二分类
nn.CrossEntropyLoss()	交叉熵损失	多分类
nn.NLLLoss()	负对数似然函数损失	多分类
nn.NLLLoss2d()	图像负对数似然函数损失	图像分割
nn.KLDivLoss()	KL 散度损失	分布学习
nn.MarginRankingLoss()	排序相似度的损失	推荐排序
nn.MultiLabelMarginLoss()	多标签分类的损失	多标签分类
nn.SoftMarginLoss()	多标签二分类问题的损失	多标签二分类

大部分损失函数的用法都是先实例化，再输入预测值和标签值，预测值形状为 (batch_size, num_class)，对应每个类的概率，标签值形状为 (batch_size, num_class) 或 (batch_size, 1)，前者代表标签平滑或多标签，后者仅仅只有目标类的索引。具体用法如下：

criterion = nn.MSELoss()
loss = criterion(pred, target) # 输出的 loss 是零维张量

其中 nn.CrossEntropyLoss() 的 target 输入必须是 LongTensor 类型，可以使用 .long() 进行类型强转。而其他 Loss 则是采用 Float 作为目标输入。

需要注意所有损失函数都会自动对 Loss 取均值，原因是默认参数 reduction='mean'，可以修改成 'sum' 进行求和，或 'none' 保留 (batch_size, 1) 形状。这会影响到梯度的大小，用 mean 时，其学习率大小应为用 sum 时的 batch_size 倍。有时候自己写的损失函数，也会在调用时加上 .mean() 来对齐结果。

此外，有的损失函数会有 weight 参数，BCEWithLogitsLoss() 有 pos_weight 参数，这些都是用来帮助模型解决样本类别不均衡问题的。需要传入张量参数 orch.from_numpy(np.array([1,15])).float()。

优化器

PyTorch 在 torch.optim 中定义了十多种优化器，这里介绍常用的几种：

'''
	params: 待优化的模型参数、或者是模型中参数组的字典（定制不同策略）
	lr: 学习率起始值，给定常数即可，如 1e-3, 5e-6
	momentum: 动量因子，默认为 0，通常设置为 0.9 0.8
	weight_decay: 权重衰减（L2惩罚的系数），默认为 0，通常设置为 1e-4 1e-3 再微调
	nesterov: 是否采用牛顿动量（NAG），较少设置
	### 变化量（速度） v = mo * v + dx, 参数更新 x -= lr * v
'''
optim.SGD(params, lr, momentum=0, weight_decay=0, nesterov=False)
'''
	lr: 全局学习率，设定后就不再改变，默认为 0.1 一般不修改（会自适应）
	### 根据梯度的大小自适应学习率，梯度大则学习率小，更新 x -= lr * dx / sqrt(dx^2)
'''
optim.Adagrad(params, lr=0.01, weight_decay=0)
'''
	alpha: 避免学习率下降太快，采用滑动平均的保留比例，默认为 0.99
	eps: 为了增加数值计算的稳定性而加到分母里的项，不用理会
	### 缓存 cache = α * cache + (1-α) * dx^2, 更新 x -= lr * dx / sqrt(cache)
'''
optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0)
'''
	betas: 在缓存的基础上再对梯度加上记忆，推荐参数 (0.9, 0.999)
	### 记忆 memory = β1 * memory + (1 ‐ β1) * dx
	### 缓存 cache = β2 * cache + (1 - β2) * dx^2
	### 更新 x -= lr * memory / sqrt(cache)
'''
optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
optim.Adamax(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) # 变体1
optim.SparseAdam(params, lr=0.001, betas=(0.9,0.999), eps=1e-08) # 变体2，适用于稀疏张量
optim.AdamW(params, lr=0.001, betas=(0.9,0.999), eps=1e-08, weight_decay=0.01) #变体3

具体用法如下：

import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9) # 用于模型训练
optimizer = optim.Adam([var1, var2], lr=0.0001) # 也可以用于任何变量

optimizer.zero_grad()	# 把上一轮 loss 关于 weight 的梯度变成 0，防止自动累计
optimizer.step()		# 进行一次优化迭代，即 x += Δ 的过程

如果想对网络中不同部分采取不同的优化策略，则需要用到参数组（param_group），其保存了各个参数及其对应的学习率、动量等，可以通过列表嵌套字典的形式来设置：

# 默认情况，一个参数组
optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
# 两个参数组，此时 len(optim.param_gruops) == 2
optim.SGD([{'params': model.base.parameters(), 'weight_decay': wd},
           {'params': model.classifier.parameters(), 'lr': 1e-3}], # 访问模型局部参数即可
          lr=1e-2, momentum=0.9)

调度器

注意到上述优化器参数都有一个固定的 lr，即使通过不同策略使学习率改变，也会有一个全局 lr。为了更好地收敛模型，PyTorch 在 torch.optim.lr_scheduler 里封装了主动进行学习率衰减的方法。通过一个调度器（Scheduler）来控制优化器，具体用法如下：

optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=1e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, verbose=True)

for epoch in range(num_epochs):
	for X, y in train_set:
		optimizer.zero_grad()
        pred = model(X)
        loss = criterion(pred, y)
        loss.backward()
        optimizer.step()
    scheduler.step()					# 调度器工作，每个 epoch 调度一次

当然还有许多不同的调度器，可以有规律地间隔调整，或者根据指标的变化情况调整。参数 verbose=True 表示每次更新都会输出一条信息，适用于所有调度器：

# 等间隔调整，每隔 step_size 个 epoch，调整为 lr * gamma
optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1)
# 多间隔调整，milestones 传入参数列表，表示在第几个 epoch 更新
optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1)
# 指数衰减调整，每个 epoch 更新为 lr * gamma ^ epoch，gamma 一般设为 0.9
optim.lr_scheduler.ExponentialLR(optimizer, gamma)
# 余弦退火函数调整，走势类似 cos(x)，当 epoch = T_max 时，学习率取最低点 eta_min
optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0)

# 根据指标调整，当给定指标在最近 patience 个 epoch 中都没有变化，学习率下降
optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10)
scheduler.step(val_loss) 	# loss 对应 mode='min'，表示当指标不再降低时更新
scheduler.step(val_acc) 	# acc 对应 mode='max'，表示当指标不再升高时更新

# 自定义倍率调整，通过匿名函数为不同参数组设置不同规则，更新为 lr * lambda(epoch)
optimizer = optim.SGD([{'params': base_params},
        			   {'params': net.fc.parameters(), 'lr': 1e-2}],
                      lr=1e-3, momentum=0.9, weight_decay=1e-4)
lambda1 = lambda epoch: epoch // 3
lambda2 = lambda epoch: 0.95 ** epoch
scheduler = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2])

评估阶段

前文提到，「评估模式」是相对于「训练模式」的模型传播路径，通常在验证集、测试集上推理时会开启。开启命令为 model.eval()，也可以用 model.train(False)。在该模式下，Dropout、BatchNorm 等模块会有所区别。并且模型不需要进行反向传播，也就无需进行梯度计算，此时使用 torch.is_grad_enabled() 会返回 False。

需要注意的是，为了防止 PyTorch 自动对张量进行梯度计算，要用 with torch.no_grad(): 包裹评估阶段的代码。关闭梯度计算后模型的推理速度也会大大增加。下面介绍具体用法。

• 验证过程：通常与训练阶段一起运行，即每个 epoch 训练完都在验证集上验证，并计算 Loss：

for epoch in range(num_epoch):
	model.train()
    for X, y in tr_set:
        # 训练过程略去

    model.eval()										# 打开评估模式
    with torch.no_grad():								# 关闭梯度计算
        total_loss = 0
        for X, y in val_set:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            loss = criterion(pred, y)
            total_loss += loss.cpu().item() * len(X) 	# loss 会默认取 mean

• 测试阶段：无需优化、无需计算损失，只需前向推理预测答案即可：

model.eval()
preds = []
with torch.no_grad():
    for X in test_set:
        X = X.to(device)
        pred = model(X)
        preds.append(pred.cpu())

TensorBoard 可视化

TensorBoard 用于可视化网络的相关信息，包括训练过程的指标、输出输出、网络内部参数等。使用前需要安装 pip install tensorboard，并在代码中导入 Writer。Writer 会将我们关心的数据保存在一个文件夹中，在显示到浏览器上。

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./path/to/log') 	# 实例化 writer
writer.close()								# 关闭 writer，代码结束时用
'''
	tag/main_tag: 可视化时图像的标签名，是图的唯一标识
	scalar_value: 代码中的变量名，必须是一个标量
	tag_scalar_dict: 传入变量字典，key 相当于可视化的 tag，value 是代码中的变量名
	global_step: 可视化时的横坐标，传入 epoch 就表示随 epoch 的指标变化
'''
writer.add_scalar(tag, scalar_value, global_step=None, walltime=None)
writer.add_scalars(main_tag, tag_scalar_dict, global_step=None, walltime=None)
'''
	img_tensor: 代码中的变量名，可以传入张量，如输入数据、各层的输入输出
	dataformats: 数据格式，如 CHW，HWC，HW
'''

writer.add_image(tag, img_tensor, global_step=None, walltime=None, dataformats='CHW')
'''
	model: 传入模型，必须是 nn.Module 类
	input_to_model: 输入给模型的数据
	verbose: 是否打印计算图结构信息
'''
writer.add_graph(model, input_to_model=None, verbose=False, use_strict_trace=True)

运行上述代码后，数据将被保存到文件夹中，接下来在终端中输入：

tensorboard --logdir=./path/to/log --port 6006

就可以在 localhost:6006 看到可视化结果。如果是在远程服务器中，则需要将服务器的 6006 端口转发到本机的一个端口，在本机浏览器中访问端口。

模型保存与加载

PyTorch 文件后缀为 .pt、.pth 和 .pkl，三者在使用上没有区别，都可以保存张量、模型（结构、权重参数），张量的保存和加载方法如下：

torch.save(t, 'save_tensor.pt')		# 保存张量包括其 device 属性
t = torch.load('save_tensor.pt')	# 是否加载到 GPU 由 save 之前的参数决定
t = torch.load('save_tensor.pt', map_location=torch.device('cpu'))	# 加载到 CPU
t = torch.load('save_tensor.pt', map_location="cuda:0")	# 加载到指定 GPU

如果有多个张量，可以通过字典的形式进行存储，将字符串映射到张量，这也是模型的保存形式：

tensor_dict = {'x': x, 'y': y}
torch.save(tensor_dict, 'save_dict.pt')
tensor_dict2 = torch.load('save_dict.pt')

模型（Model）分为网络的结构和权重参数两部分，后者位于 model.stack_dict() 中，以字典的形式存储，将每一层映射成它的参数张量，通过 .keys() 可以查看所有键。模型的保存和加载方法有两种：

# 只保存模型的权重参数，不保存模型结构，速度更快
torch.save(model.state_dict(), 'save_param.pt')
model = Model(*args, **kwargs)	# 需要重新 init 模型，模型是否加载到 GPU 由参数决定
model.load_state_dict(torch.load('save_param.pt', map_location="cuda:0"))	# 参数加载
model = model.to(device)
model.eval()

# 保存整个模型，包括网络结构和权重参数
torch.save(model, 'save_model.pt')
model = torch.load('save_model.pt')	# 模型是否加载到 GPU 由 save 之前的参数决定
model = model.to(device)
model.eval()

注意第二种方法中，网络模型定义的代码要与 load 代码写在一起，或者从其他文件中 import，这样才能成功加载。

Checkpoint

在适当的时候保存模型的断点，可以在训练意外终止后不必从头开始训练，也方便必要的时候回滚到历史状态。除了模型的参数，还包括 Loss、Epoch 等训练信息。最容易忽略的是优化器的参数，也是使用 state_dict() 访问，包括：学习率、动量值、衰减系数等参数。一般使用 tar 文件格式来保存这些检查点。

if (epoch + 1) % checkpoint_interval == 0:
    checkpoint = {'epoch': epoch+1,
                  'loss': loss,
                  'best_loss': best_loss,
                  'model_state_dict': net.state_dict(),
                  'optimizer_state_dict': optimizer.state_dict(),
                 }
    path_checkpoint = "./history/ckp_{}_epoch.pth.tar".format(epoch)
    torch.save(checkpoint, path_checkpoint)

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(path_checkpoint)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
best_loss = checkpoint['best_loss']

model.eval()

切换设备

load 提供设备重载的功能，可以在 CPU、GPU 之间切换，也可以在不同 GPU 间切换：

# 强制所有 GPU 张量加载到 CPU 中
torch.load('tensors.pt', map_location=lambda storage, loc: storage)
torch.load('tensors.pt', map_location=torch.device('cpu')) 
# 把所有的张量加载到 GPU 1 中
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
torch.load('tensors.pt', map_location='cuda:1')
# 把张量从GPU 1 移动到 GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

上述代码只有模型在一个 GPU 上训练时才生效，如果在多个 GPU 中，则需要使用 nn.DataParallel 模块，例如：

model = nn.DataParalle(model) # 自动将任务发送到所有 GPU 上执行
torch.save(model.module.state_dict(), 'model.pt')

注意到此时模型的参数名都带上了 module 前缀，如果再想把模型加载到单个 GPU 或 CPU 上，则需要手动将 key 中的前缀去除：

# 原始通过 DataParallel 保存的文件
state_dict = torch.load('model.pt')

# 创建一个不包含 module. 的新 OrderedDict
from collections import OrderedDict
new_state_dict = OrderedDict()
for key, value in state_dict.items():
    name = key[7:] # 去掉 module.
    new_state_dict[name] = value

# 加载参数
model.load_state_dict(new_state_dict)

加载部分模型

在迁移学习或者训练新的复杂模型时，加载部分模型是很常见的。利用经过训练的参数，即使只有少数参数可用，也将有助于预热训练过程，并且使模型更快收敛。

在加载部分模型参数进行预训练的时候，很可能会碰到键不匹配的情况（模型权重都是按键值对的形式保存并加载回来的）。因此，无论是缺少键还是多出键的情况，都可以通过在 load_state_dict() 函数中设定 strict 参数为 False 来忽略不匹配的键。

torch.save(modelA.state_dict(), PATH)
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False) # 忽略不匹配的键

如果想将某一层的参数加载到其他层，但是有些键不匹配，可以通过强制修改 state_dict 中参数的 key 解决。