背景
pytorch作为深度学习的计算框架正得到越来越多的应用.
我们除了在模型训练阶段应用外，最近也把pytorch应用在了部署上.
在部署时，为了减少计算量，可以考虑使用16位浮点模型，而训练时涉及到梯度计算，需要使用32位浮点，这种精度的不一致经过测试，模型性能下降有限，可以接受.
但是推断时计算量可以降低一半，同等计算资源下，并发度可提升近一倍
具体方法
在pytorch中，一般模型定义都继承torch.nn.Moudle，torch.nn.Module基类的half()方法会把所有参数转为16位浮点，所以在模型加载后，调用一下该方法即可达到模型切换的目的.接下来只需要在推断时把input的tensor切换为16位浮点即可
另外还有一个小的trick，在推理过程中模型输出的tensor自然会成为16位浮点，如果需要新创建tensor，最好调用已有tensor的new_zeros，new_full等方法而不是torch.zeros和torch.full，前者可以自动继承已有tensor的类型，这样就不需要到处增加代码判断是使用16位还是32位了，只需要针对input tensor切换.
补充：pytorch 使用amp.autocast半精度加速训练
准备工作
pytorch 1.6+
如何使用autocast？
根据官方提供的方法，
答案就是autocast + GradScaler。
如何在PyTorch中使用自动混合精度？
答案：autocast + GradScaler。
1.autocast
正如前文所说，需要使用torch.cuda.amp模块中的autocast 类。使用也是非常简单的
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
|
from torch.cuda.amp import autocast as autocast
# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
for input, target in data:
optimizer.zero_grad()
# 前向过程(model + loss)开启 autocast
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 反向传播在autocast上下文之外
loss.backward()
optimizer.step()
|
2.GradScaler
GradScaler就是梯度scaler模块，需要在训练最开始之前实例化一个GradScaler对象。
因此PyTorch中经典的AMP使用方式如下：
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
|
from torch.cuda.amp import autocast as autocast
# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 前向过程(model + loss)开启 autocast
with autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
|
3.nn.DataParallel
单卡训练的话上面的代码已经够了，亲测在2080ti上能减少至少1/3的显存，至于速度。。。
要是想多卡跑的话仅仅这样还不够，会发现在forward里面的每个结果都还是float32的，怎么办？
|
1
2
3
4
5
6
7
8
|
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, input_data_c1):
with autocast():
# code
return
|
只要把forward里面的代码用autocast代码块方式运行就好啦！
自动进行autocast的操作
如下操作中tensor会被自动转化为半精度浮点型的torch.HalfTensor：
1、matmul
2、addbmm
3、addmm
4、addmv
5、addr
6、baddbmm
7、bmm
8、chain_matmul
9、conv1d
10、conv2d
11、conv3d
12、conv_transpose1d
13、conv_transpose2d
14、conv_transpose3d
15、linear
16、matmul
17、mm
18、mv
19、prelu
那么只有这些操作才能半精度吗？不是。其他操作比如rnn也可以进行半精度运行，但是需要自己手动，暂时没有提供自动的转换。
原文链接：https://blog.csdn.net/treeswolf/article/details/105748209