盼望已久，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為例：

#getinputsandstatesforLSTMCellinputs=get_lstm_inputs()#instantiateaScriptModulecell=LSTMCell(input_size,hidden_size)#printtheoptimizedgraphusinggraph_forout=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)withprim::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)...somebroadcastsizesoperations...%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)...somebroadcastsizesoperations...return(%hy,%cy,%49,%44,%196,%199,%340,%192,%325,%185,%ingate.1,%forgetgate.1,%cellgate.1,%outgate.1,%395,%396,%287)withprim::FusionGroup_0=graph(%13:Float(*,*),%71:Tensor,%76:Tensor,%81:Tensor,%86:Tensor):...somechunks,constants,andaddoperations...%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],eachTensorisT'xCpadded=torch.utils.rnn.pad_sequence(sequences)lengths=[seq.size(0)forseqinsequences]padded#TxNxC,whereNisbatchsizeandTisthemaxofallT'model=LSTM(...)output,hiddens=model(padded)output#TxNxC

當然，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_sizeself.hidden_size=hidden_sizeself.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=stategates=(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)returnhy,(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],[])foriinrange(len(inputs)):out,state=self.cell(inputs[i],state)outputs+=[out]returntorch.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