Pytorch技法：混合精度训练（一）

写在前面：由于近期实验室显卡供不应求，同步的训练进程过多，导致训练和推理速度太慢，遂进行了混合精度训练的相关尝试。

• 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只是对原来的代码进行了简单的更改。具体理论等待后续更新。