pytorch中提升训练速度的方法

想用COCO_Captions跑我模型的预训练过程，正赶上出省开会三四天，动身前用实验室的3090开始训练，结果回来之后才跑了3个epoch（预备跑50个），我崩溃啊…

于是去跟参考论文的作者做了个比较，作者使用8块V100，每块的batch_size设置为32，40个epoch一共跑了3/4天，我的速度大概比作者慢了10倍吧…于是采用了以下解决方案：

• 转为半精度训练，根据这篇文章的实验，使用单精度训练时，V100的速度要略慢于3090，而换成半精度训练，V100的速度就变了3090的1.5倍，同时减少了内存的消耗，但并没有降低最终任务的准确度（应该会有精度损失的，但是在可接受范围内）。转换的方式是： 首先，在上下文管理器torch.cuda.amp.autocast的实例下计算模型输出和损失：

 autocast = torch.cuda.amp.autocast
 with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

然后，添加损失缩放以保留较小的梯度值：

scaler = torch.cuda.amp.GradScaler()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

大家只需要跟着文章的内容，把上面的代码放到合适的位置就可以了。不过如果使用GPU训练或者是多GPU并行训练的话，需要将输入和损失部分改为：

outputs = model(inputs).to(device)
loss = criterion(outputs, labels).to(device)