PyTorch最佳实践和代码编写风格指南

这不是PyTorch的官方风格指南。本文总结了使用PyTorch框架进行深度学习的一年多的经验中的最佳实践。请注意，我们分享的经验大多来自研究和创业的视角。本文档分为三个部分。一是Python最佳实践的快速总结，二是一些使用PyTorch的tips和建议，三是分享了一些使用其他框架的见解和经验，这些框架帮助我们总体上改进了我们的工作流程。

建议使用Python 3.6+

Python风格指南之梗概

在Python上，遵循Google的代码风格，具体参考下面这个链接
Please refer to the well-documented style guide on python code provided by Google.

一些总结如下：

命名规范

类型Type	规范Convention	栗子Example
Packages & Modules 包和模块	lower_with_under 小写+下划线	from prefetch_generator import BackgroundGenerator
Classes 类	CapWords 首字母大写	class DataLoader
Constants 常量	CAPS_WITH_UNDER 大写+下划线	BATCH_SIZE=16
Instances 实例	lower_with_under 小写+下划线	dataset = Dataset
Methods & Functions 方法函数	lower_with_under() 小写+下划线	def visualize_tensor()
Variables 变量	lower_with_under 小写+下划线	background_color='Blue’

IDEs

代码编辑器

建议使用 visual studio code 或 PyCharm. Whereas VS Code provides syntax highlighting and autocompletion in a relatively lightweight editor PyCharm has lots of advanced features for working with remote clusters.
VS Code由于其快速增长的扩展生态系统而变得非常强大。

设置VS Code与远程机器一起工作

确保你已经安装了以下扩展:

• Python (检查、自动补全、语法高亮显示、代码格式化 linting, autocompletion, syntax highlighting, code formatting)
• Remote - SSH (to work with remote machines)

1. 教程链接: https://code.visualstudio.com/docs/remote/remote-overview

设置PyCharm与远程机器一起工作

1. 登录到您的远程机器(AWS，谷歌等)
2. 创建新文件夹和新的虚拟环境
3. 在Pycharm(专业版)项目设置中设置一个远程解释器
4. 配置远程python解释器(AWS上venv的路径，谷歌等)
5. 配置代码从本地机器到远程机器的映射

如果设置正确，你可以做以下事情:

• 在你的本地电脑(笔记本电脑，桌面电脑)上进行coding(离线，在线)
• 同步本地代码与远程机器
• 额外的包将自动安装在远程机器上
• 你不需要任何数据集在你的本地机器上
• 在远程机器上运行代码并调试，就像在本地机器上运行代码一样

Jupyter Notebook vs Python 脚本

总的来说，建议使用jupyter Notebook进行初步的探索/尝试新的模型和代码。当你想要在更大的数据集上训练模型时，就应该使用Python脚本，在更大的数据集上，重现性也更重要。

建议的工作流

1. 从Jupyter Notebook开始
2. 探索数据和模型
3. 在笔记本的单元格内构建类和方法
4. 将代码移动到python脚本中
5. 训练和部署在服务器上

Jupyter Notebook	Python 脚本
+ Exploration 探索	+ Running longer jobs without interruption （不可打断）
+ Debugging 调试	+ Easy to track changes with git（容易跟踪）
- Can become a huge file 不能使用大文件（不建议）	- Debugging mostly means rerunning the whole script（调试不便）
- Can be interrupted (don’t use for long training) （可以打断）
- Prone to errors and become a mess （很乱）

Libraries包

建议使用的包

Name包名	Description描述	Used for适用于
torch	使用神经网络的基本框架	creating tensors, networks and training them using backprop
torchvision	torch数据处理工具	data preprocessing, augmentation, postprocessing
Pillow (PIL)	Python影像库	Loading images and storing them
Numpy	Python科学计算包	Data preprocessing & postprocessing
prefetch_generator	后台处理库	Loading next batch in background during computation
tqdm	进度条	Progress during training of each epoch
torchsummary	PyTorch的summary	Displays network, it’s parameters and sizes at each layer
tensorboardX	Tensorboard without tensorflow	Logging experiments and showing them in tensorboard

文件组织

不要把所有的层和模型放在同一个文件中。最佳实践是将最终的网络分离到一个单独的文件(networks.py)，并将层、损失和ops保存在各自的文件中(layers.py、loss .py、ops.py)。完成的模型(由一个或多个网络组成)应该在一个带有其名称的文件中引用(例如yolov3.py， DCGAN.py)

主例程，分别是训练脚本和测试脚本，应该只从具有模型名称的文件中导入。

在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__()
        self.block = nn.Sequential(
            nn.Conv2d(...), 
            nn.ReLU(), 
            nn.BatchNorm2d(...)
        )  
    
    def forward(self, x):
        return self.block(x)

class SimpleNetwork(nn.Module):
    def __init__(self, num_resnet_blocks=6):
        super(SimpleNetwork, self).__init__()
        # 在这里我们添加了单独的层
        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.Module
• 我们建立一个所需层的list，并最终使用**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

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来冻结网络
• 返回一个包含切片的三个输出的列表

自定义Loss

尽管PyTorch有了很多搭建好的标准的Loss损失函数，有时候你还是可能会用到你自己的损失。为此，创建一个单独的文件 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

推荐用于训练模型的代码结构

完整示例请看cifar10-example

Note that we used the following patterns:

• We use BackgroundGenerator from prefetch_generator to load next batches in background see this issue for more information
• We use tqdm to monitor training progress and show the compute efficiency. This helps us find bottlenecks in our data loading pipeline.

注意，我们使用了以下模式:

• 我们使用BackgroundGenerator在后台加载下一批。参见此了解更多
• 我们使用tqdm来监控训练进度并显示计算效率。

# import statements
import torch
import torch.nn as nn
from torch.utils import data
...

# set flags / seeds 设置flags或种子
torch.backends.cudnn.benchmark = True
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
...

# 主函数入口
if __name__ == '__main__':
    # argparse： 用于实验的附加flags
    parser = argparse.ArgumentParser(description="Train a network for ...")
    ...
    opt = parser.parse_args() 
    
    # 数据集相关的代码
    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 ...)
    ...
    
    # 实例化网络
    net = MyNetwork(...)
    ...
    
    # 损失函数
    criterion_L1 = torch.nn.L1Loss()
    ...
    
    # 是否使用cuda，自行进行选择
    use_cuda = torch.cuda.is_available()
    if use_cuda:
        net = net.cuda()
        ...
    
    # 创建优化器
    optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
    ...
    
    # 是否加载预训练权重
    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")
        ...
        
    # 多个GPU并行训练
    net = torch.nn.DataParallel(net)
    ...
    
    # 通常我们使用tensorboardX来跟踪实验
    writer = SummaryWriter(...)
    
    # 开始主循环
    n_iter = start_n_iter
    for epoch in range(start_epoch, opt.epochs):
        # 设置为训练模式.train()
        net.train()
        ...
        
        # 使用prefetch_generator和TQDM来遍历数据
        pbar = tqdm(enumerate(BackgroundGenerator(train_data_loader, ...)),
                    total=len(train_data_loader))
        start_time = time.time()
        
        # 循环遍历数据集
        for i, data in pbar:
            # 数据准备
            img, label = data
            if use_cuda:
                img = img.cuda()
                label = label.cuda()
            ...
            
            # 使用tqdm来跟踪准备时间和计算时间，以发现数据加载程序中的任何问题，这是一个很好的实践
            prepare_time = start_time-time.time()
            
            # 前向和后向传播
            optim.zero_grad()
            ...
            loss.backward()
            optim.step()
            ...
            
            # 更新 tensorboardX
            writer.add_scalar(..., n_iter)
            ...
            
            # 计算计算时间和计算效率
            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()
            
        # 每x个epochs进行一次测试
        if epoch % x == x-1:
            # evel模式
            net.eval()
            ...
            # 测试
            pbar = tqdm(enumerate(BackgroundGenerator(test_data_loader, ...)),
                    total=len(test_data_loader)) 
            for i, data in pbar:
                ...
                
            # save checkpoint if needed
            ...

PyTorch在多卡进行训练

在PyTorch中有两种使用多个gpu进行训练的方法。
从我们的经验来看，这两种模式都是有效的。然而，第一个结果是更好和更少的代码。第二种方法似乎有一点性能优势，因为gpu之间的通信更少。我在PyTorch官方论坛问了一个关于这两种方法的问题 https://discuss.pytorch.org/t/how-to-best-use-dataparallel-with-multiple-models/39289)

分解每个网络的批量输入

最常见的一种方法是简单地将所有网络的批次分割到单个gpu上。

因此，一个运行在批大小为64的1个GPU上的模型将运行在2个批大小为32的GPU上。这可以通过**nn. dataparelles (model)**自动完成。

将所有网络打包在一个super网络中，并拆分输入批处理

这种模式不太常用。实现这种方法的链接：pix2pixHD

需要做和不能做的事情

避免nn.Module中使用Numpy代码

Numpy在CPU上运行，比torch代码慢。由于torch在开发过程中与numpy类似，所以PyTorch已经支持了大多数numpy函数。

将DataLoader从主代码中分离出来

数据加载pipeline应该独立于你的主函数中的训练代码。PyTorch在后台会更有效地加载数据，而且不会干扰主函数中的训练过程。

不要每一步都记录结果

通常，我们训练我们的模型有数千个steps。因此，每第n步记录丢失和其他结果就足以减少开销。特别是，在训练过程中，将中间结果保存为图像的成本很高。

使用命令行参数

在代码执行期间使用命令行参数设置参数非常方便(批大小，学习率等)。跟踪实验参数的一个简单方法是打印从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版本< 0.4中，必须使用**.data**访问一个张量变量。

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

output = self.net.forward(input)
# 上下两个方法是不一样的
output = self.net(input)

FAQ

1. 如何保持我的实验重现性?

我们建议在代码的开头设置以下种子:

np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)

2. 如何进一步提高训练和推理速度?

在Nvidia GPUs 你可以使用如下代码. 这将允许cuda后端在第一次执行时优化你的tensor图。但是，请注意，如果您更改网络输入/输出张量大小，则每次发生更改时，图都会被优化。这可能会导致非常慢的运行时间和内存不足错误。只有当你的输入和输出总是相同的形状时才设置这个标志。通常情况下，这将有大约20%的改善。

torch.backends.cudnn.benchmark = True

3. 使用tqdm + prefetch_generator模式对计算效率有什么好处?

取决于使用的机器、预处理pipeline和网络规模。在1080Ti GPU的SSD上运行，我们看到计算效率接近1.0，这是一个理想的场景。如果使用浅层网络或慢速的硬盘，这个数字可能会下降到0.1-0.2左右，这取决于您的设置。

4. 我怎么能有一个batsize > 1，即使我没有足够的内存?

在PyTorch中，我们可以很容易地实现虚拟批处理大小。我们只是阻止优化器更新参数，并为batch_size周期性累计计算梯度。

...
# in the main loop
out = net(input)
loss = criterion(out, label)
# 前向计算，但是batchsize批次间，不进行反向梯度更新，这样就可以达到batchsize批次的目的
loss.backward() 
total_loss += loss.item() / batch_size
if n_iter % batch_size == batch_size-1:
    # 这里我们使用虚拟批大小执行优化步骤
    optim.step()
    optim.zero_grad()
    print('Total loss: ', total_loss)
    total_loss = 0.0
...

5. 如何在训练时调整学习速度?

我们可以直接使用实例化的优化器来访问学习率，如下所示:

...
for param_group in optim.param_groups:
    old_lr = param_group['lr']
    new_lr = old_lr * 0.1
    param_group['lr'] = new_lr
    print('Updated lr from {} to {}'.format(old_lr, new_lr))
...

6. 如何在训练中使用预训练的模型计算损失(不使用no_grad())

如果你想使用预先训练的模型，如VGG来计算损失，但不训练它，你可以使用以下模式:

...
# 实例化模型
pretrained_VGG = VGG19(...)

# 禁用梯度(防止训练)
for p in pretrained_VGG.parameters():  # reset requires_grad
    p.requires_grad = False
...
# 不使用no_grad()，但可以只运行模型
# no gradients will be computed for the VGG model
out_real = pretrained_VGG(input_a)
out_fake = pretrained_VGG(input_b)
loss = any_criterion(out_real, out_fake)
...

7. 为什么在PyTorch中使用*.train()和.eval()* ?

Those methods are used to set layers such as BatchNorm2d or Dropout2d from training to inference mode. Every module which inherits from nn.Module has an attribute called isTraining. .eval() and .train() just simply sets this attribute to True/ False. For more information of how this method is implemented please have a look at the module code in PyTorch

这些方法作用在某些层，如BatchNorm2d或Dropout2d从训练到推理模式。从nn继承的每个模块。模块有一个名为isTraining的属性。**.eval()和.train()**只是简单地将该属性设置为True/ False。有关如何实现此方法的更多信息，请参阅PyTorch中的模块代码

8. 为何我的模型在推理过程中使用了大量的内存/如何在PyTorch中正确运行模型进行推理?

确保在代码执行期间没有计算和存储梯度信息。你可以简单地使用以下模式来确保:

with torch.no_grad():
    # run model here
    out_tensor = net(in_tensor)

9. 如何微调一个预训练的模型?

在PyTorch中，你可以冻结部分图层。这将防止它们在优化步骤中被更新。

# you can freeze whole modules using
for p in pretrained_VGG.parameters():  # reset requires_grad
    p.requires_grad = False

10. 什么时候使用 Variable(…)?

自从PyTorch 0.4 Variable和Tensor被合并。我们不再需要显式地创建Variable对象。

11. 在c++上使用PyTorch比使用Python更快吗?

C++ 版本快10%左右

12. TorchScript / JIT可以加速我的代码吗?

待更新

13. PyTorch 使用 cudnn.benchmark=True 会加快速度吗?

根据我们的经验，可以获得大约20%的加速。但第一次运行模型需要相当长的时间构建优化的图。在某些情况下(循环在前向传递，没有固定的输入形状，if/else在前向，等等)这个标志可能导致out of memory或其他错误

14. 如何使用多卡训练?

待更新

15. .detach() 在 PyTorch如何进行工作的?

把一个张量从计算图中解放出来。一个插图在这 here

参考链接：

https://github.com/IgorSusmelj/pytorch-styleguide