PyTorch代碼規範最佳實踐和樣式指南

根據經驗,作者建議使用 Python 3.6+,因為以下功能有助於寫出乾淨簡單的代碼:

  • 支持 Python 3.6 以後的輸入。
  • 自 Python 3.6 起支持 f 字符串

Python Styleguide 概述

作者嘗試按照 Google Styleguide for Python 進行操作,這裡是 Google 提供的 python 代碼詳細樣式指南。

常見的命名約定:


PyTorch代碼規範最佳實踐和樣式指南


Jupyter Notebook與Python腳本

一般來說,建議使用 Jupyternotebook 進行初步探索和使用新的模型和代碼。如果你想在更大的數據集上訓練模型,就應該使用 Python 腳本。在這裡,複用性更為重要。

推薦使用的工作流程是:

  1. 從Jupyter筆記本開始
  2. 探索數據和模型
  3. 在 notebook 的單元格中構建類/方法
  4. 將代碼移動到python腳本中
  5. 在服務器上訓練/部署

注意,不要將所有層和模型放在同一個文件中。最佳做法是將最終網絡分離為單獨的文件(networks.py),並將層、損耗和 ops 保存在各自的文件(layers.py、losses.py、ops.py)中。完成的模型(由一個或多個網絡組成)應在一個文件中引用,文件名為 yolov3.py、dcgan.py 這樣。

在PyTorch中構建神經網絡

我們建議將網絡拆分為更小的可重用部分。網絡由操作或其它網絡模塊組成。損失函數也是神經網絡的模塊,因此可以直接集成到網絡中。

繼承自 nn.module 的類必須有一個 forward 方法來實現各個層或操作的 forward 傳遞。

使用 self.net(input),可以在輸入數據上使用 nn.module。這隻需使用對象的 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)

需要注意的是:

  • 重用簡單的、循環的構建塊,例如 ConvBlock,它由相同的循環模式(卷積、激活、歸一化)組成,並將它們放入單獨的nn.模塊中。
  • 作者構建了一個所需層的列表,最後使用 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 out

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

PyTorch中具有多個輸出的網絡

對於需要多個輸出的網絡,例如使用預訓練的 VGG 網絡構建感知損失,我們使用以下模式:

class Vgg19(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 來凍結網絡。
  • 我們返回一個包含三個輸出部分的列表。

自定義損失

雖然 PyTorch 已經有很多標準的損失函數,但有時也可能需要創建自己的損失函數。為此,請創建單獨的文件 losses.py 並擴展 nn.module 類以創建自定義的損失函數:

class CustomLoss(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 在後臺加載 batch。有關詳細信息,請參閱這裡。

我們使用 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 之間的通信較少,第二種方法似乎具有輕微的性能優勢。

分割每個網絡的批輸入

最常見的方法是簡單地將所有網絡的批劃分為單個 GPU。

因此,在批大小為 64 的 1 個 GPU 上運行的模型將在批大小為 32 的 2 個 GPU 上運行。這可以通過使用 nn.dataparallel(model)自動包裝模型來完成。

將所有網絡打包到超級網絡中並拆分輸入批

這種模式不太常用。Nvidia 的 pix2pixhd 實現中顯示了實現此方法的存儲庫。

什麼該做什麼不該做

避免在 nn.Module 的 forward 方法中使用 numpy 代碼

numpy 代碼在 CPU 上運行的速度比 torch 代碼慢。由於 torch 的開發理念和 numpy 類似,所以 pytorch 支持大多數 numpy 函數。

將數據加載器與主代碼分離

數據加載管道應該獨立於你的主要訓練代碼。PyTorch 使後臺工作人員可以更高效地加載數據,但不會干擾主要的訓練過程。

不要每個步驟都輸出結果日誌

通常,我們對模型進行數千步的訓練。因此,不要在每一步記錄結果就足以減少開銷。尤其是,在訓練過程中將中間結果保存為圖像成本高昂。

使用命令行參數

在代碼執行期間使用命令行參數設置參數(批大小、學習速率等)非常方便。跟蹤實驗參數的一個簡單方法是隻打印從 parse_args 接收到的字典:

...

# saves arguments to config.txt file

opt = parser.parse_args()

with open("config.txt", "w") as f:

f.write(opt.__str__())...

如果可能,使用 .detach()從圖表中釋放張量

pytorch跟蹤所有涉及張量的自動微分操作。使用 .detach()防止記錄不必要的操作。

使用 .item()打印標量張量

你可以直接打印變量,但是建議使用 variable.detach()或 variable.item()。在早期的 pytorch 版本中,必須使用 .data 來訪問變量的張量。

在 nn.Module 上使用 call 方法而不是 forward

這兩種方法不完全相同,下面的例子就可以看出這一點:

output = self.net.forward(input)

# they are not equal!

output = self.net(input)

如果你是想深入學習人工智能那你可以私聊我呀!帶你們玩這些項目。

PyTorch代碼規範最佳實踐和樣式指南


分享到:


相關文章: