Pytorch uses nn.ModuleList() and nn.Sequential() to write the neural network model

This blog describes how to use nn.ModuleList() and nn.Sequential() to simplify the way models are created. The traditional methods, nn.ModuleList() and nn.Sequential() are used to create a model to fit the sin function, train the model to save the results, and experience the convenience of the two functions.

1, Using traditional methods

Create model

import numpy as np
import torch
import torch.nn as nn

#Prepare data
data=np.linspace(-2*np.pi,2*np.pi,400) # -All points divided into 400 parts from 2PI to 2PI
x=torch.tensor(data.reshape(400,-1),dtype=torch.float)
y=torch.tensor(np.sin(data.reshape(400,-1)),dtype=torch.float)

# Create model
class model1(nn.Module):
    def __init__(self):
        super(model1,self).__init__()
        self.linear1=nn.Linear(1,10)
        self.activation1=nn.ReLU()
        self.linear2=nn.Linear(10,100)
        self.activation2=nn.ReLU()
        self.linear3=nn.Linear(100,10)
        self.activation3=nn.ReLU()
        self.linear4=nn.Linear(10,1)
    def forward(self,x):
        out=self.linear1(x)
        out=self.activation1(out)
        out=self.linear2(out)
        out=self.activation2(out)
        out=self.linear3(out)
        out=self.activation3(out)
        out=self.linear4(out)
        return out

# Create model instance
model=model1()
## Set optimizer and loss function
optimizer=torch.optim.Adam(model1.parameters(model),lr=0.005)
Loss=nn.MSELoss()

# Start training
for i in range(3000):
    y_predict=model(x)
    loss=Loss(y_predict,y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (i+1)%100==0:
        print("[training] step: {0} , loss: {1}".format(i+1,loss))

## Show and save pictures
import matplotlib.pyplot as plt
predict=model(x)
plt.plot(data,np.sin(data),color="red")
plt.plot(data,predict.detach().numpy(),color="blue")
plt.savefig(figsize=(10,10),fname="result.png")
plt.show()

Output results:
Red is the training data set (real data), and blue is the prediction result.

It can be seen that after 3000 times of training, the fitting is pretty good

2, Use the nn.ModuleList() function

The traditional model creation method code is as follows:

class model1(nn.Module):
    def __init__(self):
        super(model1,self).__init__()
        self.linear1=nn.Linear(1,10)
        self.activation1=nn.ReLU()
        self.linear2=nn.Linear(10,100)
        self.activation2=nn.ReLU()
        self.linear3=nn.Linear(100,10)
        self.activation3=nn.ReLU()
        self.linear4=nn.Linear(10,1)
    def forward(self,x):
        out=self.linear1(x)
        out=self.activation1(out)
        out=self.linear2(out)
        out=self.activation2(out)
        out=self.linear3(out)
        out=self.activation3(out)
        out=self.linear4(out)
        return out

Use ModuleList to simplify writing.
What needs to be said here is that the ModuleList can store multiple models. In the traditional method, a model needs to write a forward, but if they are saved to a ModuleList, a forward can be used.
ModuleList is a subclass of module. When it is used in module, it can be automatically recognized as a child module.
When nn.ModuleList is added as a member of nn.Module object (that is, when we add modules to our network), all nn.Module parameters inside nn.ModuleList are also added as parameters of our network.
Using ModuleList can also make the network structure flexible. For example, I need to set the number of layers of the network as a variable. The traditional method needs to be implemented with the help of list, which is inconvenient. Using ModuleList can simplify this operation.

Use ModuleList to replace the original method. The modified model is as follows: (input data, optimizer, loss, training process and other inconveniences are omitted here)

class model2(nn.Module):
    def __init__(self):
        super(model2, self).__init__()
        self.layers=nn.ModuleList([
            nn.Linear(1,10), nn.ReLU(),
            nn.Linear(10,100),nn.ReLU(),
            nn.Linear(100,10),nn.ReLU(),
            nn.Linear(10,1)])
    def forward(self,x):
        out=x
        for i,layer in enumerate(self.layers):
            out=layer(out)
        return out

Other usage

ModuleList has a similar usage to list. In fact, it can be regarded as a combination of Module and list.
In addition to passing in a list of modules when creating a ModuleList, you can also use the extend function and the append function to add models

extend method
Similar to list, the parameter is a list whose element is Module. The effect of this method is to add all modules in the list to ModuleList.

self.linears.extend([nn.Linear(size1, size2) for i in range(1, num_layers)])

append method

self.linears.append(nn.Linear(size1, size2)

3, Use nn.Sequential()

The modified model is as follows

class model3(nn.Module):
    def __init__(self):
        super(model3, self).__init__()
        self.network=nn.Sequential(
            nn.Linear(1,10),nn.ReLU(),
            nn.Linear(10,100),nn.ReLU(),
            nn.Linear(100,10),nn.ReLU(),
            nn.Linear(10,1)
        )
    def forward(self, x):
        return self.network(x)

You can use self.network[0] to obtain the first Linear sub model. Since each sub model has no unique name, you can only use the digital index to obtain it.

How to add sub models

Like ModuleList, the Sequential method can not only add sub models directly when creating, but also add new models after creating, and can define unique name indexes for sub models for easy access. Examples are as follows:

self.network.add_module("linear1",nn.Linear(100,100))

If you want to obtain this sub model, you can do the following:

linear=self.network.linear1

Another way to write

from collections import OrderedDict
self.network=nn.Sequential(OrderedDict(
            ("linear1",nn.Linear(1,10)),
            ("activation1",nn.ReLU()),
            ("linear2",nn.Linear(10,100))
        ))

Similarly, the acquisition sub model can be obtained by name, which is consistent with the second writing method.

Full reference code

The complete reference code is as follows. Changing the marked position can switch different building methods of the model

import numpy as np
import torch
import torch.nn as nn

#Prepare data
data=np.linspace(-2*np.pi,2*np.pi,400) # -All points divided into 400 parts from 2PI to 2PI
x=torch.tensor(data.reshape(400,-1),dtype=torch.float)
y=torch.tensor(np.sin(data.reshape(400,-1)),dtype=torch.float)

# Create model
class model1(nn.Module):
    """
    The first construction method
    """
    def __init__(self):
        super(model1,self).__init__()
        self.linear1=nn.Linear(1,10)
        self.activation1=nn.ReLU()
        self.linear2=nn.Linear(10,100)
        self.activation2=nn.ReLU()
        self.linear3=nn.Linear(100,10)
        self.activation3=nn.ReLU()
        self.linear4=nn.Linear(10,1)
    def forward(self,x):
        out=self.linear1(x)
        out=self.activation1(out)
        out=self.linear2(out)
        out=self.activation2(out)
        out=self.linear3(out)
        out=self.activation3(out)
        out=self.linear4(out)
        return out
class model2(nn.Module):
    """
    The second construction method
    """
    def __init__(self):
        super(model2, self).__init__()
        self.layers=nn.ModuleList([
            nn.Linear(1,10),
            nn.ReLU(),
            nn.Linear(10,100),
            nn.ReLU(),
            nn.Linear(100,10),
            nn.ReLU(),
            nn.Linear(10,1)
        ])

    def forward(self,x):
        out=x
        for i,layer in enumerate(self.layers):
            out=layer(out)
        return out

    def printf(self):
        for i,layer in enumerate(self.layers):
            print("The first{0}Layers:{1}".format(i,layer))

class model3(nn.Module):
    """
    The third construction method
    """
    def __init__(self):
        super(model3, self).__init__()
        self.network=nn.Sequential(
            nn.Linear(1,10),nn.ReLU(),
            nn.Linear(10,100),nn.ReLU(),
            nn.Linear(100,10),nn.ReLU(),
            nn.Linear(10,1)
        )

    def forward(self, x):
        return self.network(x)
# Create model instance
#####################################################################################################
################################      Here you can switch the mode of building the model      ############################################
#####################################################################################################
model=model3()

## Set optimizer and loss function
optimizer=torch.optim.Adam(model1.parameters(model),lr=0.005)
Loss=nn.MSELoss()

# Start training
for i in range(3000):
    y_predict=model(x)
    loss=Loss(y_predict,y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (i+1)%100==0:
        print("[training] step: {0} , loss: {1}".format(i+1,loss))

## Show and save pictures
import matplotlib.pyplot as plt
predict=model(x)
plt.plot(data,np.sin(data),color="red")
plt.plot(data,predict.detach().numpy(),color="blue")
plt.savefig(figsize=(10,10),fname="result.png")
plt.show()

--------
Copyright notice: This article is the original article of CSDN blogger "pretending to be a bad humble gentleman", which follows the CC 4.0 BY-SA copyright agreement. Please attach the original source link and this notice for reprint.
Original link:

Tags: Linux Pytorch Deep Learning Visual Studio Code

Posted on Sun, 21 Nov 2021 22:57:07 -0500 by shdt