PyTorch經驗指南：技巧與陷阱

來自專欄機器之心81 人贊了文章

選自GitHub，作者：Kaixhin，機器之心編譯。

PyTorch 的構建者表明，PyTorch 的哲學是解決當務之急，也就是說即時構建和運行計算圖。目前，PyTorch 也已經藉助這種即時運行的概念成為最受歡迎的框架之一，開發者能快速構建模型與驗證想法，並通過神經網路交換格式 ONNX 在多個框架之間快速遷移。本文從基本概念開始介紹了 PyTorch 的使用方法、訓練經驗與技巧，並展示了可能出現的問題與解決方案。

項目地址：https://github.com/Kaixhin/grokking-pytorch

PyTorch 是一種靈活的深度學習框架，它允許通過動態神經網路（例如利用動態控流——如 if 語句或 while 循環的網路）進行自動微分。它還支持 GPU 加速、分散式訓練以及各類優化任務，同時還擁有許多更簡潔的特性。以下是作者關於如何利用 PyTorch 的一些說明，裡面雖然沒有包含該庫的所有細節或最優方法，但可能會對大家有所幫助。

神經網路是計算圖的一個子類。計算圖接收輸入數據，數據被路由到對數據執行處理的節點，並可能被這些節點轉換。在深度學習中，神經網路中的神經元（節點）通常利用參數或可微函數轉換數據，這樣可以優化參數以通過梯度下降將損失最小化。更廣泛地說，函數是隨機的，圖結構可以是動態的。所以說，雖然神經網路可能非常適合數據流式編程，但 PyTorch 的 API 卻更關注命令式編程——一種編程更常考慮的形式。這令讀取代碼和推斷複雜程序變得簡單，而無需損耗不必要的性能；PyTorch 速度很快，且擁有大量優化，作為終端用戶你毫無後顧之憂。

本文其餘部分寫的是關於 grokking PyTorch 的內容，都是基於 MINIST 官網實例，應該要在學習完官網初學者教程後再查看。為便於閱讀，代碼以塊狀形式呈現，並帶有注釋，因此不會像純模塊化代碼一樣被分割成不同的函數或文件。

Pytorch 基礎

PyTorch 使用一種稱之為 imperative / eager 的範式，即每一行代碼都要求構建一個圖，以定義完整計算圖的一個部分。即使完整的計算圖還沒有構建好，我們也可以獨立地執行這些作為組件的小計算圖，這種動態計算圖被稱為「define-by-run」方法。

PyTorch 張量

正如 PyTorch 文檔所說，如果我們熟悉 NumPy 的多維數組，那麼 Torch 張量的很多操作我們能輕易地掌握。PyTorch 提供了 CPU 張量和 GPU 張量，並且極大地加速了計算的速度。

從張量的構建與運行就能體會，相比 TensorFLow，在 PyTorch 中聲明張量、初始化張量要簡潔地多。例如，使用 torch.Tensor(5, 3) 語句就能隨機初始化一個 5×3 的二維張量，因為 PyTorch 是一種動態圖，所以它聲明和真實賦值是同時進行的。

在 PyTorch 中，torch.Tensor 是一種多維矩陣，其中每個元素都是單一的數據類型，且該構造函數默認為 torch.FloatTensor。以下是具體的張量類型：

除了直接定義維度，一般我們還可以從 Python 列表或 NumPy 數組中創建張量。而且根據使用 Python 列表和元組等數據結構的習慣，我們可以使用相似的索引方式進行取值或賦值。PyTorch 同樣支持廣播（Broadcasting）操作，一般它會隱式地把一個數組的異常維度調整到與另一個運算元相匹配的維度，以實現維度兼容。

自動微分模塊

TensorFlow、Caffe 和 CNTK 等大多數框架都使用靜態計算圖，開發者必須建立或定義一個神經網路，並重複使用相同的結構來執行模型訓練。改變網路的模式就意味著我們必須從頭開始設計並定義相關的模塊。

但 PyTorch 使用的技術為自動微分（automatic differentiation）。在這種機制下，系統會有一個 Recorder 來記錄我們執行的運算，然後再反向計算對應的梯度。這種技術在構建神經網路的過程中十分強大，因為我們可以通過計算前向傳播過程中參數的微分來節省時間。

從概念上來說，Autograd 會維護一個圖並記錄對變數執行的所有運算。這會產生一個有向無環圖，其中葉結點為輸入向量，根結點為輸出向量。通過從根結點到葉結點追蹤圖的路徑，我們可以輕易地使用鏈式法則自動計算梯度。

在內部，Autograd 將這個圖表徵為 Function 對象的圖，並且可以應用 apply() 計算評估圖的結果。在計算前向傳播中，當 Autograd 在執行請求的計算時，它還會同時構建一個表徵梯度計算的圖，且每個 Variable 的 .grad_fn 屬性就是這個圖的輸入單元。在前向傳播完成後，我們可以在後向傳播中根據這個動態圖來計算梯度。

PyTorch 還有很多基礎的模塊，例如控制學習過程的最優化器、搭建深度模型的神經網路模塊和數據載入與處理等。這一節展示的張量與自動微分模塊是 PyTorch 最為核心的概念之一，讀者可查閱 PyTorch 文檔了解更詳細的內容。

下面作者以 MNIST 為例從數據載入到模型測試具體討論了 PyTorch 的使用、思考技巧與陷阱。

PyTorch 實用指南

導入

import argparseimport torchfrom torch import nn, optimfrom torch.nn import functional as Ffrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms

除了用於計算機視覺問題的 torchvision 模塊外，這些都是標準化導入。

設置

parser = argparse.ArgumentParser(description=PyTorch MNIST Example)parser.add_argument(--batch-size, type=int, default=64, metavar=N, help=input batch size for training (default: 64))parser.add_argument(--epochs, type=int, default=10, metavar=N, help=number of epochs to train (default: 10))parser.add_argument(--lr, type=float, default=0.01, metavar=LR, help=learning rate (default: 0.01))parser.add_argument(--momentum, type=float, default=0.5, metavar=M, help=SGD momentum (default: 0.5))parser.add_argument(--no-cuda, action=store_true, default=False, help=disables CUDA training)parser.add_argument(--seed, type=int, default=1, metavar=S, help=random seed (default: 1))parser.add_argument(--save-interval, type=int, default=10, metavar=N, help=how many batches to wait before checkpointing)parser.add_argument(--resume, action=store_true, default=False, help=resume training from checkpoint)args = parser.parse_args()use_cuda = torch.cuda.is_available() and not args.no_cudadevice = torch.device(cuda if use_cuda else cpu)torch.manual_seed(args.seed)if use_cuda: torch.cuda.manual_seed(args.seed)

argparse 是在 Python 中處理命令行參數的一種標準方式。

編寫與設備無關的代碼（可用時受益於 GPU 加速，不可用時會倒退回 CPU）時，選擇並保存適當的 torch.device, 不失為一種好方法，它可用於確定存儲張量的位置。關於與設備無關代碼的更多內容請參閱官網文件。PyTorch 的方法是使用戶能控制設備，這對簡單示例來說有些麻煩，但是可以更容易地找出張量所在的位置——這對於 a）調試很有用，並且 b）可有效地使用手動化設備。

對於可重複實驗，有必要為使用隨機數生成的任何數據設置隨機種子（如果也使用隨機數，則包括隨機或 numpy）。要注意，cuDNN 用的是非確定演算法，可以通過語句 torch.backends.cudnn.enabled = False 將其禁用。

數據

train_data = datasets.MNIST(data, train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))test_data = datasets.MNIST(data, train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)test_loader = DataLoader(test_data, batch_size=args.batch_size, num_workers=4, pin_memory=True)

torchvision.transforms 對於單張圖像有非常多便利的轉換工具，例如裁剪和歸一化等。

DataLoader 包含非常多的參數，除了 batch_size 和 shuffle，num_workers 和 pin_memory 對於高效載入數據同樣非常重要。例如配置 num_workers > 0 將使用子進程非同步載入數據，而不是使用一個主進程塊載入數據。參數 pin_memory 使用固定 RAM 以加速 RAM 到 GPU 的轉換，且在僅使用 CPU 時不會做任何運算。

模型

class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = self.fc2(x) return F.log_softmax(x, dim=1)model = Net().to(device)optimiser = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)if args.resume: model.load_state_dict(torch.load(model.pth)) optimiser.load_state_dict(torch.load(optimiser.pth))

神經網路初始化一般包括變數、包含可訓練參數的層級、可能獨立的可訓練參數和不可訓練的緩存器。隨後前向傳播將這些初始化參數與 F 中的函數結合，其中該函數為不包含參數的純函數。有些開發者喜歡使用完全函數化的網路（如保持所有參數獨立，使用 F.conv2d 而不是 nn.Conv2d），或者完全由 layers 函數構成的網路（如使用 nn.ReLU 而不是 F.relu）。

在將 device 設置為 GPU 時，.to(device) 是一種將設備參數（和緩存器）發送到 GPU 的便捷方式，且在將 device 設置為 CPU 時不會做任何處理。在將網路參數傳遞給優化器之前，把它們傳遞給適當的設備非常重要，不然的話優化器不能正確地追蹤參數。

神經網路（nn.Module）和優化器（optim.Optimizer）都能保存和載入它們的內部狀態，而.load_state_dict(state_dict) 是完成這一操作的推薦方法，我們可以從以前保存的狀態字典中載入兩者的狀態並恢復訓練。此外，保存整個對象可能會出錯。

這裡沒討論的一些注意事項即前向傳播可以使用控制流，例如一個成員變數或數據本身能決定 if 語句的執行。此外，在前向傳播的過程中列印張量也是可行的，這令 debug 更加簡單。最後，前向傳播可以使用多個參數。以下使用間斷的代碼塊展示這一點：

def forward(self, x, hx, drop=False): hx2 = self.rnn(x, hx) print(hx.mean().item(), hx.var().item()) if hx.max.item() > 10 or self.can_drop and drop: return hx else: return hx2

訓練

model.train()train_losses = []for i, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimiser.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() train_losses.append(loss.item()) optimiser.step() if i % 10 == 0: print(i, loss.item()) torch.save(model.state_dict(), model.pth) torch.save(optimiser.state_dict(), optimiser.pth) torch.save(train_losses, train_losses.pth)

網路模塊默認設置為訓練模式，這影響了某些模塊的工作方式，最明顯的是 dropout 和批歸一化。最好用.train() 對其進行手動設置，這樣可以把訓練標記向下傳播到所有子模塊。

在使用 loss.backward() 收集一系列新的梯度以及用 optimiser.step() 做反向傳播之前，有必要手動地將由 optimiser.zero_grad() 優化的參數梯度歸零。默認情況下，PyTorch 會累加梯度，在單次迭代中沒有足夠資源來計算所有需要的梯度時，這種做法非常便利。

PyTorch 使用一種基於 tape 的自動化梯度（autograd）系統，它收集按順序在張量上執行的運算，然後反向重放它們來執行反向模式微分。這正是為什麼 PyTorch 如此靈活並允許執行任意計算圖的原因。如果沒有張量需要做梯度更新（當你需要為該過程構建一個張量時，你必須設置 requires_grad=True），則不需要保存任何圖。然而，網路傾向於包含需要梯度更新的參數，因此任何網路輸出過程中執行的計算都將保存在圖中。因此如果想保存在該過程中得到的數據，你將需要手動禁止梯度更新，或者，更常見的做法是將其保存為一個 Python 數（通過一個 Python 標量上的.item()）或者 NumPy 數組。更多關於 autograd 的細節詳見官網文件。

截取計算圖的一種方式是使用.detach()，當通過沿時間的截斷反向傳播訓練 RNN 時，數據流傳遞到一個隱藏狀態可能會應用這個函數。當對損失函數求微分（其中一個成分是另一個網路的輸出）時，也會很方便。但另一個網路不應該用「loss - examples」的模式進行優化，包括在 GAN 訓練中從生成器的輸出訓練判別器，或使用價值函數作為基線（例如 A2C）訓練 actor-critic 演算法的策略。另一種在 GAN 訓練（從判別器訓練生成器）中能高效阻止梯度計算的方法是在整個網路參數上建立循環，並設置 param.requires_grad=False，這在微調中也很常用。

除了在控制台/日誌文件里記錄結果以外，檢查模型參數（以及優化器狀態）也是很重要的。你還可以使用 torch.save() 來保存一般的 Python 對象，但其它標準選擇還包括內建的 pickle。

測試

model.eval()test_loss, correct = 0, 0with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, size_average=False).item() pred = output.argmax(1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item()test_loss /= len(test_data)acc = correct / len(test_data)print(acc, test_loss)

為了早點響應.train()，應利用.eval() 將網路明確地設置為評估模式。

正如前文所述，計算圖通常會在使用網路時生成。通過 with torch.no_grad() 使用 no_grad 上下文管理器，可以防止這種情況發生。

其它

內存有問題？可以查看官網文件獲取幫助。

CUDA 出錯？它們很難調試，而且通常是一個邏輯問題，會在 CPU 上產生更易理解的錯誤信息。如果你計劃使用 GPU，那最好能夠在 CPU 和 GPU 之間輕鬆切換。更普遍的開發技巧是設置代碼，以便在啟動合適的項目（例如準備一個較小/合成的數據集、運行一個 train + test epoch 等）之前快速運行所有邏輯來檢查它。如果這是一個 CUDA 錯誤，或者你沒法切換到 CPU，設置 CUDA_LAUNCH_BLOCKING=1 將使 CUDA 內核同步啟動，從而提供更詳細的錯誤信息。

torch.multiprocessing，甚至只是一次運行多個 PyTorch 腳本的注意事項。因為 PyTorch 使用多線程 BLAS 庫來加速 CPU 上的線性代數計算，所以它通常需要使用多個內核。如果你想一次運行多個任務，在具有多進程或多個腳本的情況下，通過將環境變數 OMP_NUM_THREADS 設置為 1 或另一個較小的數字來手動減少線程，這樣做減少了 CPU thrashing 的可能性。官網文件還有一些其它注意事項，尤其是關於多進程。