PyTorch最佳实践，怎样才能写出一手风格优美的代码

时间 2019-12-18

标签 pytorch 最佳实践怎样才能写出一手风格优美代码繁體版

原文原文链接

【摘要】 PyTorch是最优秀的深度学习框架之一，它简单优雅，很是适合入门。本文将介绍PyTorch的最佳实践和代码风格都是怎样的。python

虽然这是一个非官方的 PyTorch 指南，但本文总结了一年多使用 PyTorch 框架的经验，尤为是用它开发深度学习相关工做的最优解决方案。请注意，咱们分享的经验大可能是从研究和实践角度出发的。git

这是一个开发的项目，欢迎其它读者改进该文档：https://github.com/IgorSusmelj/pytorch-styleguide。github

本文档主要由三个部分构成：首先，本文会简要清点 Python 中的最好装备。接着，本文会介绍一些使用 PyTorch 的技巧和建议。最后，咱们分享了一些使用其它框架的看法和经验，这些框架一般帮助咱们改进工做流。编程

清点 Python 装备服务器

建议使用 Python 3.6 以上版本网络

根据咱们的经验，咱们推荐使用 Python 3.6 以上的版本，由于它们具备如下特性，这些特性可使咱们很容易写出简洁的代码：框架

自 Python 3.6 之后支持「typing」模块dom
自 Python 3.6 之后支持格式化字符串（f string）编辑器

Python 风格指南ide

咱们试图遵循 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)

来源：原文连接