新智元原創
來源：pytorch.org、GitHub
編輯：金磊
【新智元導讀】盼望已久，Pytorch終於更新了！Pytroch 1.1.0的發布除了修復了已有bug之外，最大的亮點就是可以更快、更好的支持自定義RNN，以及TensorBoard對可視化和模型調試提供了一流的本地支持。
Pytorch 1.1.0，來了！
可以說是一大波更新來襲了，話不多說上亮點：
TorchScript(Pytorch JIT)更快、更好的支持自定義RNN；TensorBoard對可視化和模型調試提供了一流的本地支持；可以在ScriptModule上通過使用torch.jit包裝屬性來分配屬性；TorchScript現在對列表和字典類型提供了魯棒性的支持；對於更複雜的有狀態操作，TorchScript現在支持使用@torch.jit.script注釋類；nn.parallel.DistributedDataParallel：現在可以包裝多GPU模塊，它可以在一臺伺服器上實現模型並行和跨伺服器的數據並行等用例。
註：不再支持CUDA 8.0。
此更新一出，在Reddit上也引發了一波熱議，大部分網友們表示：
「贊！」、「好用！」、「愛了！」
用TorchScript優化CUDA遞歸神經網絡
Pytorch添加的一個新特性是更好地支持帶有TorchScript (PyTorch JIT)的快速自定義遞歸神經網絡(fastrnns)。
RNN是一種流行的模型，在各種NLP任務上都表現出了良好的性能。PyTorch可以實現許多最流行的變體，例如Elman RNN、GRU和LSTM，以及多層和雙向變體。
然而，許多用戶希望實現他們自己的自定義RNN。將層規範化應用於LSTM就是這樣一種用例。由於PyTorch CUDA LSTM實現使用融合內核，因此很難插入規範化甚至修改基本LSTM實現。
許多用戶已經轉向使用標準PyTorch運算符編寫自定義實現，但是這樣的代碼遭受高開銷：大多數PyTorch操作在GPU上啟動至少一個內核，並且RNN由於其重複性質通常運行許多操作。但是可以應用TorchScript來融合操作並自動優化代碼，在GPU上啟動更少、更優化的內核。
此次更新的目標之一是讓用戶能夠在TorchScript中編寫快速，自定義的RNN，而無需編寫專門的CUDA內核來實現類似的性能。接下來將提供如何使用TorchScript編寫自己的快速RNN的教程。
編寫自定義RNN
首先，可以使用下方連結中的文件作為模板來編寫自己的自定義RNN。
https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py
如果想獲得TorchScript當前提供的速度/優化（如運算符融合，批量矩陣乘法等），請遵循以下指南。
如果定製操作都是element-wise的，那就可以自動獲得PyTorch JIT操作符fusion的優勢！如果有更複雜的操作（例如，reduce和element-wise的渾南和操作），請考慮分別對reduce操作和element-wise操作進行分組。如果想知道自定義RNN中融合了什麼，可以使用graph_for檢查操作的優化圖。以LSTMCell為例：
# get inputs and states for LSTMCellinputs = get_lstm_inputs()# instantiate a ScriptModule cell = LSTMCell(input_size, hidden_size)# print the optimized graph using graph_for out = cell(inputs) print(cell.graph_for(inputs))
這將提供的專用輸入生成優化的TorchScript圖形（a.k.a PyTorch JIT IR）：
graph(%x : Float(*, *),%hx : Float(*, *), %cx : Float(*, *), %w_ih : Float(*, *), %w_hh : Float(*, *), %b_ih : Float(*), %b_hh : Float(*)): %hy : Float(*, *), %cy : Float(*, *) = prim::DifferentiableGraph_0(%cx, %b_hh, %b_ih, %hx, %w_hh, %x, %w_ih) %30 : (Float(*, *), Float(*, *)) = prim::TupleConstruct(%hy, %cy)return (%30)with prim::DifferentiableGraph_0 = graph(%13 : Float(*, *), %29 : Float(*), %33 : Float(*), %40 : Float(*, *), %43 : Float(*, *), %45 : Float(*, *), %48 : Float(*, *)): %49 : Float(*, *) = aten::t(%48) %47 : Float(*, *) = aten::mm(%45, %49) %44 : Float(*, *) = aten::t(%43) %42 : Float(*, *) = aten::mm(%40, %44) ...some broadcast sizes operations... %hy : Float(*, *), %287 : Float(*, *), %cy : Float(*, *), %outgate.1 : Float(*, *), %cellgate.1 : Float(*, *), %forgetgate.1 : Float(*, *), %ingate.1 : Float(*, *) = prim::FusionGroup_0(%13, %346, %345, %344, %343) ...some broadcast sizes operations...return (%hy, %cy, %49, %44, %196, %199, %340, %192, %325, %185, %ingate.1, %forgetgate.1, %cellgate.1, %outgate.1, %395, %396, %287)with prim::FusionGroup_0 = graph(%13 : Float(*, *), %71 : Tensor, %76 : Tensor, %81 : Tensor, %86 : Tensor): ...some chunks, constants, and add operations... %ingate.1 : Float(*, *) = aten::sigmoid(%38) %forgetgate.1 : Float(*, *) = aten::sigmoid(%34) %cellgate.1 : Float(*, *) = aten::tanh(%30) %outgate.1 : Float(*, *) = aten::sigmoid(%26) %14 : Float(*, *) = aten::mul(%forgetgate.1, %13) %11 : Float(*, *) = aten::mul(%ingate.1, %cellgate.1) %cy : Float(*, *) = aten::add(%14, %11, %69) %4 : Float(*, *) = aten::tanh(%cy) %hy : Float(*, *) = aten::mul(%outgate.1, %4)return (%hy, %4, %cy, %outgate.1, %cellgate.1, %forgetgate.1, %ingate.1)
從上圖中可以看到它有一個prim :: FusionGroup_0子圖，它融合了LSTMCell中的所有element-wise操作（轉置和矩陣乘法不是element-wise操作）。
可變長度序列最佳實踐
TorchScript不支持PackedSequence。 通常，當處理可變長度序列時，最好將它們填充到單個張量中並通過TorchScript LSTM發送該張量。 例如：
sequences = [...] # List[Tensor], each Tensor is T' x Cpadded = torch.utils.rnn.pad_sequence(sequences)lengths = [seq.size(0) for seq in sequences]padded # T x N x C, where N is batch size and T is the max of all T'model = LSTM(...)output, hiddens = model(padded)output # T x N x C
當然，output可能在填充區域中有一些垃圾數據；使用lengths來跟蹤你不需要的部分。
優化
現在將解釋PyTorch JIT為加速自定義RNN所執行的優化。 將在TorchScript中使用一個簡單的自定義LSTM模型來說明優化，但其中許多是通用的並適用於其他RNN。
為了說明所做的優化以及如何從這些優化中獲益，將運行一個用TorchScript編寫的簡單自定義LSTM模型（可以參考custom_lstm.py中的代碼或下面的代碼片段）並計算更改。
在配備2個Intel Xeon晶片和一個Nvidia P100的機器中設置環境，安裝了cuDNN v7.3，CUDA 9.2。 LSTM模型的基本設置如下：
input_size = 512hidden_size = 512mini_batch = 64numLayers = 1seq_length = 100
PyTorch JIT最重要的是將python程序編譯為PyTorch JIT IR，這是一個用於對程序圖形結構進行建模的中間表示。然後，該IR可以從整個程序優化，硬體加速中受益，並且總體上具有提供大量計算增益的潛力。
接下來，將解釋在如何提高訓練或推理性能方面所做的主要優化，從LSTMCell和LSTMLayer開始，以及一些misc優化。
LSTM Cell(前向)
LSTM中的幾乎所有計算都發生在LSTMCell中，因此重要的是看看它包含的計算以及如何提高它們的速度。 下面是TorchScript中的LSTMCell實現示例：
classLSTMCell(jit.ScriptModule):def__init__(self, input_size, hidden_size):super(LSTMCell, self).__init__() self.input_size = input_size self.hidden_size = hidden_size self.weight_ih = Parameter(torch.randn(4 * hidden_size, input_size)) self.weight_hh = Parameter(torch.randn(4 * hidden_size, hidden_size)) self.bias_ih = Parameter(torch.randn(4 * hidden_size)) self.bias_hh = Parameter(torch.randn(4 * hidden_size)) @jit.script_methoddefforward(self, input, state):# type: (Tensor, Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]] hx, cx = state gates = (torch.mm(input, self.weight_ih.t()) + self.bias_ih + torch.mm(hx, self.weight_hh.t()) + self.bias_hh) ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1) ingate = torch.sigmoid(ingate) forgetgate = torch.sigmoid(forgetgate) cellgate = torch.tanh(cellgate) outgate = torch.sigmoid(outgate) cy = (forgetgate * cx) + (ingate * cellgate) hy = outgate * torch.tanh(cy)return hy, (hy, cy)
TorchScript生成的此圖形表示（IR）可實現多種優化和可伸縮計算。 除了可以做的典型編譯器優化（CSE，常量傳播等）之外，還可以運行其他IR轉換以使代碼運行得更快。
LSTM層(前向)
classLSTMLayer(jit.ScriptModule):def__init__(self, cell, *cell_args):super(LSTMLayer, self).__init__() self.cell = cell(*cell_args) @jit.script_methoddefforward(self, input, state):# type: (Tensor, Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]] inputs = input.unbind(0) outputs = torch.jit.annotate(List[Tensor], [])for i in range(len(inputs)): out, state = self.cell(inputs[i], state) outputs += [out]return torch.stack(outputs), state
在為TorchScript LSTM生成的IR上做了一些技巧來提高性能，團隊做了一些示例優化：
循環展開(Loop Unrolling)：自動在代碼中展開循環（對於大循環，展開它的一小部分），然後授權對for循環控制流進行進一步的優化。 例如，fuser可以將循環體的迭代中的操作融合在一起，這導致對於諸如LSTM的控制流密集型模型的良好性能改進。批量矩陣乘法：對於輸入預乘的RNN（即模型具有大量相同LHS或RHS的矩陣乘法），可以將這些操作一起有效地批量處理為單個矩陣乘法，同時對輸出進行分塊以實現等效語義。
通過應用這些技術，將前向傳播的時間減少了1.6ms，達到8.4ms(1.2倍加速)，後向傳播的時間減少了7ms，達到20ms左右(1.35倍加速)。
LSTM層(後向)
「樹結構」批處理矩陣Muplication：通常情況是在LSTM反向圖中多次重複使用單個權重，形成一個樹，其中葉子是矩陣乘法，節點是相加的。 這些節點可以通過在不同維度上連接LHS和RHS來組合在一起，然後計算為單個矩陣乘法。 等價公式可表示如下：
$L1 * R1 + L2 * R2 = torch.cat((L1, L2), dim=1) * torch.cat((R1, R2), dim=0)$
Autograd是使PyTorch成為如此優雅的ML框架的關鍵組件。因此，將其應用到PyTorch JIT，但是使用了一種新的自動微分(AD)機制，該機制在IR級別上工作。JIT自動微分將把正向圖分割成符號可微分的子圖，並為這些子圖生成向後節點。以上面的IR為例，對於具有AD公式的操作，我們將圖節點分組為一個prim :: DifferentiableGraph_0。對於沒有添加到AD公式中的操作，我們將在執行期間返回到Autograd。優化反向路徑是困難的，隱式broadcasting語義使得自動微分的優化更加困難。 PyTorch可以方便地編寫張量操作，而無需通過broadcasting張量來擔心形狀。 對於性能而言，反向的痛點是需要對這種可broadcasting操作進行求和。 這導致每個可broadcasting操作的導數後跟一個求和。 由於目前無法融合減少操作，這會導致FusionGroups分成多個小組，從而導致性能下降。 要解決這個問題，請參閱Thomas Viehmann撰寫的文章：http://lernapparat.de/fast-lstm-pytorch/。
更多這方面的優化內容可參考Pytorch團隊博客原文：
https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/
更多新功能
運算符
torch.tril_indices, torch.triu_indices:添加了與NumPy具有相同行為的運算符；torch.combinations, torch.cartesian_prod: 添加了類似於itertools的新運算符；torch.repeat_interleave: 新運算符類似於numpy.repeat；torch.from_file:類似於Storage.from_file的新運算符，但返回一個張量；torch.unique_consecutive: 新的運算符，其語義類似於C ++中的std :: unique；torch.tril, torch.triu, torch.trtrs:現在支持批處理；torch.gather: 添加對sparse_grad選項的支持；torch.std, torch.max_values, torch.min_values, torch.logsumexp現在可以同時在多個維度上運行；torch.cdist:添加了與scipy.spatial.distance.cdist等效的運算符；torch.__config__.show():報告所有庫的詳細版本。
NN
nn.MultiheadedAttention:從注意力中實現MultiheadedAttention的新模塊；nn.functional.interpolate:增加了對bicubic的支持；nn.SyncBatchNorm:支持同步批量標準化；nn.Conv:通過mode ='circular'添加了對Circular Padding的支持；nn.EmbeddingBag: 現在支持可訓練的`per_sample_weights；nn.EmbeddingBag:添加對from_pretrained方法的支持，如nn.Embedding中所示；RNNs:通過enforce_sorted自動處理未排序的可變長度序列；nn.Identity:便於模型surgery的新模塊。
更多有關張量/dtypes、性能提高、bug修復、棄用的項目等內容可查看Pytorch在GitHub發布的項目原文：
https://github.com/pytorch/pytorch/releases/tag/v1.1.0
新智元春季招聘開啟，一起弄潮 AI 之巔！
崗位詳情請戳：