專案地址: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 argparse
import torch
from torch import nn, optim
from torch.nn import functional as F
from torch.utils.data import DataLoader
from 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_cuda
device = 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, 0
with 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 的可能性。官網檔案還有一些其它注意事項,尤其是關於多程式。