前言
導致pytorch的模型訓練速度比較慢的原因最有可能的是三個：1. 數據導入環節，操作複雜 2.模型本身很複雜，數據流在模型中傳遞時過於耗時 3.loss函數計算複雜。
這其中第一個環節往往是最有可能的原因，第二，三個環節其實一回事；pytorch本身的框架針對這兩個問題也做了大量的優化，如果不是專業技術過硬，建議在這兩個環節上就不要過於糾結了，設計簡單易用的模型纔是正道。
數據導入環節的優化
數據導入環節，尤其是諸如圖像等大張量從內存中的反覆讀取，以及後續的數據增廣操作往往是造成訓練速度低的主要原因。針對這個環節的加速其實有一些trick可以試用。
- 常規的做法
常規的做法主要是訓練中各種現成的pytorch工具使用以及訓練參數的設置，主要有如下的幾種方案：
- 採用pytorch 自帶的Dataloader而不是自己編寫張量的導入類，Dataloader可以方便的設置cpu多線程，很多操作諸如張量的縮放都經過了優化
- 採用更大的batch size，設置cpu或者GPU能夠承受的最大batch size，這最大的用處在於節省了後續梯度傳遞時花的時間
- 可以使用累積梯度，其實就是在cpu或者GPU能承受範圍內，多次循環batch 再進行梯度計算
- 保存圖，就像下面這樣
losses = []
...
losses.append(loss)
print(f'current loss: {torch.mean(losses)'})
5. 使用多個GPU
當然這裏推薦一個用於加速運行的插件，叫做pytorch lighting （https://github.com/williamFalcon/pytorch-lightning）
它的使用也是比較簡單的
from pytorch_lightning import Trainer
model = LightningModule(…)
trainer = Trainer()
trainer.fit(model)
- 非常規的做法
非常規的做法就得視場合而定了，這些做法並不是對所有的應用場景有效，在不適合的場景裏可能造成嚴重的訓練質量下降。
- 半精度或者混合精度訓練， 該方法在一些本來對張量精度要求不是很高的領域比較適用，可以顯著的提高訓練速度，同時顯著降低運算顯存開銷，但是並不是所有領域都適合。關於半精度以及混合精度，可以採用apex library 在英偉達顯卡上方便的實現。
- 提前規範化數據，比如大量的圖像張量導入時，可以將圖像提前縮放成2的n次方類型，這主要是因爲大量的框架優化對於這種尺度的圖像處理優化效果明顯，而對於任意尺寸輸入的圖像不敢保證；但是提前的規範化有可能造成一些細節的變形
- 使用hdf5格式，提前將數據轉成hdf5格式，這種格式對於cpu運算較爲友好，同時也是受限比較小的一種方式；但是我在使用中發現，hdf5的解析有賴於自己寫的方式，如果技術不過硬這裏有可能還是解決不了問題。我這裏有一個示例類，可供參考
class AdobePatchDataHDF5(data.Dataset):
def __init__(self, root, cropsize = 256, outputsize = 256):
fgfile = h5py.File(root, 'r')
self.root = root
self.fgfile = fgfile
self.cropsize = cropsize
self.outputsize = outputsize
def __getitem__(self, index):
# read image
fgimg = self.fgfile['img'][index, ...]
# random crop and resize, random flip with cv2
# toTensors
fgimg = fgimg.astype(np.float32) / 255.0
fgimg = torch.from_numpy(fg.transpose((2, 0, 1)))
# norm [0, 1] to [-1, 1]
return fgimg, label
def __len__(self):
return self.fgfile['img'].shape[0]
還有一些是同gpu綁定的方法，比如使用Nvidia DALI，（https://github.com/NVIDIA/DALI），這在預處理階段可以進行極大的加速，但是目前的穩定版本(截止12.25)好像只能支持有限型號的顯卡。
模型訓練環節以及loss環節優化
這兩部分的優化就比較專業了，需要過硬的本事來做平臺上的優化，一般而言很難以取得效果，相反有可能造成較大的問題；比如我曾經手動從頭書寫DenseNet而不是採用pytorch自身的源代碼，結果發現不但速度降低而且還造成顯存消耗劇增；當然，當時由於時間問題，這事情就沒有深究了。
一般來說，這一塊的優化，主要是採用半精度或者混合精度的訓練來達成；當然如果硬件允許，其實使用tensorRt來進行訓練也是非常不錯的選擇。
綜合
基本上訓練緩慢的原因集中在第一點，對於這裏的優化方案可供參考的也層出不窮；一般不推薦針對後兩種的優化，那對於一般人來說較爲複雜，稍不留意可能適得其反。