點擊上方“算法猿的成長“，關注公眾號，選擇加“星標“或“置頂”

總第 132?篇文章，本文大約 7000?字，閱讀大約需要 20?分鐘

原文：https://github.com/vahidk/EffectivePyTorch

作者：vahidk

前言

這是一份 PyTorch 教程和最佳實踐筆記，目錄如下所示：

PyTorch 基礎

將模型封裝為模塊

廣播機制的優缺點

使用好重載的運算符

采用 TorchScript 優化運行時間

構建高效的自定義數據加載類

PyTorch 的數值穩定性

因為原文太長所以分為上下兩篇文章進行介紹，本文介紹前四點，從基礎開始介紹到使用重載的運算符。

首先 PyTorch 的安裝可以根據官方文檔進行操作：

https://pytorch.org/

pip?install?torch?torchvision

---

1. PyTorch 基礎

PyTorch 是數值計算方面其中一個最流行的庫，同時也是機器學習研究方面最廣泛使用的框架。在很多方面，它和 NumPy 都非常相似，但是它可以在不需要代碼做多大改變的情況下，在 CPUs，GPUs，TPUs 上實現計算，以及非常容易實現分布式計算的操作。PyTorch 的其中一個最重要的特征就是自動微分。它可以讓需要采用梯度下降算法進行訓練的機器學習算法的實現更加方便，可以更高效的自動計算函數的梯度。我們的目標是提供更好的 PyTorch 介紹以及討論使用 PyTorch 的一些最佳實踐。

對于 PyTorch 第一個需要學習的就是張量（Tensors）的概念，張量就是多維數組，它和 numpy 的數組非常相似，但多了一些函數功能。

一個張量可以存儲一個標量數值、一個數組、一個矩陣：

import?torch #?標量數值 a?=?torch.tensor(3) print(a)??#?tensor(3) #?數組 b?=?torch.tensor([1,?2]) print(b)??#?tensor([1,?2]) #?矩陣 c?=?torch.zeros([2,?2]) print(c)??#?tensor([[0.,?0.],?[0.,?0.]]) #?任意維度的張量 d?=?torch.rand([2,?2,?2])

張量還可以高效的執行代數的運算。機器學習應用中最常見的運算就是矩陣乘法。例如希望將兩個隨機矩陣進行相乘，維度分別是 和 ，這個運算可以通過矩陣相乘運算實現（@）：

import?torchx?=?torch.randn([3,?5]) y?=?torch.randn([5,?4]) z?=?x?@?yprint(z)

對于向量相加，如下所示：

z?=?x?+?y

將張量轉換為 numpy 數組，可以調用 numpy() 方法：

print(z.numpy())

當然，反過來 numpy 數組轉換為張量是可以的：

x?=?torch.tensor(np.random.normal([3,?5]))

自動微分

PyTorch 中相比 numpy ?最大優點就是可以實現自動微分，這對于優化神經網絡參數的應用非常有幫助。下面通過一個例子來幫助理解這個優點。

假設現在有一個復合函數：g(u(x)) ，為了計算 g 對 x 的導數，這里可以采用鏈式法則，即

而 PyTorch 可以自動實現這個求導的過程。

為了在 PyTorch 中計算導數，首先要創建一個張量，并設置其 requires_grad = True ，然后利用張量運算來定義函數，這里假設 u 是一個二次方的函數，而 g 是一個簡單的線性函數，代碼如下所示：

x?=?torch.tensor(1.0,?requires_grad=True)def?u(x):return?x?*?xdef?g(u):return?-u

在這個例子中，復合函數就是 ，所以導數是 ，如果 x=1 ，那么可以得到 -2 。

在 PyTorch 中調用梯度函數：

dgdx?=?torch.autograd.grad(g(u(x)),?x)[0] print(dgdx)??#?tensor(-2.)

擬合曲線

為了展示自動微分有多么強大，這里介紹另一個例子。

首先假設我們有一些服從一個曲線（也就是函數 ）的樣本，然后希望基于這些樣本來評估這個函數 f(x) 。我們先定義一個帶參數的函數:

函數的輸入是 x，然后 w 是參數，目標是找到合適的參數使得下列式子成立：

實現的一個方法可以是通過優化下面的損失函數來實現：

盡管這個問題里有一個正式的函數（即 f(x) 是一個具體的函數），但這里我們還是采用一個更加通用的方法，可以應用到任何一個可微分的函數，并采用隨機梯度下降法，即通過計算 L(w) 對于每個參數 w 的梯度的平均值，然后不斷從相反反向移動。

利用 PyTorch 實現的代碼如下所示：

import?numpy?as?np import?torch#?Assuming?we?know?that?the?desired?function?is?a?polynomial?of?2nd?degree,?we #?allocate?a?vector?of?size?3?to?hold?the?coefficients?and?initialize?it?with #?random?noise. w?=?torch.tensor(torch.randn([3,?1]),?requires_grad=True)#?We?use?the?Adam?optimizer?with?learning?rate?set?to?0.1?to?minimize?the?loss. opt?=?torch.optim.Adam([w],?0.1)def?model(x):#?We?define?yhat?to?be?our?estimate?of?y.f?=?torch.stack([x?*?x,?x,?torch.ones_like(x)],?1)yhat?=?torch.squeeze(f?@?w,?1)return?yhatdef?compute_loss(y,?yhat):#?The?loss?is?defined?to?be?the?mean?squared?error?distance?between?our#?estimate?of?y?and?its?true?value.?loss?=?torch.nn.functional.mse_loss(yhat,?y)return?lossdef?generate_data():#?Generate?some?training?data?based?on?the?true?functionx?=?torch.rand(100)?*?20?-?10y?=?5?*?x?*?x?+?3return?x,?ydef?train_step():x,?y?=?generate_data()yhat?=?model(x)loss?=?compute_loss(y,?yhat)opt.zero_grad()loss.backward()opt.step()for?_?in?range(1000):train_step()print(w.detach().numpy())

運行上述代碼，可以得到和下面相近的結果：

[4.9924135, 0.00040895029, 3.4504161]

這和我們的參數非常接近。

上述只是 PyTorch 可以做的事情的冰山一角。很多問題，比如優化一個帶有上百萬參數的神經網絡，都可以用 PyTorch 高效的用幾行代碼實現，PyTorch 可以跨多個設備和線程進行拓展，并且支持多個平臺。

---

2. 將模型封裝為模塊

在之前的例子中，我們構建模型的方式是直接實現張量間的運算操作。但為了讓代碼看起來更加有組織，推薦采用 PyTorch 的 modules 模塊。一個模塊實際上是一個包含參數和壓縮模型運算的容器。

比如，如果想實現一個線性模型 ，那么實現的代碼可以如下所示：

import?torchclass?Net(torch.nn.Module):def?__init__(self):super().__init__()self.a?=?torch.nn.Parameter(torch.rand(1))self.b?=?torch.nn.Parameter(torch.rand(1))def?forward(self,?x):yhat?=?self.a?*?x?+?self.breturn?yhat

使用的例子如下所示，需要實例化聲明的模型，并且像調用函數一樣使用它：

x?=?torch.arange(100,?dtype=torch.float32)net?=?Net() y?=?net(x)

參數都是設置 requires_grad 為 true 的張量。通過模型的 parameters() 方法可以很方便的訪問和使用參數，如下所示：

for?p?in?net.parameters():print(p)

現在，假設是一個未知的函數 y=5x+3+n ，注意這里的 n 是表示噪音，然后希望優化模型參數來擬合這個函數，首先可以簡單從這個函數進行采樣，得到一些樣本數據：

x?=?torch.arange(100,?dtype=torch.float32)?/?100 y?=?5?*?x?+?3?+?torch.rand(100)?*?0.3

和上一個例子類似，需要定義一個損失函數并優化模型的參數，如下所示：

criterion?=?torch.nn.MSELoss() optimizer?=?torch.optim.SGD(net.parameters(),?lr=0.01)for?i?in?range(10000):net.zero_grad()yhat?=?net(x)loss?=?criterion(yhat,?y)loss.backward()optimizer.step()print(net.a,?net.b)?#?Should?be?close?to?5?and?3

在 PyTorch 中已經實現了很多預定義好的模塊。比如 torch.nn.Linear 就是一個類似上述例子中定義的一個更加通用的線性函數，所以我們可以采用這個函數來重寫我們的模型代碼，如下所示：

class?Net(torch.nn.Module):def?__init__(self):super().__init__()self.linear?=?torch.nn.Linear(1,?1)def?forward(self,?x):yhat?=?self.linear(x.unsqueeze(1)).squeeze(1)return?yhat

這里用到了兩個函數，squeeze 和 unsqueeze ，主要是torch.nn.Linear 會對一批向量而不是數值進行操作。

同樣，默認調用 parameters() 會返回其所有子模塊的參數：

net?=?Net() for?p?in?net.parameters():print(p)

當然也有一些預定義的模塊是作為包容其他模塊的容器，最常用的就是 torch.nn.Sequential ，它的名字就暗示了它主要用于堆疊多個模塊（或者網絡層），例如堆疊兩個線性網絡層，中間是一個非線性函數 ReLU ，如下所示：

model?=?torch.nn.Sequential(torch.nn.Linear(64,?32),torch.nn.ReLU(),torch.nn.Linear(32,?10), )

---

3. 廣播機制的優缺點

優點

PyTorch 支持廣播的元素積運算。正常情況下，當想執行類似加法和乘法操作的時候，你需要確認操作數的形狀是匹配的，比如無法進行一個 [3, 2] 大小的張量和 [3, 4] 大小的張量的加法操作。

但是存在一種特殊的情況：只有單一維度的時候，PyTorch 會隱式的根據另一個操作數的維度來拓展只有單一維度的操作數張量。因此，實現 [3,2] 大小的張量和 [3,1] 大小的張量相加的操作是合法的。

如下代碼展示了一個加法的例子：

import?torcha?=?torch.tensor([[1.,?2.],?[3.,?4.]]) b?=?torch.tensor([[1.],?[2.]]) #?c?=?a?+?b.repeat([1,?2]) c?=?a?+?bprint(c)

廣播機制可以實現隱式的維度復制操作（repeat 操作），并且代碼更短，內存使用上也更加高效，因為不需要存儲復制的數據的結果。這個機制非常適合用于結合多個維度不同的特征的時候。

為了拼接不同維度的特征，通常的做法是先對輸入張量進行維度上的復制，然后拼接后使用非線性激活函數。整個過程的代碼實現如下所示：

a?=?torch.rand([5,?3,?5]) b?=?torch.rand([5,?1,?6])linear?=?torch.nn.Linear(11,?10)#?concat?a?and?b?and?apply?nonlinearity tiled_b?=?b.repeat([1,?3,?1])?#?b?shape:??[5,?3,?6] c?=?torch.cat([a,?tiled_b],?2)?#?c?shape:?[5,?3,?11] d?=?torch.nn.functional.relu(linear(c))print(d.shape)??#?torch.Size([5,?3,?10])

但實際上通過廣播機制可以實現得更加高效，即 f(m(x+y)) 是等同于 f(mx+my) 的，也就是我們可以先分別做線性操作，然后通過廣播機制來做隱式的拼接操作，如下所示：

a?=?torch.rand([5,?3,?5]) b?=?torch.rand([5,?1,?6])linear1?=?torch.nn.Linear(5,?10) linear2?=?torch.nn.Linear(6,?10)pa?=?linear1(a)?#?pa?shape:?[5,?3,?10] pb?=?linear2(b)?#?pb?shape:?[5,?1,?10] d?=?torch.nn.functional.relu(pa?+?pb)print(d.shape)??#?torch.Size([5,?3,?10])

實際上這段代碼非常通用，可以用于任意維度大小的張量，只要它們之間是可以實現廣播機制的，如下所示：

class?Merge(torch.nn.Module):def?__init__(self,?in_features1,?in_features2,?out_features,?activation=None):super().__init__()self.linear1?=?torch.nn.Linear(in_features1,?out_features)self.linear2?=?torch.nn.Linear(in_features2,?out_features)self.activation?=?activationdef?forward(self,?a,?b):pa?=?self.linear1(a)pb?=?self.linear2(b)c?=?pa?+?pbif?self.activation?is?not?None:c?=?self.activation(c)return?c

缺點

到目前為止，我們討論的都是廣播機制的優點。但它的缺點是什么呢？原因也是出現在隱式的操作，這種做法非常不利于進行代碼的調試。

這里給出一個代碼例子：

a?=?torch.tensor([[1.],?[2.]]) b?=?torch.tensor([1.,?2.]) c?=?torch.sum(a?+?b)print(c)

所以上述代碼的輸出結果 c 是什么呢？你可能覺得是 6，但這是錯的，正確答案是 12 。這是因為當兩個張量的維度不匹配的時候，PyTorch 會自動將維度低的張量的第一個維度進行拓展，然后在進行元素之間的運算，所以這里會將b ?先拓展為 [[1, 2], [1, 2]]，然后 a+b 的結果應該是 [[2,3], [3, 4]] ，然后sum 操作是將所有元素求和得到結果 12。

那么避免這種結果的方法就是顯式的操作，比如在這個例子中就需要指定好想要求和的維度，這樣進行代碼調試會更簡單，代碼修改后如下所示：

a?=?torch.tensor([[1.],?[2.]]) b?=?torch.tensor([1.,?2.]) c?=?torch.sum(a?+?b,?0)print(c)

這里得到的 c 的結果是 [5, 7]，而我們基于結果的維度可以知道出現了錯誤。

這有個通用的做法，就是在做累加（ reduction ）操作或者使用 torch.squeeze 的時候總是指定好維度。

---

4. 使用好重載的運算符

和 NumPy 一樣，PyTorch 會重載 python 的一些運算符來讓 PyTorch 代碼更簡短和更有可讀性。

例如，切片操作就是其中一個重載的運算符，可以更容易的對張量進行索引操作，如下所示：

z?=?x[begin:end]??#?z?=?torch.narrow(0,?begin,?end-begin)

但需要謹慎使用這個運算符，它和其他運算符一樣，也有一些副作用。正因為它是一個非常常用的運算操作，如果過度使用可以導致代碼變得低效。

這里給出一個例子來展示它是如何導致代碼變得低效的。這個例子中我們希望對一個矩陣手動實現行之間的累加操作：

import?torch import?timex?=?torch.rand([500,?10])z?=?torch.zeros([10])start?=?time.time() for?i?in?range(500):z?+=?x[i] print("Took?%f?seconds."?%?(time.time()?-?start))

上述代碼的運行速度會非常慢，因為總共調用了 500 次的切片操作，這就是過度使用了。一個更好的做法是采用 torch.unbind 運算符在每次循環中將矩陣切片為一個向量的列表，如下所示：

z?=?torch.zeros([10]) for?x_i?in?torch.unbind(x):z?+=?x_i

這個改進會提高一些速度（在作者的機器上是提高了大約30%）。

但正確的做法應該是采用 torch.sum 來一步實現累加的操作：

z?=?torch.sum(x,?dim=0)

這種實現速度就非常的快（在作者的機器上提高了100%的速度）。

其他重載的算數和邏輯運算符分別是：

z?=?-x??#?z?=?torch.neg(x) z?=?x?+?y??#?z?=?torch.add(x,?y) z?=?x?-?y z?=?x?*?y??#?z?=?torch.mul(x,?y) z?=?x?/?y??#?z?=?torch.div(x,?y) z?=?x?//?y z?=?x?%?y z?=?x?**?y??#?z?=?torch.pow(x,?y) z?=?x?@?y??#?z?=?torch.matmul(x,?y) z?=?x?>?y z?=?x?>=?y z?=?x?<?y z?=?x?<=?y z?=?abs(x)??#?z?=?torch.abs(x) z?=?x?&?y z?=?x?|?y z?=?x?^?y??#?z?=?torch.logical_xor(x,?y) z?=?~x??#?z?=?torch.logical_not(x) z?=?x?==?y??#?z?=?torch.eq(x,?y) z?=?x?!=?y??#?z?=?torch.ne(x,?y)

還可以使用這些運算符的遞增版本，比如 x += y 和 x **=2 都是合法的。

另外，Python 并不允許重載 and 、or 和 not 三個關鍵詞。

---

精選AI文章

1.?10個實用的機器學習建議

2.?深度學習算法簡要綜述(上)

3.?深度學習算法簡要綜述(上)

4.?常見的數據增強項目和論文介紹

5.?實戰|手把手教你訓練一個基于Keras的多標簽圖像分類器

精選python文章

1.??python數據模型

2.?python版代碼整潔之道

3.?快速入門 Jupyter notebook

4.?Jupyter 進階教程

5.?10個高效的pandas技巧

精選教程資源文章

1.?[資源分享] TensorFlow 官方中文版教程來了

2.?[資源]推薦一些Python書籍和教程，入門和進階的都有！

3.?[Github項目推薦] 推薦三個助你更好利用Github的工具

4.?Github上的各大高校資料以及國外公開課視頻

5.?GitHub上有哪些比較好的計算機視覺/機器視覺的項目？

歡迎關注我的微信公眾號--算法猿的成長，或者掃描下方的二維碼，大家一起交流，學習和進步！

?

如果覺得不錯，在看、轉發就是對小編的一個支持！

創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎

總結

以上是生活随笔為你收集整理的编写高效的PyTorch代码技巧（上）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。