背景：
pytorch從1.6版本開始，已經內置了torch.cuda.amp，采用自動混合精度訓練就不需要加載第三方NVIDIA的apex庫了。本文主要從三個方面來介紹AMP：
一．什么是AMP?
二．為什么要使用AMP？
三．如何使用AMP?
四. 注意事項
正文：
一．什么是AMP?
默認情況下，大多數深度學習框架都采用32位浮點算法進行訓練。２０１７年，NVIDIA研究了一種用於混合精度訓練的方法，該方法在訓練網絡時將單精度（FP32）與半精度(FP16)結合在一起，並使用相同的超參數實現了與FP32幾乎相同的精度。
在介紹AMP之前，先來理解下FP16與FP32，FP16也即半精度是一種計算機使用的二進制浮點數據類型，使用2字節存儲。而FLOAT就是FP32。
其中，sign位表示正負，exponent位表示指數2^(n-15+1(n=0))，fraction位表示分數(m/1024)。
一般情況下，我們在pytorch中創建一個Tensor:
>>import torch >>tensor1=torch.zeros(30,20) >>tensor1.type() 'torch.FloatTensor' >>tensor2=torch.Tensor([1,2]) >>tensor2.type() 'torch.FlatTensor'
可以看到，默認創建的tensor都是FloatTensor類型。而在Pytorch中，一共有10種類型的tensor:
torch.FloatTensor(32bit floating point)
torch.DoubleTensor(64bit floating point)
torch.HalfTensor(16bit floating piont1)
torch.BFloat16Tensor(16bit floating piont2)
torch.ByteTensor(8bit integer(unsigned)
torch.CharTensor(8bit integer(signed))
torch.ShortTensor(16bit integer(signed))
torch.IntTensor(32bit integer(signed))
torch.LongTensor(64bit integer(signed))
torch.BoolTensor(Boolean)
默認Tensor是32bit floating point，這就是32位浮點型精度的tensor。
AMP(自動混合精度）的關鍵詞有兩個：自動，混合精度。
自動：Tensor的dtype類型會自動變化，框架按需自動調整tensor的dtype,當然有些地方還需手動干預。
混合精度：采用不止一種精度的Tensor，torch.FloatTensor和torch.HalfTensor
pytorch1.6的新包：torch.cuda.amp，是ＮVIDIA開發人員貢獻到pytorch里的。只有支持tensor core的CUDA硬件才能享受到AMP帶來的優勢。Tensor core是一種矩陣乘累加的計算單元，每個tensor core時針執行64個浮點混合精度操作（FP16矩陣相乘和FP32累加）。
二、為什么要使用AMP?
前面已介紹，AMP其實就是Float32與Float16的混合，那為什么不單獨使用Float32或Float16，而是兩種類型混合呢？原因是：在某些情況下Float32有優勢，而在另外一些情況下Float16有優勢。這里先介紹下FP16：
優勢有三個：
１．減少顯存占用；
２．加快訓練和推斷的計算，能帶來多一倍速的體驗；
３．張量核心的普及（NVIDIA Tensor Core）,低精度計算是未來深度學習的一個重要趨勢。
但凡事都有兩面性，FP16也帶來了些問題：１．溢出錯誤；２．舍入誤差；
１．溢出錯誤：由於FP16的動態范圍比FP32位的狹窄很多，因此，在計算過程中很容易出現上溢出和下溢出，溢出之后就會出現"NaN"的問題。在深度學習中，由於激活函數的梯度往往要比權重梯度小，更易出現下溢出的情況
2.舍入誤差
舍入誤差指的是當梯度過小時，小於當前區間內的最小間隔時，該次梯度更新可能會失敗：
為了消除torch.HalfTensor也就是FP16的問題，需要使用以下兩種方法：
１）混合精度訓練
在內存中用FP16做儲存和乘法從而加速計算，而用FP32做累加避免舍入誤差。混合精度訓練的策略有效地緩解了舍入誤差的問題。
什么時候用torch.FloatTensor,什么時候用torch.HalfTensor呢？這是由pytorch框架決定的，在pytorch1.6的AMP上下文中，以下操作中Tensor會被自動轉化為半精度浮點型torch.HalfTensor：
__matmul__ addbmm addmm addmv addr baddbmm bmm chain_matmul conv1d conv2d conv3d conv_transpose1d conv_transpose2d conv_transpose3d linear matmul mm mv prelu
2)損失放大（Loss scaling)
即使了混合精度訓練，還是存在無法收斂的情況，原因是激活梯度的值太小，造成了溢出。可以通過使用torch.cuda.amp.GradScaler，通過放大loss的值來防止梯度的underflow（只在BP時傳遞梯度信息使用，真正更新權重時還是要把放大的梯度再unscale回去）；
反向傳播前，將損失變化手動增大2^k倍，因此反向傳播時得到的中間變量（激活函數梯度）則不會溢出；
反向傳播后，將權重梯度縮小2^k倍，恢復正常值。
三．如何使用AMP?
目前有兩種版本：pytorch1.5之前使用的NVIDIA的三方包apex.amp和pytorch1.6自帶的torch.cuda.amp
１．pytorch1.5之前的版本（包括１．５）
使用方法如下：
from apex import amp model,optimizer = amp.initial(model,optimizer,opt_level="O1") #注意是O,不是０ with amp.scale_loss(loss,optimizer) as scaled_loss: scaled_loss.backward()
取代
loss.backward()
其中,opt_level配置如下：
O0:純FP32訓練，可作為accuracy的baseline；
O1:混合精度訓練（推薦使用），根據黑白名單自動決定使用FP16(GEMM,卷積）還是FP32（softmax)進行計算。
O2:幾乎FP16，混合精度訓練，不存在黑白名單 ，除了bacthnorm，幾乎都是用FP16計算；
O3:純FP16訓練，很不穩定，但是可以作為speed的baseline；
動態損失放大（dynamic loss scaling)部分，為了充分利用FP16的范圍，緩解舍入誤差，盡量使用最高的放大倍數2^24,如果產生上溢出，則跳出參數更新，縮小放大倍數使其不溢出。在一定步數后再嘗試使用大的scale來充分利用FP16的范圍。
分布式訓練：
import argparse import apex import amp import apex.parallel import convert_syncbn_model import apex.parallel import DistributedDataParallel as DDP 定義超參數： def parse(): parser=argparse.ArgumentParser() parser.add_argument('--local_rank',type=int, default=0) #local_rank指定了輸出設備，默認為GPU可用列表中的第一個GPU，必須加上。 ... args = parser.parser.parse_args() return args 主函數寫： def main(): args = parse() torch.cuda.set_device(args.local_rank) #必須寫在下一句的前面 torch.distributed.init_process_group( 'nccl', init_method='env://') 導入數據接口，需要用DistributedSampler dataset = ... num_workers = 4 if cuda else 0 train_sampler=torch.utils.data.distributed.DistributedSampler(dataset) loader = DataLoader(dataset, batch_size=batchsize, shuflle=False, num_worker=num_workers,pin_memory=cuda, drop_last=True, sampler=train_sampler) 定義模型： net = XXXNet(using_amp=True) net.train() net= convert_syncbn_model(net) device=torch.device('cuda:{}'.format(args.local_rank)) net=net.to(device) 定義優化器，損失函數，定義優化器一定要把模型搬運到GPU之上 apt = Adam([{'params':params_low_lr,'lr':4e-5}, {'params':params_high_lr,'lr':1e-4}],weight_decay=settings.WEIGHT_DECAY) crit = nn.BCELoss().to(device) 多GPU設置
import torch.nn.parallel.DistributedDataParallel as DDP net,opt = amp.initialize(net,opt,opt_level='o1') net=DDP(net,delay_allreduce=True)
loss使用方法：
opt.zero_grad()
with amp.scale_loss(loss, opt) as scaled_loss:
scaled_loss.backward()
opt.step()
加入主入口：
if __name__ == '__main__':
main()
無論是apex支持的DDP還是pytorch自身支持的DDP,都需使用torch.distributed.launch來使用，方法如下：
CUDA_VISIBLE_DIVECES=1,2,4 python -m torch.distributed.launch --nproc_per_node=3 train.py
1,2,4是GPU編號，nproc_per_node是指定用了哪些GPU,記得開頭說的local_rank，是因為torch.distributed.launch會調用這個local_ran
分布式訓練時保存模型注意點：
如果直接在代碼中寫torch.save來保存模型，則每個進程都會保存一次相同的模型，會存在寫文件寫到一半，會被個進程寫覆蓋的情況。如何避免呢?
可以用local_rank == 0來僅僅在第一個GPU上執行進程來保存模型文件。
雖然是多個進程，但每個進程上模型的參數值都是一樣的，而默認代號為０的進程是主進程
if arg.local_rank == 0: torch.save(xxx)
２．pytorch1.6及以上版本
有兩個接口：autocast和Gradscaler
1) autocast
導入pytorch中模塊torch.cuda.amp的類autocast
from torch.cuda.amp import autocast as autocast model=Net().cuda() optimizer=optim.SGD(model.parameters(),...) for input,target in data: optimizer.zero_grad() with autocast(): output=model(input) loss = loss_fn(output,target) loss.backward() optimizer.step()
可以使用autocast的context managers語義（如上），也可以使用decorators語義。當進入autocast上下文后，在這之后的cuda ops會把tensor的數據類型轉換為半精度浮點型，從而在不損失訓練精度的情況下加快運算。而不需要手動調用.half(),框架會自動完成轉換。
不過，autocast上下文只能包含網絡的前向過程(包括loss的計算），不能包含反向傳播，因為BP的op會使用和前向op相同的類型。
當然，有時在autocast中的代碼會報錯：
Traceback (most recent call last): ...... File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl result = self.forward(*input, ** kwargs) ...... RuntimeError: expected scalar type float but found c10::Half
對於RuntimeError:expected scaler type float but found c10:Half,應該是個bug,可在tensor上手動調用.float()來讓type匹配。
２）GradScaler
使用前，需要在訓練最開始前實例化一個GradScaler對象，例程如下：
from torch.cuda.amp import autocast as autocast model=Net().cuda() optimizer=optim.SGD(model.parameters(),...) scaler = GradScaler() #訓練前實例化一個GradScaler對象 for epoch in epochs: for input,target in data: optimizer.zero_grad() with autocast(): ＃前后開啟autocast output=model(input) loss = loss_fn(output,targt) scaler.scale(loss).backward() #為了梯度放大 #scaler.step() 首先把梯度值unscale回來，如果梯度值不是inf或NaN,則調用optimizer.step()來更新權重，否則，忽略step調用，從而保證權重不更新。
scaler.step(optimizer) scaler.update() #准備着，看是否要增大scaler
scaler的大小在每次迭代中動態估計，為了盡可能減少梯度underflow，scaler應該更大；但太大，半精度浮點型又容易overflow（變成inf或NaN).所以，動態估計原理就是在不出現if或NaN梯度的情況下，盡可能的增大scaler值。在每次scaler.step(optimizer)中，都會檢查是否有inf或NaN的梯度出現：
１．如果出現inf或NaN,scaler.step(optimizer)會忽略此次權重更新(optimizer.step()），並將scaler的大小縮小（乘上backoff_factor)；
２．如果沒有出現inf或NaN,那么權重正常更新，並且當連續多次(growth_interval指定)沒有出現inf或NaN，則scaler.update()會將scaler的大小增加(乘上growth_factor)。
對於分布式訓練，由於autocast是thread local的，要注意以下情形：
１）torch.nn.DataParallel：
以下代碼分布式是不生效的
model = MyModel() dp_model = nn.DataParallel(model) with autocast(): output=dp_model(input) loss=loss_fn(output)
需使用autocast裝飾model的forward函數
MyModel(nn.Module): @autocast() def forward(self, input): ... #alternatively MyModel(nn.Module): def forward(self, input): with autocast(): ... model = MyModel() dp_model=nn.DataParallel(model) with autocast(): output=dp_model(input) loss = loss_fn(output)
2）torch.nn.DistributedDataParallel:
同樣，對於多GPU,也需要autocast裝飾model的forward方法，保證autocast在進程內部生效。
四. 注意事例：
在使用AMP時，由於報錯信息並不明顯，給調試帶來了一定的難度。但只要注意以下一些點，相信會少走很多彎路。
１．判斷GPU是否支持FP16，支持Tensor core的GPU（2080Ti,Titan,Tesla等），不支持的(Pascal系列）不建議；
1080Ti與2080Ti對比
gtx 1080ti: 半精度浮點數：0.17TFLOPS 單精度浮點數：11.34TFLOPS 雙精度浮點數：0.33TFLOPS rtx 2080ti: 半精度浮點數：20.14TFLOPS 單精度浮點數：10.07TFLOPS 雙精度浮點數：0.31TFLOPS
半精度浮點數即FP16，單精度浮點數即FP32，雙精度浮點數即FP64。
在不使用apex的pytorch訓練過程中，一般默認均為單精度浮點數，從上面的數據可以看到1080ti和2080ti的單精度浮點數運算能力差不多，因此不使用apex時用1080ti和2080ti訓練模型時間上差別很小。
使用apex時用1個2080ti訓練時一個epoch是2h31min，兩者時間幾乎一樣，但是卻少用了一張2080ti。這是因為在pytorch訓練中使用apex時，此時大多數運算均為半精度浮點數運算，而2080ti的半精度浮點數運算能力是其單精度浮點數運算能力的兩倍
２．常數范圍：為了保證計算不溢出，首先保證人工設定的常數不溢出。如epsilon,INF等；
３．Dimension最好是8的倍數:維度是８的倍數，性能最好；
４．涉及sum的操作要小心，容易溢出，softmax操作，建議用官方API，並定義成layer寫在模型初始化里；
５．模型書寫要規范：自定義的Layer寫在模型初始化函數里，graph計算寫在forward里；
６．一些不常用的函數，使用前要注冊：amp.register_float_function(torch,'sogmoid')
７．某些函數不支持FP16加速，建議不要用；
８．需要操作梯度的模塊必須在optimizer的step里，不然AMP不能判斷grad是否為NaN