Pytorch使用GradScaler进行训练加速

起因是一次参考一个github项目时，发现该项目训练和验证一个epoch耗时30s，而我的项目训练和验证一个epoch耗时53s，当训练多个epoch时，这个差异就很大了。通过研究发现github项目使用了GradScaler来进行加速，所以这里总结一下。

1、Pytorch的GradScaler

GradScaler在文章Pytorch自动混合精度(AMP)介绍与使用中有详细的介绍，也即是如果tensor全是torch.float32，计算成本会大一些，但是单精度（FP32）与半精度(FP16)结合在一起，在保证精度的情况下，会加快训练。想看原理的去上面这篇文章，这里总结一下使用。

2、如何使用

下面给出代码，省略了一下与GradScaler无关的代码，同时#标注的是于GradScaler不想关的代码。

from torch.cuda.amp import GradScaler

scaler = GradScaler()	# 实例化对象

for imgs,labels in dataloader:
	#imgs=imgs.to(device)		# Tensor:(16,10,1,40,40) torch.float32
	#labels=labels.to(device)		# Tensor:(16,) torch.int64
	
	with autocast():
		#predict = model(imgs)
		#loss = criterion(predict, labels)
	
		scaler.scale(loss).backward()	# 计算梯度
		scaler.step(optimizer)	# 调整lr
        scaler.update()		# 更新梯度
        # optimizer.zero_grad()	 # 梯度清零(不涉及复杂计算，所以不需要GradScaler参与)