训练提速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 丢弃这批数据(例如跳过梯度更新) ，将损失乘数减半，并重置其倍增时间。

（编辑：网站开发网_安阳站长网）

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