PyTorch最佳實踐，怎樣才能寫出一手風格優美的代碼

PyTorch是最優秀的深度學習框架之一，它簡單優雅，非常適合入門。本文將介紹PyTorch的最佳實踐和代碼風格都是怎樣的。

雖然這是一個非官方的 PyTorch指南，但本文總結了一年多使用 PyTorch 框架的經驗，尤其是用它開發深度學習相關工作的最優解決方案。請注意，我們分享的經驗大多是從研究和實踐角度出發的。

這是一個開發的項目，歡迎其它讀者改進該文檔：https://github.com/IgorSusmelj/pytorch-styleguide。

本文檔主要由三個部分構成：首先，本文會簡要清點 Python 中的最好裝備。接著，本文會介紹一些使用 PyTorch 的技巧和建議。最後，我們分享了一些使用其它框架的見解和經驗，這些框架通常幫助我們改進工作流。

清點 Python 裝備

建議使用 Python 3.6 以上版本

根據我們的經驗，我們推薦使用 Python 3.6 以上的版本，因為它們具有以下特性，這些特性可以使我們很容易寫出簡潔的代碼：

• 自 Python 3.6 以後支持「typing」模塊
• 自 Python 3.6 以後支持格式化字元串（f string）

Python 風格指南

我們試圖遵循 Google 的 Python 編程風格。請參閱 Google 提供的優秀的 python 編碼風格指南：

地址：https://github.com/google/styleguide/blob/gh-pages/pyguide.md。

在這裡，我們會給出一個最常用命名規範小結：

集成開發環境

一般來說，我們建議使用 visual studio 或 PyCharm這樣的集成開發環境。而 VS Code 在相對輕量級的編輯器中提供語法高亮和自動補全功能，PyCharm 則擁有許多用於處理遠程集群任務的高級特性。

Jupyter Notebooks VS Python 腳本

一般來說，我們建議使用 Jupyter Notebook 進行初步的探索，或嘗試新的模型和代碼。如果你想在更大的數據集上訓練該模型，就應該使用 Python 腳本，因為在更大的數據集上，復現性更加重要。

我們推薦你採取下面的工作流程：

• 在開始的階段，使用 Jupyter Notebook
• 對數據和模型進行探索
• 在 notebook 的單元中構建你的類/方法
• 將代碼移植到 Python 腳本中
• 在伺服器上訓練/部署

開發常備庫

常用的程序庫有：

文件組織

不要將所有的層和模型放在同一個文件中。最好的做法是將最終的網路分離到獨立的文件（networks.py）中，並將層、損失函數以及各種操作保存在各自的文件中（layers.py，losses.py，ops.py）。最終得到的模型（由一個或多個網路組成）應該用該模型的名稱命名（例如，yolov3.py，DCGAN.py），且引用各個模塊。

主程序、單獨的訓練和測試腳本應該只需要導入帶有模型名字的 Python 文件。

PyTorch 開發風格與技巧

我們建議將網路分解為更小的可復用的片段。一個 nn.Module 網路包含各種操作或其它構建模塊。損失函數也是包含在 nn.Module 內，因此它們可以被直接整合到網路中。

繼承 nn.Module 的類必須擁有一個「forward」方法，它實現了各個層或操作的前向傳導。

一個 nn.module 可以通過「self.net(input)」處理輸入數據。在這裡直接使用了對象的「call()」方法將輸入數據傳遞給模塊。

output = self.net(input)

PyTorch 環境下的一個簡單網路

使用下面的模式可以實現具有單個輸入和輸出的簡單網路：

class ConvBlock(nn.Module):
def __init__(self):
super(ConvBlock, self).__init__()
block = [nn.Conv2d(...)]
block += [nn.ReLU()]
block += [nn.BatchNorm2d(...)]
self.block = nn.Sequential(*block)
def forward(self, x):
return self.block(x)
class SimpleNetwork(nn.Module):
def __init__(self, num_resnet_blocks=6):
super(SimpleNetwork, self).__init__()
# here we add the individual layers
layers = [ConvBlock(...)]
for i in range(num_resnet_blocks):
layers += [ResBlock(...)]
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)

請注意以下幾點：

• 我們復用了簡單的循環構建模塊（如卷積塊 ConvBlocks），它們由相同的循環模式（卷積、
• 激活函數
• 、歸一化）組成，並裝入獨立的 nn.Module 中。
• 我們構建了一個所需要層的列表，並最終使用「nn.Sequential()」將所有層級組合到了一個模型中。我們在 list 對象前使用「*」操作來展開它。
• 在前向傳導過程中，我們直接使用輸入數據運行模型。

PyTorch 環境下的簡單殘差網路

class ResnetBlock(nn.Module):
def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
super(ResnetBlock, self).__init__()
self.conv_block = self.build_conv_block(...)
def build_conv_block(self, ...):
conv_block = []
conv_block += [nn.Conv2d(...),
norm_layer(...),
nn.ReLU()]
if use_dropout:
conv_block += [nn.Dropout(...)]
conv_block += [nn.Conv2d(...),
norm_layer(...)]
return nn.Sequential(*conv_block)
def forward(self, x):
out = x + self.conv_block(x)
return ou

在這裡，ResNet 模塊的跳躍連接直接在前向傳導過程中實現了，PyTorch 允許在前向傳導過程中進行動態操作。

PyTorch 環境下的帶多個輸出的網路

對於有多個輸出的網路（例如使用一個預訓練好的 VGG 網路構建感知損失），我們使用以下模式:

class Vgg19(torch.nn.Module):
def __init__(self, requires_grad=False):
super(Vgg19, self).__init__()
vgg_pretrained_features = models.vgg19(pretrained=True).features
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
for x in range(7):
self.slice1.add_module(str(x), vgg_pretrained_features[x])
for x in range(7, 21):
self.slice2.add_module(str(x), vgg_pretrained_features[x])
for x in range(21, 30):
self.slice3.add_module(str(x), vgg_pretrained_features[x])
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
def forward(self, x):
h_relu1 = self.slice1(x)
h_relu2 = self.slice2(h_relu1)
h_relu3 = self.slice3(h_relu2)
out = [h_relu1, h_relu2, h_relu3]
return out

請注意以下幾點：

• 我們使用由「torchvision」包提供的預訓練模型
• 我們將一個網路切分成三個模塊，每個模塊由預訓練模型中的層組成
• 我們通過設置「requires_grad = False」來固定網路權重
• 我們返回一個帶有三個模塊輸出的 list

自定義損失函數

即使 PyTorch 已經具有了大量標準損失函數，你有時也可能需要創建自己的損失函數。為了做到這一點，你需要創建一個獨立的「losses.py」文件，並且通過擴展「nn.Module」創建你的自定義損失函數：

class CustomLoss(torch.nn.Module):
def __init__(self):
super(CustomLoss,self).__init__()
def forward(self,x,y):
loss = torch.mean((x - y)**2)
return loss

訓練模型的最佳代碼結構

對於訓練的最佳代碼結構，我們需要使用以下兩種模式：

• 使用 prefetch_generator 中的 BackgroundGenerator 來載入下一個批量數據
• 使用 tqdm 監控訓練過程，並展示計算效率，這能幫助我們找到數據載入流程中的瓶頸

# import statements
import torch
import torch.nn as nn
from torch.utils import data
...
# set flags / seeds
torch.backends.cudnn.benchmark = True
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
...
# Start with main code
if __name__ == "__main__":
# argparse for additional flags for experiment
parser = argparse.ArgumentParser(description="Train a network for ...")
...
opt = parser.parse_args()
# add code for datasets (we always use train and validation/ test set)
data_transforms = transforms.Compose([
transforms.Resize((opt.img_size, opt.img_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = datasets.ImageFolder(
root=os.path.join(opt.path_to_data, "train"),
transform=data_transforms)
train_data_loader = data.DataLoader(train_dataset, ...)
test_dataset = datasets.ImageFolder(
root=os.path.join(opt.path_to_data, "test"),
transform=data_transforms)
test_data_loader = data.DataLoader(test_dataset ...)
...
# instantiate network (which has been imported from *networks.py*)
net = MyNetwork(...)
...
# create losses (criterion in pytorch)
criterion_L1 = torch.nn.L1Loss()
...
# if running on GPU and we want to use cuda move model there
use_cuda = torch.cuda.is_available()
if use_cuda:
net = net.cuda()
...
# create optimizers
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
...
# load checkpoint if needed/ wanted
start_n_iter = 0
start_epoch = 0
if opt.resume:
ckpt = load_checkpoint(opt.path_to_checkpoint) # custom method for loading last checkpoint
net.load_state_dict(ckpt["net"])
start_epoch = ckpt["epoch"]
start_n_iter = ckpt["n_iter"]
optim.load_state_dict(ckpt["optim"])
print("last checkpoint restored")
...
# if we want to run experiment on multiple GPUs we move the models there
net = torch.nn.DataParallel(net)
...
# typically we use tensorboardX to keep track of experiments
writer = SummaryWriter(...)
# now we start the main loop
n_iter = start_n_iter
for epoch in range(start_epoch, opt.epochs):
# set models to train mode
net.train()
...
# use prefetch_generator and tqdm for iterating through data
pbar = tqdm(enumerate(BackgroundGenerator(train_data_loader, ...)),
total=len(train_data_loader))
start_time = time.time()
# for loop going through dataset
for i, data in pbar:
# data preparation
img, label = data
if use_cuda:
img = img.cuda()
label = label.cuda()
...
# It"s very good practice to keep track of preparation time and computation time using tqdm to find any issues in your dataloader
prepare_time = start_time-time.time()
# forward and backward pass
optim.zero_grad()
...
loss.backward()
optim.step()
...
# udpate tensorboardX
writer.add_scalar(..., n_iter)
...
# compute computation time and *compute_efficiency*
process_time = start_time-time.time()-prepare_time
pbar.set_description("Compute efficiency: {:.2f}, epoch: {}/{}:".format(
process_time/(process_time+prepare_time), epoch, opt.epochs))
start_time = time.time()
# maybe do a test pass every x epochs
if epoch % x == x-1:
# bring models to evaluation mode
net.eval()
...
#do some tests
pbar = tqdm(enumerate(BackgroundGenerator(test_data_loader, ...)),
total=len(test_data_loader))
for i, data in pbar:
...
# save checkpoint if needed
...

PyTorch 的多 GPU 訓練

PyTorch 中有兩種使用多 GPU 進行訓練的模式。

根據我們的經驗，這兩種模式都是有效的。然而，第一種方法得到的結果更好、需要的代碼更少。由於第二種方法中的 GPU 間的通信更少，似乎具有輕微的性能優勢。

對每個網路輸入的 batch 進行切分

最常見的一種做法是直接將所有網路的輸入切分為不同的批量數據，並分配給各個 GPU。

這樣一來，在 1 個 GPU 上運行批量大小為 64 的模型，在 2 個 GPU 上運行時，每個 batch 的大小就變成了 32。這個過程可以使用「nn.DataParallel(model)」包裝器自動完成。

將所有網路打包到一個超級網路中，並對輸入 batch 進行切分

這種模式不太常用。下面的代碼倉庫向大家展示了 Nvidia 實現的 pix2pixHD，它有這種方法的實現。

地址：https://github.com/NVIDIA/pix2pixHD

PyTorch 中該做和不該做的

在「nn.Module」的「forward」方法中避免使用 Numpy 代碼

Numpy 是在 CPU 上運行的，它比 torch 的代碼運行得要慢一些。由於 torch 的開發思路與 numpy 相似，所以大多數

Numpy

中的函數已經在 PyTorch 中得到了支持。

將「DataLoader」從主程序的代碼中分離

載入數據的工作流程應該獨立於你的主訓練程序代碼。PyTorch 使用「background」進程更加高效地載入數據，而不會干擾到主訓練進程。

不要在每一步中都記錄結果

通常而言，我們要訓練我們的模型好幾千步。因此，為了減小計算開銷，每隔 n 步對損失和其它的計算結果進行記錄就足夠了。尤其是，在訓練過程中將中間結果保存成圖像，這種開銷是非常大的。

使用命令行參數

使用命令行參數設置代碼執行時使用的參數（batch 的大小、學習率等）非常方便。一個簡單的實驗參數跟蹤方法，即直接把從「parse_args」接收到的字典（dict 數據）列印出來：

# saves arguments to config.txt file
opt = parser.parse_args()with open("config.txt", "w") as f:
f.write(opt.__str__())

如果可能的話，請使用「Use .detach()」從計算圖中釋放張量

為了實現自動微分，PyTorch 會跟蹤所有涉及張量的操作。請使用「.detach()」來防止記錄不必要的操作。

使用「.item()」列印出標量張量

你可以直接列印變數。然而，我們建議你使用「variable.detach()」或「variable.item()」。在早期版本的 PyTorch（< 0.4）中，你必須使用「.data」訪問變數中的張量值。

使用「call」方法代替「nn.Module」中的「forward」方法

這兩種方式並不完全相同，正如下面的 GitHub 問題單所指出的：https://github.com/IgorSusmelj/pytorch-styleguide/issues/3

output = self.net.forward(input)
# they are not equal!
output = self.net(input)

原文鏈接：https://github.com/IgorSusmelj/pytorch-styleguide

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！