训练提速60%！只需5行代码，PyTorch 1.6即将原生支持自动混合精度训练

self.train() X = torch.tensor(X, dtype=torch.float32) y = torch.tensor(y, dtype=torch.float32)   optimizer = torch.optim.Adam(self.parameters(), lr=self.max_lr) scheduler = torch.optim.lr_scheduler.OneCycleLR(     optimizer, self.max_lr,     cycle_momentum=False,     epochs=self.n_epochs,     steps_per_epoch=int(np.ceil(len(X) / self.batch_size)), ) batches = torch.utils.data.DataLoader(     torch.utils.data.TensorDataset(X, y),     batch_size=self.batch_size, shuffle=True )   # NEW scaler = torch.cuda.amp.GradScaler()   for epoch in range(self.n_epochs):     for i, (X_batch, y_batch) in enumerate(batches):         X_batch = X_batch.cuda()         y_batch = y_batch.cuda()         optimizer.zero_grad()           # NEW         with torch.cuda.amp.autocast():             y_pred = model(X_batch).squeeze()             loss = self.loss_fn(y_pred, y_batch)           # NEW         scaler.scale(loss).backward()         lv = loss.detach().cpu().numpy()         if i % 100 == 0:             print(f"Epoch {epoch + 1}/{self.n_epochs}; Batch {i}; Loss {lv}")           # NEW         scaler.step(optimizer)         scaler.update()          scheduler.s 

新的 PyTorch GradScaler 对象是 PyTorch 实现的损失缩放。回想一下在“混合精度如何工作”一节中提到，在训练期间，为了防止梯度变小到0，某种形式的缩放是必要的。最佳的损失乘数得足够高以保留非常小的梯度，同时不能太高以至于导致非常大的梯度四舍五入到 inf产生相反的问题。

PyTorch使用指数退避（exponential backoff）来解决这个问题。Gradscalar 以一个小的损失乘数开始，这个乘数每次会翻倍。这种逐渐加倍的行为一直持续到 GradScalar 遇到包含 inf 值的梯度更新。Gradscalar 丢弃这批数据(例如跳过梯度更新) ，将损失乘数减半，并重置其倍增时间。

通过这种方式逐级上下移动损失乘数，PyTorch 可以随着时间的推移近似得到合适的损失乘数。熟悉 TCP 拥塞控制的读者应该会发现这里的核心思想非常熟悉！该算法使用的准确数字是可配置的，你可以直接从docstring中看到默认值:

torch.cuda.amp.GradScaler(     init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5,     growth_interval=2000, enabled=True )  

Gradscalar 需要对梯度更新计算(检查是否溢出)和优化器(将丢弃的batches转换为 no-op)进行控制，以实现其操作。这就是为什么 loss.backwards()被 scaler.scale(loss).backwards()取代， 以及 optimizer.step()被 scaler.step(optimizer)替换的原因。

值得注意的是，GradScalar 可以检测并停止overflows(因为 inf 总是坏的) ，但是它无法检测和停止underflows(因为0通常是一个合法值)。如果你选择的初始值太低，增长间隔太长，你的网络可能会在 GradScalar 介入之前underflow并发散。由于这个原因，选择一个非常大的初始值可能是一个好主意。

最后，注意 GradScalar 是一个有状态对象。使用此功能保存模型checkpoint需要和模型权重一起写入和读取磁盘。用 state _ dict 和 load _ state _ dict 对象方法(在 PyTorch 文档中有介绍)可以很容易地做到这一点。

自动混合精度训练拼图的另一半是 torch.cuda.amp.autocast 上下文管理器。Autocast实现了 fp32-> fp16转换。回想一下“混合精度是如何工作的“中的内容，由于不同的操作以不同的速率累积误差，并非所有的操作都可以在 fp16中安全运行。下面的截图来自 amp 模块文档，介绍了autocast如何处理 PyTorch 中可用的各种操作:

这个列表主要由矩阵乘法和卷积两部分组成，还有简单的线性函数。

这些操作在 fp16中是安全的，但是在输入有 fp16和 fp32混合的情况下，这些操作具有向上适配(up-casting)规则，以确保它们不会出问题。注意，这个列表还包括另外两个基本的线性代数运算: 矩阵/向量点积和向量叉积。

对数、指数、三角函数、正规函数、离散函数和(大)和在 fp16中是不安全的，必须在 fp32中执行。

通过浏览这个列表，在我看来，大多数层都会从autocasting中受益，这要归功于它们内部对基本线性代数操作的依赖，但大多数激活函数却不是。卷积层是最大赢家。

启用sutocasting非常简单。你只需要做的就是使用autocast上下文管理器包好模型的正向传播:

（编辑：南平站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!