背景
pytorch作為深度學習的計算框架正得到越來越多的應用.
我們除了在模型訓練階段應用外，最近也把pytorch應用在了部署上.
在部署時，為了減少計算量，可以考慮使用16位浮點模型，而訓練時涉及到梯度計算，需要使用32位浮點，這種精度的不一致經(jīng)過測試，模型性能下降有限，可以接受.
但是推斷時計算量可以降低一半，同等計算資源下，并發(fā)度可提升近一倍
具體方法
在pytorch中，一般模型定義都繼承torch.nn.Moudle，torch.nn.Module基類的half()方法會把所有參數(shù)轉(zhuǎn)為16位浮點，所以在模型加載后，調(diào)用一下該方法即可達到模型切換的目的.接下來只需要在推斷時把input的tensor切換為16位浮點即可
另外還有一個小的trick，在推理過程中模型輸出的tensor自然會成為16位浮點，如果需要新創(chuàng)建tensor，最好調(diào)用已有tensor的new_zeros，new_full等方法而不是torch.zeros和torch.full，前者可以自動繼承已有tensor的類型，這樣就不需要到處增加代碼判斷是使用16位還是32位了，只需要針對input tensor切換.
補充：pytorch 使用amp.autocast半精度加速訓練
準備工作
pytorch 1.6+
如何使用autocast？
根據(jù)官方提供的方法，
答案就是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# 創(chuàng)建model，默認是torch.FloatTensormodel = 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中經(jīng)典的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# 創(chuàng)建model，默認是torch.FloatTensormodel = 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
單卡訓練的話上面的代碼已經(jīng)夠了，親測在2080ti上能減少至少1/3的顯存，至于速度。。。
要是想多卡跑的話僅僅這樣還不夠，會發(fā)現(xiàn)在forward里面的每個結(jié)果都還是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會被自動轉(zhuǎn)化為半精度浮點型的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也可以進行半精度運行，但是需要自己手動，暫時沒有提供自動的轉(zhuǎn)換。
原文鏈接：https://blog.csdn.net/treeswolf/article/details/105748209
