PyTorch 1.6 釋出：原生支援自動混合精度訓練並進入穩定階段

PyTorch 1.6 穩定版已釋出，此版本增加了許多新的 API、用於效能改進和效能分析的工具、以及對基於分散式資料並行(Distributed Data Parallel, DDP)和基於遠端過程呼叫(Remote Procedure Call, RPC)的分散式訓練的重大更新。部分更新亮點包括：

• 原生支援自動混合精度訓練(automatic mixed-precision training)，並已進入穩定階段
• 為 tensor-aware 增加對 TensorPipe 的原生支援
• 在前端 API 增加了對 complex tensor 的支援
• 新的分析工具提供了張量級的記憶體消耗資訊
• 針對分散式資料並行訓練和遠端過程呼叫的多項改進和新功能

此外，從該版本起，新功能的狀態將分為三種，分別是穩定、Beta 和原型。原型功能即為新功能提案，團隊如果能從使用者獲取對此的良好反饋，那麼原型功能就會進入 Beta 階段，否則就停止開發。另外，原型功能不包含在二進位制發行版中，而是通過從 Nightly 版本原始碼構建或通過 compiler flag 使用。詳情檢視此部落格。

原生支援自動混合精度訓練

由 Nvidia 貢獻的自動混合精度訓練功能已經進入穩定階段，AMP 訓練能在 Tensor Core GPU 上實現更高的效能並節省多達 50％ 的記憶體。

AMP API 提供方便使用混合精度的方法。官方提到，在像線性層(Linear Layer)或是卷積操作上，float16 運算較快，但像 Reduction 運算又需要 float32 的動態範圍，而現在有了 AMP 功能，便可以在部分運算操作使用 float16，另一部分則使用 float32，混合精度功能會嘗試為每個運算使用相匹配的資料型別。

改進分散式訓練

PyTorch 支援兩種強大的正規化：用於對模型進行完全同步資料並行訓練的 DDP 和支援分散式模型並行的 RPC 框架。過去這兩個功能獨立執行，使用者無法混合和匹配它們來嘗試混合並行正規化。

從 PyTorch 1.6 開始，DDP 和 RPC 可以無縫協作，使用者可以結合這兩種技術來實現資料並行和模型並行。官方舉了一個例子，使用者希望將大型嵌入表放置在引數伺服器上，並使用 RPC 框架嵌入查詢，但希望將較小的密集引數儲存在訓練器上，並使用 DDP 來同步密集引數，下面是示例程式碼：

 // On each trainer
 
 remote_emb = create_emb(on="ps", ...)
 ddp_model = DDP(dense_model)
 
 for data in batch:
    with torch.distributed.autograd.context():
       res = remote_emb(data)
       loss = ddp_model(res)
       torch.distributed.autograd.backward([loss])

支援 Complex Tensor

PyTorch 1.6 帶來了對 complex tensor 的 Beta 支援，包含 torch.complex64 和 torch.complex128 dtypes 兩種型別。Beta 階段支援通用的 PyTorch 和 complex tensor，以及 Torchaudio、ESPnet 等所需的功能。

 >>> x = torch.randn(2,2, dtype=torch.cfloat)
 >>> x
 tensor([[-0.4621-0.0303j, -0.2438-0.5874j],
      [ 0.7706+0.1421j,  1.2110+0.1918j]])

詳情檢視 https://pytorch.org/blog/pytorch-1.6-released/