PyTorch model container and AlexNet construction

Articles and codes have been archived in [Github warehouse: https://github.com/timerring/dive-into-AI ] or the public account [AIShareLab] can also be obtained by replying to the pytorch tutorial .

Model container and AlexNet construction

In addition to the above-mentioned modules, another important concept is the model container (Containers). There are 3 commonly used containers, and these containers are all inherited from nn.Module.

  • nn.Sequential: Wrap multiple network layers sequentially
  • nn.ModuleList: wraps multiple network layers like python's list, and can be iterated
  • nn.ModuleDict: wraps multiple network layers like python's dict, and assigns a name to each network layer by means of (key, value).

No. Sequetial

In deep learning, the two steps of feature extraction and classifier are integrated into a neural network. In the convolutional neural network, the previous convolutional layer and pooling layer can be considered as the feature extraction part, while the latter fully connected layer can be considered as the classifier part. For example, LeNet can be divided into two parts: feature extraction and classifier , both of which can be nn.Seuqtialpackaged separately.

code show as below:

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

At initialization time, the method nn.Sequetialis called __init__()to add each submodule to its own _modulesproperties. As you can see here, the parameter we pass in can be a list or an OrderDict. If it is an OrderDict, then use the key in the OrderDict, otherwise use the number as the key.

    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)

After the network initialization is complete, there are two subsections module: featuresand classifier.

featuresThe sub-modules in are as follows, and each network layer uses the serial number as the key :

When performing forward propagation, it will enter forward()the function of LeNet, and first call the first Sequetialcontainer: self.features, since self.featuresit is also a module, it will call __call__()the function, which calls

result = self.forward(*input, **kwargs), the entered nn.Seuqetialfunction forward(), where all modules are called in turn. The output of the previous module is the input of the next module.

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

As can be seen above nn.Sequetial, each sub-network layer module inside is indexed by a serial number, that is, a number is used as a key.

Once the number of network layers increases, it is difficult to find a specific network layer. In this case, OrderDict (ordered dictionary) can be used. You can compare it with the above code

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),
        }))
        ...
        ...
        ...

Summarize

nn.SequetialYes nn.Modulecontainer, used to package a set of network layers in order, has the following two characteristics.

  • Sequence: The network layers are constructed strictly in order. When we build the network, we must pay attention to whether the shapes of the input and output data between the front and rear network layers match.
  • Built-in forward()function: in nn.Sequetialthe forward()function, each network layer is read in turn through a for loop, and the forward propagation operation is performed. This makes the model we build more compact

nn.ModuleList

nn.ModuleListIt is nn.Modulea container, which is used to wrap a set of network layers and call the network layers in an iterative manner. There are three main methods:

  • append(): Add a network layer after ModuleList
  • extend(): Splicing two ModuleLists
  • insert(): Insert the network layer in the specified position of ModuleList

The following code loops and iterates to create 20 fully connected layers through list generation, which is very convenient, but it forward()needs to manually call each network layer in the function.

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.ModuleDictIt is nn.Modulea container, used to wrap a group of network layers, and call the network layer by index. There are mainly the following five methods:

  • clear(): Clear ModuleDict
  • items(): returns an iterable key-value pair (key, value)
  • keys(): returns all keys of the dictionary
  • values(): returns all values ​​of the dictionary
  • pop(): returns a pair of key values ​​and removes them from the dictionary

The following model creates two ModuleDict: self.choicesand self.activations, and executes the corresponding network layer by passing in the corresponding key during forward propagation.

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)

container summary

  • nn.Sequetial: Sequence, each network layer is executed strictly in order, often used in block construction, code calls during forward propagation become concise
  • nn.ModuleList: Iterate rows, often used for a large number of repetitive network construction, and achieve repeated construction through for loops
  • nn.ModuleDict: indexed, often used for optional network layers

Implementation of AlexNet

The features of AlexNet are as follows:

  • Use ReLU to replace the saturated activation function to reduce gradient disappearance
  • Use LRN (Local Response Normalization) to locally normalize the data to reduce gradient disappearance
  • Use Dropout to improve the robustness of the network and increase the generalization ability
  • Using Data Augmentation, including TenCrop and some color modification

The network structure of AlexNet can be divided into two parts: features and classifier.

torchvision.modelsThe code of AlexNet can be found in the computer vision library , and it can be seen that it is used nn.Sequentialto encapsulate the network layer.

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

Guess you like

Origin blog.csdn.net/m0_52316372/article/details/131674512