This article is a study of some functions. Functions mainly include the following four conveniences:

  • Function of model construction:add_module.add_module.add_module
  • Access submodules:add_module.add_module.add_module.add_module
  • Network traversal:

add_module,add_module

  • Model saving and loading:add_module.add_module.add_module

1. Model building function

Burn.nn. Module is the base class for all networks and is inherited from the Classes in the PyTorch implementation model (as mentioned in the previous course). In building a Module, a Module is a Module that contains other modules. Similarly, you can define a small network Module and then treat it as a component of another network. Therefore, the network structure is a presentation tree structure.

Let’s start by simply defining a network:

import torch.nn as nn
import torch 
class MyNet(nn.Module) :
    def __init__(self) :
        super(MyNet,self).__init__()
        self.conv1 = nn.Conv2d(3.64.3)
        self.conv2 = nn.Conv2d(64.64.3)

    def forward(self,x) :
        x = self.conv1(x)
        x = self.conv2(x)
        return x
net = MyNet()
print(net)
Copy the code

Output results:MyNetThere are two properties inconv1andconv2It’s two convolution layers, propagating forwardforward, the two convolution layers are called in turn to realize the functions of the network.

1.1 add_module

This is the most common way to define a network, and in some projects you’ll see the method add_module. Let’s rewrite the above network in this way:

class MyNet(nn.Module) :
    def __init__(self) :
        super(MyNet,self).__init__()
        self.add_module('conv1',nn.Conv2d(3.64.3))
        self.add_module('conv2',nn.Conv2d(64.64.3))

    def forward(self,x) :
        x = self.conv1(x)
        x = self.conv2(x)
        return x
Copy the code

Add_module (name,layer) implements the same function as self.name=layer. Add_module (name,layer) implements the same function as self.name=layer. Anyway, let’s get familiar with this.

Both of the above methods are layer by layer, and if the network is complex, you need to write a lot of repetitive code. So let’s look at the construction of network modules,torch.nn.ModuleList and torch.nn.Sequential

1.2 ModuleList

ModuleList literally holds the network layer as a list. In this way, you can first build the layer needed for the network, save it in a list, and then add it to the network through the ModuleList method.

class MyNet(nn.Module) :
    def __init__(self) :
        super(MyNet,self).__init__()
        self.linears = nn.ModuleList(
            [nn.Linear(10.10) for i in range(5)])def forward(self,x) :
        for l in self.linears:
            x = l(x)
        return x
net = MyNet()
print(net)
Copy the code

The output is:

This ModuleList is used to read the config file to build the network model. The following example uses the VGG model as an example:

vgg_cfg = [64.64.'M'.128.128.'M'.256.256.256.'C'.512.512.512.'M'.512.512.512.'M']

def vgg(cfg, i, batch_norm=False) :
    layers = []
    in_channels = i
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        elif v == 'C':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return layers

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

        self.vgg = nn.ModuleList(vgg(vgg_cfg,3))

    def forward(self,x) :

        for l in self.vgg:
            x = l(x)
m1 = Model1()
print(m1)
Copy the code

Read the configuration file vgg_cfg and create a Layer list based on this file. Then use ModuleList to add to the network. This allows you to quickly create different networks (in the case of the above example, you can quickly modify the configuration file and then change the network structure).

1.3 Sequential

Sequential actually uses Sequential more often in small projects that you do yourself. Again, rewrite the original simplest example:

class MyNet(nn.Module) :
    def __init__(self) :
        super(MyNet,self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3.64.3),
            nn.Conv2d(64.64.3))def forward(self,x) :
        x = self.conv(x)
        return x
net = MyNet()
print(net)
Copy the code

Running result:

If you look closely, the network layer in Seqential is numbered by default, which we started withself.conv1andself.conv2In conv1 and conv2.

How do we change the name of the Sequential network layer? Here you need to use the collections.OrderedDict ordered dictionary. Sequential supports ordered dictionary construction.

from collections import OrderedDict 
class MyNet(nn.Module) :
    def __init__(self) :
        super(MyNet,self).__init__()
        self.conv = nn.Sequential(OrderedDict([
            ('conv1',nn.Conv2d(3.64.3)),
            ('conv2',nn.Conv2d(64.64.3))))def forward(self,x) :
        x = self.conv(x)
        return x
net = MyNet()
print(net)
Copy the code

Output results:

1.4 small summary

  • A separate network layer or submodule can be addedadd_moduleOr directly attribute;
  • ModuleListYou can add a Module List to the network.
  • SequentialGenerate a Module Module in sequence.The recommended custom here is to build using OrderedDict methods. Add a canonical name to the network layer to facilitate subsequent lookup and traversal

2 Traversal model structure

3 Save and load