PyTorch 训练一个简单神经网络的要点

原创，转载请注明

本文属于 PyTorch 入门教程， 主要说在用 PyTorch 训练神经网络时要注意的一些地方。 完整代码请参考 Notebook 。

以下本文主要讨论的代码：

def train(train_loader, model, criterion, optimizer, device):
    '''
    Function for the training step of the training loop
    '''
    model.train()
    running_loss = 0
    
    for X, y_true in train_loader:

        # Pytorch accumulate gradient, so we should zero it at first 
        optimizer.zero_grad()
        
        X = X.to(device)
        y_true = y_true.to(device)
    
        # Forward pass
        y_hat, _ = model(X) 
        loss = criterion(y_hat, y_true) 
        running_loss += loss.item() * X.size(0)

        # Backward pass
        loss.backward()
        optimizer.step()
        
    epoch_loss = running_loss / len(train_loader.dataset)
    return model, optimizer, epoch_loss

1 把模型设置到训练模式

model.train()

这里model 是 nn.Model 派生类的一个实例。 train() 函数将模型设置为训练模式， 在这种模式下， 一些只应该对训练有效的Layer 将会被激活， 比如 dropout， batchnorm 等。 但是如果在部署时， 就不应该打开这些功能。 这时就需要利用

model.eval()

将模型设置为部署模式，让模型发挥最大的推理能力。

2 Gradient 归 0

optimizer.zero_grad()

网络的每个参数， 都有一个grad变量。PyTorch 有时候需要累积grad， 所以它不会主动清理 grad（比如训练 RNN）。在进行后向传播之前， 则需要把所有的 grad 清零。

3 反向传播

loss.backward()

criterion 计算出 loss 以后，会返回loss function 对象 loss， 调用loss.backward() 则会计算 loss 对于每个 x 的 梯度： （只计算 requires_grad=True 的参数的梯度）， 计算结果存储在每个参数x 的 grad 属性中 。用伪代码表示

x.grad += dloss/dx

4 更新网络参数

optimizer.step()

根据梯度， 更新参数, 对于 SDG， 可以表示为：

x += -lr * x.grad