1. 程式人生 > >PyTorch 1.0 釋出,JIT、全新的分散式庫、C++ 前端

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 階段。

此外還有一系列更新資訊,詳情檢視釋出公告

下載地址: