解决Pytorch半精度浮点型网络训练的问题

解决 Pytorch 半精度浮点型网络训练的问题需要注意以下几点：

1. 使用合适的半精度浮点类型
2. 防止数值溢出
3. 对于早期的 Pytorch 版本，需要额外安装 apex 库

下面我会详细讲解具体的攻略。

使用合适的半精度浮点类型

Pytorch 提供了两种半精度浮点类型：torch.float16 和 torch.bfloat16，前者占用 16 位，后者占用 16 位，但精度会更接近于单精度浮点型。

根据模型和数据的特点，选择合适的半精度浮点类型很重要。如果模型中有很小的值或需要准确计算的地方，建议选择 torch.bfloat16。如果需要减少内存占用，可以选择 torch.float16。

在 Pytorch 中使用半精度浮点型可以通过以下方式：

# 定义模型时使用半精度浮点型
model = Model().half()

# 将数据转换为半精度浮点型
input_data = input_data.half()

# 定义优化器时设置半精度浮点型
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9).half()

# 在训练过程中，需要将 loss 值转换为单精度浮点型
loss = criterion(output.float(), target)

防止数值溢出

在使用半精度浮点型进行训练时，由于精度的限制，可能会出现数值溢出的情况。为了解决这个问题，可以使用以下方法：

1. 改变梯度的缩放因子

训练时可以将梯度缩小一定的因子，避免数值溢出。一般来说，梯度缩放因子的选择是当前 batch 数据中绝对值最大的值。

# 计算梯度缩放因子
clip_norm = torch.tensor(0.1)
total_norm = 0
for p in model.parameters():
    if p.grad is not None:
        param_norm = p.grad.data.norm(2)
        total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
clip_norm = clip_norm / torch.max(clip_norm, total_norm)
clip_norm_item = clip_norm.item()

# 缩放梯度
for p in model.parameters():
    if p.grad is not None:
        p.grad.data.mul_(clip_norm)

1. 改变优化器的参数

在使用半精度浮点型进行训练时，可以尝试调整优化器的参数，比如设置 loss_scale 参数。

# 修改优化器的方法
# 定义优化器时设置 loss_scale 参数
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4, nesterov=True, loss_scale=128.0)
...
# 在每次前向传播时进行 loss scale
pred = model(inputs)
loss = loss_fn(pred, targets) * loss_scale
...
# 在做 backward 操作前将需要 loss scale
(loss * (1.0 / loss_scale)).backward()
optimizer.step()

1. 使用内置的 Fixup 初始化方法

Pytorch 提供了内置的 Fixup 初始化方法，这个方法可以有效避免数值偏移问题。

# 使用 Fixup 初始化方法
from torch.nn.init import kaiming_normal_
def fixup_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        kaiming_normal_(m.weight.data)
        if m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
        # bias 乘以 Fixup-fanin 的值
        m.bias.data.mul_(2.0)
        m.weight.data.mul_(1.0 / m.weight.data.reshape(m.weight.data.size(0), -1).std(1, keepdim=True))

model.apply(fixup_init)

1. 梯度累积

在某些情况下，即使使用了以上的技巧，仍然无法解决数值溢出问题。这种情况下可以考虑梯度累积的方法，将 batch_size 改为原来的 n 倍，每次只更新 n 次参数。

batch_size = 16
accum_step = 64 // batch_size  # 累积 64 个样本的梯度
for i, (inputs, targets) in enumerate(trainloader):
    inputs = inputs.to(device)
    targets = targets.to(device)

    optimizer.zero_grad()
    for j in range(accum_step):
        start_index = j * batch_size
        end_index = start_index + batch_size
        pred = model(inputs[start_index:end_index])
        loss = loss_fn(pred, targets[start_index:end_index])
        loss = loss / accum_step
        loss.backward()
    optimizer.step()

安装 apex 库

如果你的 Pytorch 版本比较早，可能需要额外安装 apex 库（链接）。

apex 库中提供了一个叫 amp 的模块，可以用于自动化半精度浮点类型的训练过程。

使用 amp 模块时，只需要将模型、优化器、loss 函数全部使用 amp.initialize 进行包裹即可。

# 修改代码以适应 amp 库
from apex import amp
...

model = Model().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()

model, optimizer = amp.initialize(model, optimizer, opt_level='O1')

for i, (inputs, targets) in enumerate(trainloader):
    inputs = inputs.to(device)
    targets = targets.to(device)

    optimizer.zero_grad()

    # 前向传播
    outputs = model(inputs)

    # 计算 loss
    loss = criterion(outputs, targets)

    # 计算梯度并做反向传播
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_loss.backward()

    # 更新参数
    optimizer.step()

这就是解决 Pytorch 半精度浮点型网络训练的完整攻略，其中包含了选择合适的半精度浮点类型、防止数值溢出、使用内置的 Fixup 初始化方法和使用 apex 库等多种方案。