[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

張賢同學發表於2020-08-28
PyTorch筆記模型

本章程式碼：https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson3/module_containers.py


這篇文章來看下 PyTorch 中網路模型的建立步驟。網路模型的內容如下，包括模型建立和權值初始化，這些內容都在nn.Module中有實現。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

網路模型的建立步驟

建立模型有 2 個要素：構建子模組和拼接子模組。如 LeNet 裡包含很多卷積層、池化層、全連線層，當我們構建好所有的子模組之後，按照一定的順序拼接起來。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

這裡以上一篇文章中 `lenet.py`的 LeNet 為例，繼承`nn.Module`，必須實現`__init__()` 方法和`forward()`方法。其中`__init__()` 方法裡建立子模組，在`forward()`方法裡拼接子模組。 
class LeNet(nn.Module):
	# 子模組建立
    def __init__(self, classes):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, classes)
	# 子模組拼接
    def forward(self, x):
        out = F.relu(self.conv1(x))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        return out

當我們呼叫net = LeNet(classes=2)建立模型時，會呼叫__init__()方法建立模型的子模組。

當我們在訓練時呼叫outputs = net(inputs)時，會進入module.py的call()函式中：

    def __call__(self, *input, **kwargs):
        for hook in self._forward_pre_hooks.values():
            result = hook(self, input)
            if result is not None:
                if not isinstance(result, tuple):
                    result = (result,)
                input = result
        if torch._C._get_tracing_state():
            result = self._slow_forward(*input, **kwargs)
        else:
            result = self.forward(*input, **kwargs)
        ...
        ...
        ...

最終會呼叫result = self.forward(*input, **kwargs)函式，該函式會進入模型的forward()函式中，進行前向傳播。

在 torch.nn中包含 4 個模組，如下圖所示。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

其中所有網路模型都是繼承於`nn.Module`的，下面重點分析`nn.Module`模組。

nn.Module

nn.Module 有 8 個屬性，都是OrderDict(有序字典)。在 LeNet 的__init__()方法中會呼叫父類nn.Module的__init__()方法，建立這 8 個屬性。

    def __init__(self):
        """
        Initializes internal Module state, shared by both nn.Module and ScriptModule.
        """
        torch._C._log_api_usage_once("python.nn_module")

        self.training = True
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._backward_hooks = OrderedDict()
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()
  • _parameters 屬性：儲存管理 nn.Parameter 型別的引數
  • _modules 屬性：儲存管理 nn.Module 型別的引數
  • _buffers 屬性：儲存管理緩衝屬性，如 BN 層中的 running_mean
  • 5 個 ***_hooks 屬性：儲存管理鉤子函式

其中比較重要的是parameters和modules屬性。

在 LeNet 的__init__()中建立了 5 個子模組，nn.Conv2d()和nn.Linear()都是 繼承於nn.module，也就是說一個 module 都是包含多個子 module 的。

class LeNet(nn.Module):
	# 子模組建立
    def __init__(self, classes):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, classes)
        ...
        ...
        ...

當呼叫net = LeNet(classes=2)建立模型後，net物件的 modules 屬性就包含了這 5 個子網路模組。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

下面看下每個子模組是如何新增到 LeNet 的`_modules` 屬性中的。以`self.conv1 = nn.Conv2d(3, 6, 5)`為例，當我們執行到這一行時，首先 Step Into 進入 `Conv2d`的構造，然後 Step Out。右鍵`Evaluate Expression`檢視`nn.Conv2d(3, 6, 5)`的屬性。
[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

上面說了`Conv2d`也是一個 module，裡面的`_modules`屬性為空，`_parameters`屬性裡包含了該卷積層的可學習引數，這些引數的型別是 Parameter，繼承自 Tensor。
[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

此時只是完成了`nn.Conv2d(3, 6, 5)` module 的建立。還沒有賦值給`self.conv1 `。在`nn.Module`裡有一個機制，會攔截所有的類屬性賦值操作(`self.conv1`是類屬性)，進入到`__setattr__()`函式中。我們再次 Step Into 就可以進入`__setattr__()`。 
    def __setattr__(self, name, value):
        def remove_from(*dicts):
            for d in dicts:
                if name in d:
                    del d[name]

        params = self.__dict__.get('_parameters')
        if isinstance(value, Parameter):
            if params is None:
                raise AttributeError(
                    "cannot assign parameters before Module.__init__() call")
            remove_from(self.__dict__, self._buffers, self._modules)
            self.register_parameter(name, value)
        elif params is not None and name in params:
            if value is not None:
                raise TypeError("cannot assign '{}' as parameter '{}' "
                                "(torch.nn.Parameter or None expected)"
                                .format(torch.typename(value), name))
            self.register_parameter(name, value)
        else:
            modules = self.__dict__.get('_modules')
            if isinstance(value, Module):
                if modules is None:
                    raise AttributeError(
                        "cannot assign module before Module.__init__() call")
                remove_from(self.__dict__, self._parameters, self._buffers)
                modules[name] = value
            elif modules is not None and name in modules:
                if value is not None:
                    raise TypeError("cannot assign '{}' as child module '{}' "
                                    "(torch.nn.Module or None expected)"
                                    .format(torch.typename(value), name))
                modules[name] = value
            ...
            ...
            ...

在這裡判斷 value 的型別是Parameter還是Module，儲存到對應的有序字典中。

這裡nn.Conv2d(3, 6, 5)的型別是Module，因此會執行modules[name] = value，key 是類屬性的名字conv1，value 就是nn.Conv2d(3, 6, 5)。

總結

  • 一個 module 裡可包含多個子 module。比如 LeNet 是一個 Module，裡面包括多個卷積層、池化層、全連線層等子 module
  • 一個 module 相當於一個運算，必須實現 forward() 函式
  • 每個 module 都有 8 個字典管理自己的屬性

模型容器

除了上述的模組之外，還有一個重要的概念是模型容器 (Containers)，常用的容器有 3 個，這些容器都是繼承自nn.Module。

  • nn.Sequetial：按照順序包裝多個網路層
  • nn.ModuleList：像 python 的 list 一樣包裝多個網路層，可以迭代
  • nn.ModuleDict：像 python 的 dict 一樣包裝多個網路層，通過 (key, value) 的方式為每個網路層指定名稱。

nn.Sequetial

在傳統的機器學習中，有一個步驟是特徵工程，我們需要從資料中認為地提取特徵，然後把特徵輸入到分類器中預測。在深度學習的時代，特徵工程的概念被弱化了，特徵提取和分類器這兩步被融合到了一個神經網路中。在卷積神經網路中，前面的卷積層以及池化層可以認為是特徵提取部分，而後面的全連線層可以認為是分類器部分。比如 LeNet 就可以分為特徵提取和分類器兩部分，這 2 部分都可以分別使用 nn.Seuqtial 來包裝。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

程式碼如下： 
class LeNetSequetial(nn.Module):
    def __init__(self, classes):
        super(LeNet2, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

在初始化時，nn.Sequetial會呼叫__init__()方法，將每一個子 module 新增到 自身的_modules屬性中。這裡可以看到，我們傳入的引數可以是一個 list，或者一個 OrderDict。如果是一個 OrderDict，那麼則使用 OrderDict 裡的 key，否則使用數字作為 key (OrderDict 的情況會在下面提及)。

    def __init__(self, *args):
        super(Sequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)

網路初始化完成後有兩個子 module：features和classifier。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

而`features`中的子 module 如下，每個網路層以序號作為 key：
[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

在進行前向傳播時，會進入 LeNet 的`forward()`函式，首先呼叫第一個`Sequetial`容器：`self.features`，由於`self.features`也是一個 module，因此會呼叫`__call__()`函式，裡面呼叫

result = self.forward(*input, **kwargs)，進入nn.Seuqetial的forward()函式，在這裡依次呼叫所有的 module。

    def forward(self, input):
        for module in self:
            input = module(input)
        return input

在上面可以看到在nn.Sequetial中，裡面的每個子網路層 module 是使用序號來索引的，即使用數字來作為 key。一旦網路層增多，難以查詢特定的網路層，這種情況可以使用 OrderDict (有序字典)。程式碼中使用

class LeNetSequentialOrderDict(nn.Module):
    def __init__(self, classes):
        super(LeNetSequentialOrderDict, self).__init__()

        self.features = nn.Sequential(OrderedDict({
            'conv1': nn.Conv2d(3, 6, 5),
            'relu1': nn.ReLU(inplace=True),
            'pool1': nn.MaxPool2d(kernel_size=2, stride=2),

            'conv2': nn.Conv2d(6, 16, 5),
            'relu2': nn.ReLU(inplace=True),
            'pool2': nn.MaxPool2d(kernel_size=2, stride=2),
        }))

        self.classifier = nn.Sequential(OrderedDict({
            'fc1': nn.Linear(16*5*5, 120),
            'relu3': nn.ReLU(),

            'fc2': nn.Linear(120, 84),
            'relu4': nn.ReLU(inplace=True),

            'fc3': nn.Linear(84, classes),
        }))
        ...
        ...
        ...

總結

nn.Sequetial是nn.Module的容器，用於按順序包裝一組網路層，有以下兩個特性。

  • 順序性：各網路層之間嚴格按照順序構建，我們在構建網路時，一定要注意前後網路層之間輸入和輸出資料之間的形狀是否匹配
  • 自帶forward()函式：在nn.Sequetial的forward()函式裡通過 for 迴圈依次讀取每個網路層，執行前向傳播運算。這使得我們我們構建的模型更加簡潔

nn.ModuleList

nn.ModuleList是nn.Module的容器，用於包裝一組網路層，以迭代的方式呼叫網路層，主要有以下 3 個方法：

  • append()：在 ModuleList 後面新增網路層
  • extend()：拼接兩個 ModuleList
  • insert()：在 ModuleList 的指定位置中插入網路層

下面的程式碼通過列表生成式來迴圈迭代建立 20 個全連線層，非常方便，只是在 forward()函式中需要手動呼叫每個網路層。

class ModuleList(nn.Module):
    def __init__(self):
        super(ModuleList, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(20)])

    def forward(self, x):
        for i, linear in enumerate(self.linears):
            x = linear(x)
        return x


net = ModuleList()

print(net)

fake_data = torch.ones((10, 10))

output = net(fake_data)

print(output)

nn.ModuleDict

nn.ModuleDict是nn.Module的容器，用於包裝一組網路層，以索引的方式呼叫網路層，主要有以下 5 個方法：

  • clear()：清空 ModuleDict
  • items()：返回可迭代的鍵值對 (key, value)
  • keys()：返回字典的所有 key
  • values()：返回字典的所有 value
  • pop()：返回一對鍵值，並從字典中刪除

下面的模型建立了兩個ModuleDict：self.choices和self.activations，在前向傳播時通過傳入對應的 key 來執行對應的網路層。

class ModuleDict(nn.Module):
    def __init__(self):
        super(ModuleDict, self).__init__()
        self.choices = nn.ModuleDict({
            'conv': nn.Conv2d(10, 10, 3),
            'pool': nn.MaxPool2d(3)
        })

        self.activations = nn.ModuleDict({
            'relu': nn.ReLU(),
            'prelu': nn.PReLU()
        })

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x


net = ModuleDict()

fake_img = torch.randn((4, 10, 32, 32))

output = net(fake_img, 'conv', 'relu')
# output = net(fake_img, 'conv', 'prelu')
print(output)

容器總結

  • nn.Sequetial：順序性，各網路層之間嚴格按照順序執行，常用於 block 構建，在前向傳播時的程式碼呼叫變得簡潔
  • nn.ModuleList：迭代行，常用於大量重複網路構建，通過 for 迴圈實現重複構建
  • nn.ModuleDict：索引性，常用於可選擇的網路層

PyTorch 中的 AlexNet

AlexNet 是 Hinton 和他的學生等人在 2012 年提出的卷積神經網路，以高出第二名 10 多個百分點的準確率獲得 ImageNet 分類任務冠軍，從此卷積神經網路開始在世界上流行，是劃時代的貢獻。

AlexNet 特點如下：

  • 採用 ReLU 替換飽和啟用 函式，減輕梯度消失
  • 採用 LRN (Local Response Normalization) 對資料進行區域性歸一化，減輕梯度消失
  • 採用 Dropout 提高網路的魯棒性，增加泛化能力
  • 使用 Data Augmentation，包括 TenCrop 和一些色彩修改

AlexNet 的網路結構可以分為兩部分：features 和 classifier。

[PyTorch 學習筆記] 3.1 模型建立步驟與 nn.Module

在`PyTorch`的計算機視覺庫`torchvision.models`中的 AlexNet 的程式碼中，使用了`nn.Sequential`來封裝網路層。 
class AlexNet(nn.Module):

    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

參考資料


如果你覺得這篇文章對你有幫助，不妨點個贊，讓我有更多動力寫出好文章。

相關文章