使用apex进行半精度计算

from apex import amp

# 定义网络和优化器
model = ...
optimizer = ...

# 将优化器转换为半精度
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")

# 正常的训练循环
for inputs, labels in data_loader:
    # 前向传播
    outputs = model(inputs)

    # 计算损失
    loss = ...

    # 反向传播和优化
    optimizer.zero_grad()
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_loss.backward()
    optimizer.step()

其中只有一个opt_level需要用户自行配置：

O0：纯FP32训练，可以作为accuracy的baseline；

O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（对 Tensor Core 友好的操作, 如GEMM, 卷积）还是FP32（Softmax）进行计算。

O2：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。

O3：纯FP16训练，很不稳定，但是可以作为speed的baseline；

模型先放GPU

再用apex.initialize进行混合精度初始化

最后再并行化，顺序不能错