PyTorch 1.0 釋出,JIT、全新的分散式庫、C++ 前端
Python 張量與動態神經網路 PyTorch 1.0 釋出了。
此版本的主要亮點包括JIT 編譯、全新並且更快的分散式庫與 C++ 前端等。
JIT 編譯器
JIT(Just-In-Time)是一組編譯工具,用於彌合 PyTorch 研究與生產之間的差距。
它允許建立可以在不依賴 Python 直譯器的情況下執行的模型,並且可以更積極地進行優化。使用程式註解可以將現有模型轉換為 PyTorch 可以直接執行的 Python 子集 Torch Script。模型程式碼仍然是有效的 Python 程式碼,可以使用標準的 Python 工具鏈進行除錯。
PyTorch 1.0 提供了 torch.jit.trace 和 torch.jit.script 兩種方式使現有程式碼與 JIT 相容。一經註解,Torch Script 程式碼便可以被積極地優化,並且可以被序列化以在新的 C++ API 中使用,並且 C++ API 不依賴於 Python。
# Write in Python, run anywhere! @torch.jit.script def RNN(x, h, W_h, U_h, b_h): y = [] for t in range(x.size(0)): h = torch.tanh(x[t] @ W_h + h @ U_h + b_h) y += [h] return torch.stack(y), h
全新並且更快的分散式庫
torch.distributed 軟體包和 torch.nn.parallel.DistributedDataParallel 模組採用了重新設計的全新分散式庫,亮點包括:
新的 torch.distributed 是效能驅動的,並且對所有後端完全非同步操作,包括:Gloo、NCCL 和 MPI。
顯著的分散式資料並行效能改進,尤其適用於網路較慢的主機,如基於乙太網的主機。
為 torch.distributed 包中的所有分散式集合操作新增非同步支援。
在 Gloo 後端添加了一些 CPU 操作:send、recv、reduce、all_gather、gather 與 scatter。
在 NCCL 後端新增 barrier 操作。
在 NCCL 後端新增 new_group 支援。
C++ 前端
C++ 前端是 PyTorch 後端的純 C++ 介面,它遵循已建立的 Python 前端的 API 和體系結構,旨在實現高效能、低延遲和裸機 C++ 應用的研究。它提供了 torch.nn、torch.optim、torch.data 和 Python 前端的其它元件的等價物。下邊是兩種語言前端的簡單比較:
Python | C++ |
---|---|
import torch model = torch.nn.Linear(5, 1) optimizer = torch.optim.SGD(model.parameters(), lr=0.1) prediction = model.forward(torch.randn(3, 5)) loss = torch.nn.functional.mse_loss(prediction, torch.ones(3, 1)) loss.backward() optimizer.step() |
#include <torch/torch.h> torch::nn::Linear model(5, 1); torch::optim::SGD optimizer(model->parameters(), /*lr=*/0.1); torch::Tensor prediction = model->forward(torch::randn({3, 5})); auto loss = torch::mse_loss(prediction, torch::ones({3, 1})); loss.backward(); optimizer.step(); |
注意,目前 C++ API 還處於 unstable 階段。
此外還有一系列更新資訊,詳情檢視釋出公告。
下載地址: