雖然這是一個非官方的 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