深度学习笔记【实践篇】二、基本训练框架构建
2020/3/5大约 4 分钟
这篇整理了用Pytorch构建基本的训练框架。
相较于之前接触过的TensorFlow的"标准化"，整体感觉Pytorch的框架似乎更加自由？教材、B站网课的几位大佬写的代码风格会有些差异，我按自己的理解，整理了我浅显认知中的Pytorch基本框架。
框架组成部分
1. 模型架构定义
【推荐】一般用class继承nn.Module来定义模型，需要重写__init__和forward方法。
class Model(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5) ... # 定义好每一层 def forward(self, x): x = self.conv1(x) ... return x【不推荐】纯线性模型也可以直接使用nn.Sequential（本身也继承自nn.Module）。
优点：
- 不用写forward，简单、适合做demo
缺点：
- 不适合非线性模型
- 不好观察内部以及调试
model = nn.Sequential( nn.Conv2d(1, 6, 5, padding=2), ..., nn.Linear(84, 10) )【也还行】结合一下1和2
class Model(nn.Module): def __init__(self): # 网络结构定义 super().__init__() self.backbone = nn.Sequential(...) self.head = nn.Sequential(...) def forward(self, x): x = self.backbone(x) return self.head(x)模型权重初始化(在3的基础上)
class Model(nn.Module): def __init__(self): super().__init__() self.backbone = nn.Sequential(...) self.head = nn.Sequential(...) ... # 定义好每一层 # 网络参数初始化 self.apply(self.init_weights) @staticmethod def init_weights(m): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) if m.bias is not None: nn.init.constant_(m.bias, 0) def forward(self, x): x = self.backbone(x) return self.head(x)网络参数初始化原因：
如果初始化不合适（例如纯随机），网络可能：
- 学得极慢
- 梯度爆炸 / 消失
- 甚至根本学不会（loss 不降）
如果初始化合适（例如kaiming），网络可能：
- loss 一开始就下降
- 梯度稳定
- 学习率可设大一点
- 收敛快
场景 推荐初始化 典型写法 Conv + ReLU Kaiming (He) kaiming_normal_Linear + ReLU Kaiming kaiming_normal_Linear + Sigmoid/Tanh Xavier xavier_normal_分类头（FC） Normal(0, 0.01) normal_BatchNorm 默认 不手动初始化
2. 数据集准备
两种方式导入：
- a. 自定义数据集 -> 继承
torch.utils.data.Dataset - b. 公开数据集 -> 用现有的包，例如
torchvision.datasets
- a. 自定义数据集 -> 继承
数据集导入后，使用
Data.random_split进行数据集划分一般train:val=8:2，或者train:val:test=8:1:1
使用
DataLoader进行数据加载
3. 模型训练
设定训练设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) # cpu或者单卡设置优化器
具体原理和差异，见深度学习笔记【基础篇】二、损失函数
常见的两种优化器和参数如下：
Momentum:
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)Adam:
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)统计参数的初始化
训练集、验证集的损失列表（一个item就是一个epoch的平均损失）等
正式训练
伪代码如下：
for epoch in range(num_epochs): ... # 初始化训练集、验证集的loss，数量，准确数等统计参数 # ====== training ====== # 设置模型为训练模式 model.train() for batch_idx, (b_x, b_y) in enumerate(train_loader): # 注意：数据也要像模型一样，放到设备上！ b_x, b_y = b_x.to(device), b_y.to(device) # 前向传播 output = model(b_x) # 计算损失 loss = loss_fn(output, b_y) # 清零梯度。pytorch会自动累积梯度，所以每次前向传播前都要清零 optimizer.zero_grad() # 反向传播，计算梯度 loss.backward() # 更新参数 optimizer.step() # 更新统计参数 ... # ====== validation ====== # 设置模型为验证模式 model.eval() # 用with torch.no_grad()包裹验证过程。临时关闭梯度追踪，并禁用autograd with torch.no_grad(): for batch_idx, (b_x, b_y) in enumerate(val_loader): # 注意：数据也要像模型一样，放到设备上！ b_x, b_y = b_x.to(device), b_y.to(device) # 前向传播 output = model(b_x) # 计算损失 loss = loss_fn(output, b_y) # 更新统计参数 ... # 选择最优模型，并保存 torch.save(model.state_dict(), "model.pth") # 返回统计参数，用于画图注意点：
- model.train()，设置模型为训练模式。训练模式会：
- 启用Dropout：随机丢弃神经元
- 启用BatchNorm：计算 mini-batch 的均值和方差，并更新running statistics
- .to(device)：将模型和数据都放到设备上
- output = model(b_x)：前向传播
- 但不只是调用forward
- 会启用autograd，构建反向传播所需的计算图，所以后续的loss都不需要赋值给变量了
- with torch.no_grad()
- 优势：不计算梯度，节省显存和加快计算速度
- 使用场景：验证 / 评估 / 推理 阶段
- model.state_dict()
- 这个是模型"可学习状态"的字典表示，里面包含了params（weights和biases），buffers（比如BatchNorm的running_mean和running_var）
- model本身是python对象结构，与代码不解耦、不轻量
- 配套用法：
- 保存模型状态：torch.save(model.state_dict(), "model.pth")
- 加载模型状态：model.load_state_dict(torch.load("model.pth"))
- model.train()，设置模型为训练模式。训练模式会：