詳解Pytorch如何利用yaml定義折積網路

2022-10-22 14:00:52

大多數折積神經網路都是直接通過寫一個Model類來定義的，這樣寫的程式碼其實是比較好懂的，特別是在魔改網路的時候也很方便。然後也有一些會通過cfg組態檔進行模型的定義。在yolov5中可以看到是通過yaml檔案進行網路的定義【個人感覺通過組態檔魔改網路有些不方便，當然每個人習慣不同】，可能很多人也用過，如果自己去寫一個yaml檔案，自己能不能定義出來呢？很多人不知道是如何具體通過yaml檔案將裡面的引數傳入自己定義的網路中，這也就給自己修改網路帶來了不便。這篇文章將仿照yolov5的方式，利用yaml定義一個自己的網路。

定義折積塊

我們可以先定義一個折積塊CBL，C指折積Conv，B指BN層，L為啟用函數，這裡我用ReLu.

class BaseConv(nn.Module):
    def __init__(self, in_channels, out_channels, k=1, s=1, p=None):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.conv = nn.Conv2d(in_channels, out_channels, k, s, autopad(k, p))
        self.bn = nn.BatchNorm2d(out_channels)
        self.act_fn = nn.ReLU(inplace=True)
 
    def forward(self, x):
        return self.act_fn(self.bn(self.conv(x)))

折積中的autopad是自動補充pad，程式碼如下：

def autopad(k, p=None):
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
    return p

定義一個Bottleneck

可以仿照yolov5定義一個Bottleneck，參考了殘差塊的思想。

class Bottleneck(nn.Module):
    def __init__(self, in_channels, out_channels, shortcut=True):
        super(Bottleneck, self).__init__()
        self.conv1 = BaseConv(in_channels, out_channels, k=1, s=1)
        self.conv2 = BaseConv(out_channels, out_channels, k=3, s=1)
        self.add = shortcut and in_channels == out_channels
 
    def forward(self, x):
        """
        x-->conv1-->conv2-->add
          |_________________|
        """
        return x + self.conv2(self.conv1(x)) if self.add else self.conv2(self.conv1(x))

攥寫yaml組態檔

然後我們來寫一下yaml組態檔，網路不要很複雜，就由兩個折積和兩個Bottleneck組成就行。同理，仿v5的方法，我們的網路中的backone也是個列表，每行為一個折積層，每列有4個引數，分別代表from(指該層的輸入通道數為上一層的輸出通道數，所以是-1)，number【yaml中的1，1，2指該層的深度，或者說是重複幾次】，Module_nams【該層的名字】，args【網路引數，包含輸出通道數，k，s，p等設定】

# define own model
backbone:
  [[-1, 1, BaseConv, [32, 3, 1]],  # out_channles=32, k=3, s=1
   [-1, 1, BaseConv, [64, 1, 1]],
   [-1, 2, Bottleneck, [64]]
  ]

我們現在用yaml工具來開啟我們的組態檔，看看都有什麼內容

    import yaml
    # 獲得yaml檔案名字
    yaml_file = Path('Model.yaml').name
    with open(yaml_file,errors='ignore') as f:
        yaml_ = yaml.safe_load(f)
    print(yaml_)

輸出：

{'backbone': [[-1, 1, 'BaseConv', [32, 3, 1]], [-1, 1, 'BaseConv', [64, 1, 1]], [-1, 2, 'Bottleneck', [64]]]}

然後我們可以定義下自己Model類，也就是定義自己的網路。可以看到與前面讀取yaml檔案相比，多了一行 ch = self.yaml["ch"] = self.yaml["ch"] = 3 這個是在原yaml內容中加入一個key和valuse,3指的3通道，因為我們的影象是3通道。parse_model是下面要說的傳參過程。

class Model(nn.Module):
    def __init__(self, cfg='./Model.yaml', ch=3, ):
        super().__init__()
        self.yaml = cfg
        import yaml
        yaml_file = Path(cfg).name
        with open(yaml_file, errors='ignore')as f:
            self.yaml = yaml.safe_load(f)
 
        ch = self.yaml["ch"] = self.yaml["ch"] = 3
        self.backbone = parse_model(deepcopy(self.yaml), ch=[ch])
 
    def forward(self, x):
        output = self.backbone(x)
        return output

傳入引數

這一步也是最關鍵的一步，我們需要定義傳參的函數，將yaml中的折積引數傳入我們定義的網路中，這裡會用的一個非常非常重要的函數eval(),後面也會介紹到這個函數的用法。

這裡先附上完整程式碼：

def parse_model(yaml_cfg, ch):
    """
    :param yaml_cfg: yaml file
    :param ch: init in_channels default is 3
    :return: model
    """
 
    layer, out_channels = [], ch[-1]
    for i, (f, number, Module_name, args) in enumerate(yaml_cfg['backbone']):
        """
        f:上一層輸出通道
        number:該模組有幾層，就是該模組要重複幾次
        Mdule_name：折積層名字
        args:引數，包含輸出通道數，k,s,p等
        """
        # 通過eval，將str型別轉自己定義的BaseConv
        m = eval(Module_name) if isinstance(Module_name, str) else Module_name
        for j, a in enumerate(args):
            # 通過eval，將str轉int,獲得輸出通道數
            args[j] = eval(a) if isinstance(a, str) else a
        # 更新通道
        # args[0]是輸出通道
        if m in [BaseConv, Bottleneck]:
            in_channels, out_channels = ch[f], args[0]
            args = [in_channels, out_channels, *args[1:]]  # args=[in_channels, out_channels, k, s, p]
 
        # 將引數傳入模型
        model_ = nn.Sequential(*[m(*args) for _ in range(number)]) if number > 1 else m(*args)
        # 更新通道列表，每次獲取輸出通道
        ch.append(out_channels)
        layer.append(model_)
    return nn.Sequential(*layer)

下面開始分析程式碼。

這行程式碼是通過列表用來存放每層內容以及輸出通道數。

# 這行程式碼是通過列表用來存放每層內容以及輸出通道數
layer, out_channels = [], ch[-1]

然後進入我們的for迴圈，在每一次迴圈中可以獲得我們yaml檔案中的每一層網路：f是上一層網路的輸出通道【用來作為本層的輸入通道】，number【網路深度，也就是該層重複幾次而已】，Module_name是該層的名字，args是該層的一些引數。

for i, (f, number, Module_name, args) in enumerate(yaml_cfg['backbone']):

接下來會碰到一個很重要的函數eval()。下行的程式碼首先需要判斷一下我們的Module_name型別是不是字串型別，也就是判斷一下yaml中“BaseConv”是不是字串型別，如果是，則用eval進行對應型別的轉化，轉成我們的BaseConv型別。

m = eval(Module_name) if isinstance(Module_name, str) else Module_name

這裡我將對eval函數在深入點，如果知道這個函數用法的，就可以略去這部分。

我們先舉個例子，比如我現在有個變數a="123"，這個a的型別是什麼呢？他是一個str型別，不是int型別。現在我們用eval函數轉一下，看看會變成什麼樣子。

>>> b = eval(a) if isinstance(a,str) else a
>>> b
123
>>> type(b)
<class 'int'>

我們可以看到，經過eval函數以後，會自動識別並轉為int型別。那麼我繼續舉例子，如果現在a="BaseConv",經過eval以後會變成什麼？可以看到，這裡報錯了！這是為什麼？這是因為我們沒有匯入BaseConv這個類，所以eval函數並不知道我們希望轉為什麼型別。所以我們需要用import匯入BaseConv這個類才可以。

>>> a="BaseConv"
>>> b = eval(a) if isinstance(a,str) else a
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<string>", line 1, in <module>
NameError: name 'BaseConv' is not defined

當我們匯入BaseConv以後，在經過eval就可以獲得：

<class 'models.BaseConv'>

接下來是獲得args中的網路引數，也是通過eval進行轉化

        for j, a in enumerate(args):
            # 通過eval，將str轉int,獲得輸出通道數
            args[j] = eval(a) if isinstance(a, str) else a

獲取通道數，並在每次迴圈中對通道進行更新：可以仔細看一下ch[f]指的上一層輸出通道，剛開始預設為[3]，那麼ch[-1]=3，我們yaml中第一層的BaseConv args[0]為32，表示輸出32通道。因此在第一次迴圈中有in_channels = 3,out_channels=32。args也要更新，*args前面的"*"並不是指標的意思，也不是乘的意思，而是解壓操作，因此我們第一次迴圈中得到的args=[3,32,3,1]。

# 更新通道
# args[0]是輸出通道
if m in [BaseConv, Bottleneck]:
    in_channels, out_channels = ch[f], args[0]
    args = [in_channels, out_channels, *args[1:]]  # args=[in_channels, out_channels, k, s, p]

將引數傳入模型

這裡用for _ in range(number)來判斷網路的深度【或者說該模組重複幾次】，這裡的m就是前面經過eval轉化的 <class 'models.BaseConv'>。通過*args解壓操作將args列表中的內容放入m中，再通過*解壓操作放入nn.Sequential。

model_ = nn.Sequential(*[m(*args) for _ in range(number)]) if number > 1 else m(*args)

這樣就可以獲得我們第一次迴圈BaseConv了。後面的迴圈也是同樣的反覆操作而已。

BaseConv(
  (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (act_fn): ReLU(inplace=True)
)

然後是更新通道列表和layer列表，為的是獲取每次迴圈的輸出通道，沒有這一步，再下一次迴圈的時候將不能正確得到通道數。

# 更新通道列表，每次獲取輸出通道
ch.append(out_channels)
layer.append(model_)

然後我們就可以對模型呼叫進行範例化了，可以列印下模型：

Model(
  (backbone): Sequential(
    (0): BaseConv(
      (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (act_fn): ReLU(inplace=True)
    )
    (1): BaseConv(
      (conv): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (act_fn): ReLU(inplace=True)
    )
    (2): Sequential(
      (0): Bottleneck(
        (conv1): BaseConv(
          (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (act_fn): ReLU(inplace=True)
        )
        (conv2): BaseConv(
          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (act_fn): ReLU(inplace=True)
        )
      )
      (1): Bottleneck(
        (conv1): BaseConv(
          (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (act_fn): ReLU(inplace=True)
        )
        (conv2): BaseConv(
          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (act_fn): ReLU(inplace=True)
        )
      )
    )
  )
)

同時我們也可以對模型每層視覺化看一下。可以看到和我們定義的模型是一樣的。

完整的程式碼

from copy import deepcopy
 
from models import BaseConv, Bottleneck
import torch.nn as nn
import os
 
path = os.getcwd()
from pathlib import Path
import torch
 
 
def parse_model(yaml_cfg, ch):
    """
    :param yaml_cfg: yaml file
    :param ch: init in_channels default is 3
    :return: model
    """
 
    layer, out_channels = [], ch[-1]
    for i, (f, number, Module_name, args) in enumerate(yaml_cfg['backbone']):
        """
        f:上一層輸出通道
        number:該模組有幾層，就是該模組要重複幾次
        Mdule_name：折積層名字
        args:引數，包含輸出通道數，k,s,p等
        """
        # 通過eval，將str型別轉自己定義的BaseConv
        m = eval(Module_name) if isinstance(Module_name, str) else Module_name
        for j, a in enumerate(args):
            # 通過eval，將str轉int,獲得輸出通道數
            args[j] = eval(a) if isinstance(a, str) else a
        # 更新通道
        # args[0]是輸出通道
        if m in [BaseConv, Bottleneck]:
            in_channels, out_channels = ch[f], args[0]
            args = [in_channels, out_channels, *args[1:]]  # args=[in_channels, out_channels, k, s, p]
 
        # 將引數傳入模型
        model_ = nn.Sequential(*[m(*args) for _ in range(number)]) if number > 1 else m(*args)
        # 更新通道列表，每次獲取輸出通道
        ch.append(out_channels)
        layer.append(model_)
    return nn.Sequential(*layer)
 
 
class Model(nn.Module):
    def __init__(self, cfg='./Model.yaml', ch=3, ):
        super().__init__()
        self.yaml = cfg
        import yaml
        yaml_file = Path(cfg).name
        with open(yaml_file, errors='ignore')as f:
            self.yaml = yaml.safe_load(f)
 
        ch = self.yaml["ch"] = self.yaml["ch"] = 3
        self.backbone = parse_model(deepcopy(self.yaml), ch=[ch])
 
    def forward(self, x):
        output = self.backbone(x)
        return output
 
 
if __name__ == "__main__":
    cfg = path + '/Model.yaml'
    model = Model()
    model.eval()
    print(model)
    x = torch.ones(1, 3, 512, 512)
    output = model(x)
    torch.save(model, "model.pth")
 
 
 
    # model = torch.load('model.pth')
    # model.eval()
    # x = torch.ones(1,3,512,512)
    # input_name = ['input']
    # output_name = ['output']
    # torch.onnx.export(model, x, 'myonnx.onnx', verbose=True)

以上就是詳解Pytorch如何利用yaml定義折積網路的詳細內容，更多關於Pytorch yaml定義折積網路的資料請關注it145.com其它相關文章！