App下載

pytorch怎么部署半精度模型?如何操作?

猿友 2021-07-15 13:44:55 瀏覽數(shù) (4100)
反饋

在進行深度學習的模型訓練的時候,很多小伙伴的電腦的配置可能不是那么優(yōu)秀,在該階段要花費很多的運行時間。但在學習深度學習的過程中,有時候我們其實并不需要模型擁有太高的精度,這時候我們就可以設置通過降低精度的方法減少運算量,這就是pytorch加速模型訓練的辦法,那么pytorch如何設置精度呢?通過下文的學習,你將學會pytorch怎么部署半精度模型。

背景

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 類。使用也是非常簡單的

from torch.cuda.amp import autocast as autocast

# 創(chuàng)建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中經(jīng)典的AMP使用方式如下:

from torch.cuda.amp import autocast as autocast

# 創(chuàng)建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

單卡訓練的話上面的代碼已經(jīng)夠了,親測在2080ti上能減少至少1/3的顯存,至于速度。。。

要是想多卡跑的話僅僅這樣還不夠,會發(fā)現(xiàn)在forward里面的每個結(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代碼塊方式運行就好啦!

自動進行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)換。

小結(jié)

以上就是pytorch怎么部署半精度模型的全部內(nèi)容,希望能給大家一個參考,也希望大家多多支持W3Cschool。


0 人點贊