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

本文作者： skura

2019-04-16 13:02

導語：PyTorch的非官方風格指南和最佳實踐摘要

雷鋒網 AI 科技評論按，本文不是 Python 的官方風格指南。本文總結了使用 PyTorch 框架進行深入學習的一年多經驗中的最佳實踐。本文分享的知識主要是以研究的角度來看的，它來源于一個開元的 github 項目。

根據經驗，作者建議使用 Python 3.6+，因為以下功能有助于寫出干凈簡單的代碼：

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

Python Styleguide 概述

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

常見的命名約定：

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

Jupyter Notebook與Python腳本

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

推薦使用的工作流程是：

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

注意，不要將所有層和模型放在同一個文件中。最佳做法是將最終網絡分離為單獨的文件（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
            ...