PyTorch最佳實踐,怎樣才能寫出一手風格優美的程式碼

機器之心發表於2019-04-29

雖然這是一個非官方的 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。

在這裡,我們會給出一個最常用命名規範小結:

PyTorch最佳實踐,怎樣才能寫出一手風格優美的程式碼

整合開發環境

一般來說,我們建議使用 visual studio 或 PyCharm 這樣的整合開發環境。而 VS Code 在相對輕量級的編輯器中提供語法高亮和自動補全功能,PyCharm 則擁有許多用於處理遠端叢集任務的高階特性。

Jupyter Notebooks VS Python 指令碼

一般來說,我們建議使用 Jupyter Notebook 進行初步的探索,或嘗試新的模型和程式碼。如果你想在更大的資料集上訓練該模型,就應該使用 Python 指令碼,因為在更大的資料集上,復現性更加重要。

我們推薦你採取下面的工作流程:

  • 在開始的階段,使用 Jupyter Notebook

  • 對資料和模型進行探索

  • 在 notebook 的單元中構建你的類/方法

  • 將程式碼移植到 Python 指令碼中

  • 在伺服器上訓練/部署

PyTorch最佳實踐,怎樣才能寫出一手風格優美的程式碼

開發常備庫

常用的程式庫有:

PyTorch最佳實踐,怎樣才能寫出一手風格優美的程式碼

檔案組織

不要將所有的層和模型放在同一個檔案中。最好的做法是將最終的網路分離到獨立的檔案(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)

原文連結:https://github.com/IgorSusmelj/pytorch-styleguide

相關文章