GradScaler¶

class GradScaler(init_scale=2.0 ** 4, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000)[源代码]¶

在 autocast 模式下执行梯度缩放以防止数据溢出的helper类。

参数

init_scale (float) – 初始缩放因子。
growth_factor (float) – 在实际 update 阶段 scale 乘以的因子。如果 growth_factor 为0，scale_factor 将不会更新。
backoff_factor (float) – 遇到梯度溢出时 scale 乘以的因子。
growth_interval (int) – 两个 scale 更新阶段之间的间隔。

示例

gm = GradManager()
opt = ...
scaler = GradScaler()

gm.attach(model.parameters())

@autocast()
def train_step(image, label):
    with gm:
        logits = model(image)
        loss = F.nn.cross_entropy(logits, label)
        scaler.backward(gm, loss)
    opt.step().clear_grad()
    return loss

如果需要更灵活地使用，可以拆分 scaler.backward 成三行：

@autocast()
def train_step(image, label):
    with gm:
        logits = model(image)
        loss = F.nn.cross_entropy(logits, label)
        gm.backward(loss， dy=megengine.tensor(scaler.scale_factor))
    scaler.unscale(gm.attached_tensors())
    scaler.update()
    opt.step().clear_grad()
    return loss

可用于为多个 batch 积累梯度。

backward(gm, y=None, dy=None, *, unscale_grad=True, update_scale='if_unscale_grad')[源代码]¶

backward 的包装器，用于缩放 y 的梯度和反缩放参数的梯度。

参数

gm (GradManager) – 要包装的 GradManager。
y (Union[Tensor, List[Tensor], None]) – 和 GradManager backward 的 y 一样。
dy (Union[Tensor, List[Tensor], None]) – 与 GradManager backward 的 dy 相同。将乘以 scale_factor。
unscale_grad (bool) – 是否同时进行 unscale 。如果需要累积梯度，可能为 False。
update_scale (bool) – 和 unscale 的 update 一样。如果 unscale_grad 为 False 将被忽略。

unscale(grad_tensors)[源代码]¶

反缩放所有 grad_tensors 的梯度。

参数: grad_tensors (Iterable[Tensor]) – 需要反缩放梯度的张量。应该是所有在 GradManager 反向中受 target 张量影响的张量。

update(new_scale=None)[源代码]¶: 根据是否遇到梯度溢出来更新 scale 因子。如果提供了new_scale，内部更新机制将被忽略。

megengine.amp

autocast