PyTorch最佳實踐，怎樣才能寫出一手風格優美的代碼

【摘要】 PyTorch是最優秀的深度學習框架之一，它簡單優雅，很是適合入門。本文將介紹PyTorch的最佳實踐和代碼風格都是怎樣的。

雖然這是一個非官方的 PyTorch 指南，但本文總結了一年多使用 PyTorch 框架的經驗，尤爲是用它開發深度學習相關工做的最優解決方案。請注意，咱們分享的經驗大可能是從研究和實踐角度出發的。

這是一個開發的項目，歡迎其它讀者改進該文檔：https://github.com/IgorSusmelj/pytorch-styleguide。

本文檔主要由三個部分構成：首先，本文會簡要清點 Python 中的最好裝備。接着，本文會介紹一些使用 PyTorch 的技巧和建議。最後，咱們分享了一些使用其它框架的看法和經驗，這些框架一般幫助咱們改進工做流。

清點 Python 裝備

建議使用 Python 3.6 以上版本

根據咱們的經驗，咱們推薦使用 Python 3.6 以上的版本，由於它們具備如下特性，這些特性可使咱們很容易寫出簡潔的代碼：

• 自 Python 3.6 之後支持「typing」模塊

• 自 Python 3.6 之後支持格式化字符串（f string）

Python 風格指南

咱們試圖遵循 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)

來源：原文連接