写在前面:由于近期实验室显卡供不应求,同步的训练进程过多,导致训练和推理速度太慢,遂进行了混合精度训练的相关尝试。
- pytorch混合精度训练代码
- 这里意思就是在训练的时候用上你机器的FP16运算,不仅用默认的FP32。
self.scalar = torch.cuda.amp.GradScaler() #1.这是你需要添加的
for epoch in range(self.trans_epochs):
for data_tensor,label_tensor in self.train_generator:
data_tensor = data_tensor.cuda()
data_label = label_tensor.cuda()
self.optimizer.zero_grad()
with torch.cuda.amp.autocast(): #2.这是你需要添加的
y_pred = model(data_tensor)
loss = loss_f(y_pred, label_tensor)
self.scalar.scale(loss).backward() #3.这是你需要添加的,进行损失缩放工作
self.scalar.step(self.optimizer) #4.这是你需要添加的
self.scalar.update() #5.这是你需要添加的
- 注意:仅注释部分代码是进行更新的,相较于非混合精度训练,实际上只多了三行代码,也就是我在上面标注的1,2,5,其中3和4只是对原来的代码进行了简单的更改。具体理论等待后续更新。