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
            ...