pytorch半精度，混合精度，單精度訓練的區別amp.initialize

mixed_precision = True
try:  # Mixed precision training https://github.com/NVIDIA/apex
    from apex import amp
except:
    mixed_precision = False  # not installed

 model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=1)

爲了幫助提高Pytorch的訓練效率，英偉達提供了混合精度訓練工具Apex。號稱能夠在不降低性能的情況下，將模型訓練的速度提升2-4倍，訓練顯存消耗減少爲之前的一半。該項目開源於：https://github.com/NVIDIA/apex ,文檔地址是：https://nvidia.github.io/apex/index.html

該 工具 提供了三個功能，amp、parallel和normalization。由於目前該工具還是0.1版本，功能還是很基礎的，在最後一個normalization功能中只提供了LayerNorm層的復現，實際上在後續的使用過程中會發現，出現問題最多的是pytorch的BN層。

第二個工具是pytorch的分佈式訓練的復現，在文檔中描述的是和pytorch中的實現等價，在代碼中可以選擇任意一個使用，實際使用過程中發現，在使用混合精度訓練時，使用Apex復現的parallel工具，能避免一些bug。

默認訓練方式是 單精度float32

import torch
model = torch.nn.Linear(D_in, D_out)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for img, label in dataloader:
	out = model(img)
	loss = LOSS(out, label)
	loss.backward()
	optimizer.step()
	optimizer.zero_grad()

半精度 model(img.half())

import torch
model = torch.nn.Linear(D_in, D_out).half()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for img, label in dataloader:
	out = model(img.half())
	loss = LOSS(out, label)
	loss.backward()
	optimizer.step()
	optimizer.zero_grad()

接下來是混合精度的實現，這裏主要用到Apex的amp工具。代碼修改爲：
加上這一句封裝，model, optimizer = amp.initialize(model, optimizer, opt_level=“O1”)

import torch
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

for img, label in dataloader:
	out = model(img)
	loss = LOSS(out, label)
	# loss.backward()
	with amp.scale_loss(loss, optimizer) as scaled_loss:
    	scaled_loss.backward()

	optimizer.step()
	optimizer.zero_grad()

實際流程爲：調用amp.initialize按照預定的opt_level對model和optimizer進行設置。在計算loss時使用amp.scale_loss進行回傳。

需要注意以下幾點：

在調用amp.initialize之前，模型需要放在GPU上，也就是需要調用cuda()或者to()。
在調用amp.initialize之前，模型不能調用任何分佈式設置函數。
此時輸入數據不需要在轉換爲半精度。

在使用混合精度進行計算時，最關鍵的參數是opt_level。他一共含有四種設置值：‘00’，‘01’，‘02’，‘03’。實際上整個amp.initialize的輸入參數很多：

但是在實際使用過程中發現，設置opt_level即可，這也是文檔中例子的使用方法，甚至在不同的opt_level設置條件下，其他的參數會變成無效。（已知BUG：使用‘01’時設置keep_batchnorm_fp32的值會報錯）

概括起來：00相當於原始的單精度訓練。01在大部分計算時採用半精度，但是所有的模型參數依然保持單精度，對於少數單精度較好的計算（如softmax）依然保持單精度。02相比於01，將模型參數也變爲半精度。03基本等於最開始實驗的全半精度的運算。值得一提的是，不論在優化過程中，模型是否採用半精度，保存下來的模型均爲單精度模型，能夠保證模型在其他應用中的正常使用。這也是Apex的一大賣點。

在Pytorch中，BN層分爲train和eval兩種操作。實現時若爲單精度網絡，會調用CUDNN進行計算加速。常規訓練過程中BN層會被設爲train。Apex優化了這種情況，通過設置keep_batchnorm_fp32參數，能夠保證此時BN層使用CUDNN進行計算，達到最好的計算速度。但是在一些fine tunning場景下，BN層會被設爲eval（我的模型就是這種情況）。此時keep_batchnorm_fp32的設置並不起作用，訓練會產生數據類型不正確的bug。此時需要人爲的將所有BN層設置爲半精度，這樣將不能使用CUDNN加速。一個設置的參考代碼如下：

def fix_bn(m):
	classname = m.__class__.__name__
    if classname.find('BatchNorm') != -1:
    	m.eval().half()

model.apply(fix_bn)

實際測試下來，最後的模型準確度上感覺差別不大，可能有輕微下降；時間上變化不大，這可能會因不同的模型有差別；顯存開銷上確實有很大的降低。