pytorch中提升训练速度的方法

最新推荐文章于 2023-12-13 15:25:30 发布

反科研pua所所长

最新推荐文章于 2023-12-13 15:25:30 发布

阅读量574

点赞数

分类专栏： python(pytorch为主) 文章标签： pytorch 深度学习神经网络

本文链接：https://blog.csdn.net/YasmineC/article/details/120564051

版权

36 篇文章 3 订阅

订阅专栏

想用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)

关注

专栏目录