背景
pytorch作為深度學(xué)習(xí)的計(jì)算框架正得到越來越多的應(yīng)用.
我們除了在模型訓(xùn)練階段應(yīng)用外,最近也把pytorch應(yīng)用在了部署上.
在部署時(shí),為了減少計(jì)算量,可以考慮使用16位浮點(diǎn)模型,而訓(xùn)練時(shí)涉及到梯度計(jì)算,需要使用32位浮點(diǎn),這種精度的不一致經(jīng)過測試,模型性能下降有限,可以接受.
但是推斷時(shí)計(jì)算量可以降低一半,同等計(jì)算資源下,并發(fā)度可提升近一倍
具體方法
在pytorch中,一般模型定義都繼承torch.nn.Moudle,torch.nn.Module基類的half()方法會(huì)把所有參數(shù)轉(zhuǎn)為16位浮點(diǎn),所以在模型加載后,調(diào)用一下該方法即可達(dá)到模型切換的目的.接下來只需要在推斷時(shí)把input的tensor切換為16位浮點(diǎn)即可
另外還有一個(gè)小的trick,在推理過程中模型輸出的tensor自然會(huì)成為16位浮點(diǎn),如果需要新創(chuàng)建tensor,最好調(diào)用已有tensor的new_zeros,new_full等方法而不是torch.zeros和torch.full,前者可以自動(dòng)繼承已有tensor的類型,這樣就不需要到處增加代碼判斷是使用16位還是32位了,只需要針對input tensor切換.
補(bǔ)充:pytorch 使用amp.autocast半精度加速訓(xùn)練
準(zhǔn)備工作
pytorch 1.6+
如何使用autocast?
根據(jù)官方提供的方法,
答案就是autocast + GradScaler。
如何在PyTorch中使用自動(dòng)混合精度?
答案:autocast + GradScaler。
1.autocast
正如前文所說,需要使用torch.cuda.amp模塊中的autocast 類。使用也是非常簡單的
from torch.cuda.amp import autocast as autocast
# 創(chuàng)建model,默認(rèn)是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模塊,需要在訓(xùn)練最開始之前實(shí)例化一個(gè)GradScaler對象。
因此PyTorch中經(jīng)典的AMP使用方式如下:
from torch.cuda.amp import autocast as autocast
# 創(chuàng)建model,默認(rèn)是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在訓(xùn)練最開始之前實(shí)例化一個(gè)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
單卡訓(xùn)練的話上面的代碼已經(jīng)夠了,親測在2080ti上能減少至少1/3的顯存,至于速度。。。
要是想多卡跑的話僅僅這樣還不夠,會(huì)發(fā)現(xiàn)在forward里面的每個(gè)結(jié)果都還是float32的,怎么辦?
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, input_data_c1):
with autocast():
# code
return
只要把forward里面的代碼用autocast代碼塊方式運(yùn)行就好啦!
自動(dòng)進(jìn)行autocast的操作
如下操作中tensor會(huì)被自動(dòng)轉(zhuǎn)化為半精度浮點(diǎ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也可以進(jìn)行半精度運(yùn)行,但是需要自己手動(dòng),暫時(shí)沒有提供自動(dòng)的轉(zhuǎn)換。
您可能感興趣的文章:- 關(guān)于pytorch多GPU訓(xùn)練實(shí)例與性能對比分析
- 在pytorch 中計(jì)算精度、回歸率、F1 score等指標(biāo)的實(shí)例
- pytorch訓(xùn)練imagenet分類的方法
- PyTorch預(yù)訓(xùn)練Bert模型的示例