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 trainerremote_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)>>> xtensor([[-0.4621-0.0303j, -0.2438-0.5874j], [ 0.7706+0.1421j, 1.2110+0.1918j]])
詳情查看 https://pytorch.org/blog/pytorch-1.6-released/