使用基础 PyTorch 进行混合精度训练

您将使用低精度浮点数据类型来加速语言翻译模型的训练。例如，16 位浮点数据类型（float16）只有其 32 位对应类型（float32）的一半大小。这会加速矩阵乘法和卷积的计算。回顾一下，这需要对梯度进行缩放，并将运算转换为 16 位浮点。

已预加载以下对象：dataset、model、dataloader 和 optimizer。

本练习是课程的一部分

使用 PyTorch 高效训练 AI 模型

查看课程

练习说明

在循环之前，使用 torch.amp.GradScaler 为梯度定义一个 scaler。
在循环中，使用 torch.autocast 作为上下文管理器，将运算转换为 16 位浮点数据类型。
在循环中，对损失进行缩放并调用 .backward()，以创建缩放后的梯度。

交互式实操练习

通过完成这段示例代码来试试这个练习。

# Define a scaler for the gradients
scaler = torch.amp.____()
for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    # Casts operations to mixed precision
    with torch.____(device_type="cpu", dtype=torch.____):
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
    # Compute scaled gradients
    scaler.____(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

编辑并运行代码