详解Pytorch如何利用yaml定义卷积网络

我们可以先定义一个卷积块CBL,C指卷积Conv,B指BN层,L为激活函数,这里我用ReLu.

1

2

3

4

5

6

7

8

9

10

11

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,代码如下:

1

2

3

4

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,参考了残差块的思想。

1

2

3

4

5

6

7

8

9

10

11

12

13

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等设置】

1

2

3

4

5

6

# 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工具来打开我们的配置文件,看看都有什么内容

1

2

3

4

5

6

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是下面要说的传参过程。

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

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(),后面也会介绍到这个函数的用法。

这里先附上完整代码:

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

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)

下面开始分析代码 。

这行代码是通过列表用来存放每层内容以及输出通道数。

1

2

# 这行代码是通过列表用来存放每层内容以及输出通道数

layer, out_channels = [], ch[-1]

然后进入我们的for循环,在每一次循环中可以获得我们yaml文件中的每一层网络:f是上一层网络的输出通道【用来作为本层的输入通道】,number【网络深度,也就是该层重复几次而已】,Module_name是该层的名字,args是该层的一些参数。

1

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

接下来会碰到一个很重要的函数eval()。下行的代码首先需要判断一下我们的Module_name类型是不是字符串类型,也就是判断一下yaml中“BaseConv”是不是字符串类型,如果是,则用eval进行对应类型的转化,转成我们的BaseConv类型。 

1

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

这里我将对eval函数在深入点,如果知道这个函数用法的,就可以略去这部分。

我们先举个例子,比如我现在有个变量a="123",这个a的类型是什么呢?他是一个str类型,不是int类型。 现在我们用eval函数转一下,看看会变成什么样子。

1

2

3

4

5

>>> b = eval(a) if isinstance(a,str) else a

>>> b

123

>>> type(b)

<class 'int'>

我们可以看到,经过eval函数以后,会自动识别并转为int类型。那么我继续举例子,如果现在a="BaseConv",经过eval以后会变成什么?可以看到,这里报错了!这是为什么?这是因为我们没有导入BaseConv这个类,所以eval函数并不知道我们希望转为什么类型。所以我们需要用import导入BaseConv这个类才可以。

1

2

3

4

5

6

>>> a="BaseConv"

>>> b = eval(a) if isinstance(a,str) else a

Traceback (most recent call last):

  File "", line 1, in

  File "", line 1, in

NameError: name 'BaseConv' is not defined

当我们导入BaseConv以后,在经过eval就可以获得:

1

<class 'models.BaseConv'>

接下来是获得args中的网络参数,也是通过eval进行转化

1

2

3

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]。

1

2

3

4

5

# 更新通道

# 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转化的 。通过*args解压操作将args列表中的内容放入m中,再通过*解压操作放入nn.Sequential。

1

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

这样就可以获得我们第一次循环BaseConv了。后面的循环也是同样的反复操作而已。

1

2

3

4

5

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列表,为的是获取每次循环的输出通道,没有这一步,再下一次循环的时候将不能正确得到通道数。

1

2

3

# 更新通道列表,每次获取输出通道

ch.append(out_channels)

layer.append(model_)

然后我们就可以对模型调用进行实例化了,可以打印下模型:

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

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)

        )

      )

    )

  )

)

同时我们也可以对模型每层可视化看一下。可以看到和我们定义的模型是一样的。

完整的代码

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

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,网络,深度学习)