GradScaler

class GradScaler(init_scale=2.0 ** 4, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000)

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

参数:

• init_scale (float) – Initial scale factor.

• growth_factor (float) – Factor that the scale is multiplied by in actual update stage. If growth_factor is 0, scale_factor will not update.

• backoff_factor (float) – Factor that the scale is multiplied by when encountering overflow grad.

• growth_interval (int) – The interval between two scale update stages.

示例

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，内部更新机制将被忽略。