What Modules Are Needed to Build an AI Model with PyTorch

🌐 中文

English

This article aims to introduce the modules needed to efficiently build an AI model with PyTorch. It does not cover data preprocessing, model training, etc., but only focuses on model construction, serving as a quick guide to understanding the modules required to build an AI model with PyTorch.

Environment Required to Build an AI Model with PyTorch

You can refer to the article I wrote earlier; I won’t spend time reiterating it here.

Libraries Required to Build an AI Model with PyTorch (only some commonly used libraries are introduced)

1. torch

torch is the core library of PyTorch, providing various functionalities for building and training neural networks. The torch library contains many classes and functions for constructing neural networks, including tensors, automatic differentiation, neural network modules, optimizers, etc.

1
import torch

2. torch.nn

torch.nn is the model skeleton in PyTorch for building neural networks. It provides many classes and functions for constructing neural networks, including linear layers, convolutional layers, activation functions, loss functions, etc.

1
import torch.nn as nn

3. torch.optim

torch.optim is the optimizer module in PyTorch for neural networks. It provides many classes and functions for optimizing neural networks for backpropagation, including the Adam optimizer, SGD optimizer, etc.

1
import torch.optim as optim

Modules Required to Build an AI Model with PyTorch (only some commonly used modules and their corresponding parameters are introduced)

After understanding these basic libraries, we can start learning about the modules needed to build a simple model.

nn.Module

First is our most important model skeleton: nn.Module.
nn.Module is the base class and skeleton in PyTorch for building neural networks. It is used to define the structure and parameters of a neural network. In subclasses of nn.Module, we can define the layers, parameters, and forward propagation process of the neural network.

1
2
3
4
5
6
7
8
9
10
import torch.nn as nn
class ByolioModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # Define the layers of the model. In the example below, a linear layer is defined.
        self.fc = nn.Linear(10, 1)
    def forward(self, x):
        # Define the forward propagation process
        x = self.fc(x)
        return x

How to call it:

1
2
model = ByolioModel()
outputs = model(inputs)         # Call the model via the __call__ method

Convolutional Layer:

nn.Conv2d

nn.Conv2d is a class in PyTorch for defining two-dimensional convolutional layers. It accepts parameters such as input channels, output channels, kernel size, stride, and padding, and internally creates a convolution kernel and a bias vector. During forward propagation, it performs convolution operations on the input with the convolution kernel, adds the bias vector, and then applies an activation function for nonlinear transformation.

1
conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)

Parameters:

• in_channels: Number of input channels
• out_channels: Number of output channels
• kernel_size: Size of the convolution kernel
• stride: Stride of the convolution
• padding: Amount of padding
• dilation: Dilation rate of the convolution kernel
• groups: Number of groups for grouped convolution
• bias: Whether to use a bias vector (default is True)

  Fully Connected Layer:

  nn.Linear

  nn.Linear is a class in PyTorch for defining fully connected layers. It accepts input and output dimensions as parameters and internally creates a weight matrix and a bias vector. During forward propagation, it multiplies the input with the weight matrix, adds the bias vector, and then applies an activation function for nonlinear transformation.

  1
  linear = nn.Linear(in_features, out_features, bias=True)
  

  Parameters:

• in_features: Dimensionality of input features
• out_features: Dimensionality of output features
• bias: Whether to use a bias vector (default is True)

  Activation Layer:

  nn.ReLU

  nn.ReLU is a class in PyTorch for defining the ReLU activation function. It accepts input as a parameter and internally computes the value of the ReLU function. During forward propagation, it applies the ReLU function to the input and then performs nonlinear transformation via the activation function.

  1
  relu = nn.ReLU(inplace=False)
  

  Parameters:

• inplace: Whether to replace the input tensor’s value directly (default is False)

  nn.Softmax

  nn.Softmax is a class in PyTorch for defining the Softmax activation function. It accepts input as a parameter and internally computes the value of the Softmax function. During forward propagation, it applies the Softmax function to the input and then performs nonlinear transformation via the activation function.

  1
  softmax = nn.Softmax(dim=None)
  

  Parameters:

• dim: The dimension along which Softmax is computed (default is None)

  nn.Sigmoid

  nn.Sigmoid is a class in PyTorch for defining the Sigmoid activation function. It accepts input as a parameter and internally computes the value of the Sigmoid function. During forward propagation, it applies the Sigmoid function to the input and then performs nonlinear transformation via the activation function.

  1
  sigmoid = nn.Sigmoid()
  

  Pooling Layer:

  nn.MaxPool2d

  nn.MaxPool2d is a class in PyTorch for defining max pooling layers. It accepts parameters such as pooling window size, stride, and padding, and internally computes the max pooling operation. During forward propagation, it applies the pooling window to the input and then performs downsampling via the pooling operation.

  1
  maxpool2d = nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)
  

  Parameters:

• kernel_size: Size of the pooling window
• stride: Stride of the pooling operation (default is kernel_size)
• padding: Amount of padding
• dilation: Dilation rate of the pooling operation
• return_indices: Whether to return the indices of the pooling operation (default is False)
• ceil_mode: Whether to use ceiling mode for the pooling operation (default is False)

  Batch Normalization Layer:

  nn.BatchNorm2d (Please set the model to training mode during training and evaluation mode during testing)

  nn.BatchNorm2d is a class in PyTorch for defining batch normalization layers. It accepts the number of input channels as a parameter and internally computes the batch normalization operation. During forward propagation, it applies the batch normalization operation to the input and then performs normalization via the batch normalization operation.

  1
  batchnorm2d = nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  

  Parameters:

• num_features: Number of input channels
• eps: A small value for numerical stability
• momentum: Momentum for computing running mean and variance
• affine: Whether to use learnable affine transformation (default is True)
• track_running_stats: Whether to track running mean and variance (default is True)

  Flatten Layer:

  nn.Flatten

  nn.Flatten is a class in PyTorch for defining flatten layers. It accepts the input dimensions as parameters and internally flattens the input into a one-dimensional vector. During forward propagation, it flattens the input into a one-dimensional vector and then performs downsampling via the flattening operation.

  1
  flatten = nn.Flatten(start_dim=1, end_dim=-1)
  

  Parameters:

• start_dim: The starting dimension for flattening
• end_dim: The ending dimension for flattening (default is -1), flattening to the last dimension

  Loss Function:

  nn.CrossEntropyLoss

  nn.CrossEntropyLoss is a class in PyTorch for defining the cross-entropy loss function. It accepts input and target labels as parameters and internally computes the cross-entropy loss. During forward propagation, it applies the loss function to the input and target labels and then computes the loss.

  1
  crossentropyloss = nn.CrossEntropyLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')
  

  Parameters:

• weight: Weight for each class (default is None)
• size_average: Whether to compute the average loss (default is None)
• ignore_index: Index of target labels to ignore (default is -100)
• reduce: Whether to reduce the loss (default is None)
• reduction: Reduction method for the loss (default is ‘mean’)

  nn.MSELoss

  nn.MSELoss is a class in PyTorch for defining the mean squared error loss function. It accepts input and target labels as parameters and internally computes the mean squared error loss. During forward propagation, it applies the loss function to the input and target labels and then computes the loss.

  1
  mse = nn.MSELoss(size_average=None, reduce=None, reduction='mean')
  

  Parameters:

• size_average: Whether to compute the average loss (default is None)
• reduce: Whether to reduce the loss (default is None)
• reduction: Reduction method for the loss (default is ‘mean’)

  Optimizer:

  torch.optim.SGD

  torch.optim.SGD is a class in PyTorch for defining the stochastic gradient descent optimizer. It accepts hyperparameters such as model parameters and learning rate and updates the model parameters according to the rules of stochastic gradient descent. After backpropagation, the optimizer uses the gradient information of the model parameters to perform parameter updates, thereby minimizing the loss function.

  1
  optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
  

  Parameters:

• model.parameters(): Parameters of the model
• lr: Learning rate
• momentum: Momentum (default is 0)

How to Use These Modules (does not include data preprocessing, model training, etc., only involves model construction)

1. Parts defined in the subclass of nn.Module and their order:
   1. Repeatable part (you can remove unnecessary layers; adjust parameters accordingly):
      Convolutional Layer + Activation Layer + Pooling Layer + Batch Normalization Layer
   2. Non-repeatable part: Fully Connected Layer + Activation Layer
2. Parts not defined in the subclass of nn.Module:
   Loss Function + Optimizer
   The usage steps are as follows (can be optimized; the following is the unoptimized version):
   1. Define:

      1
      2
      3
      4
       # Loss function (taking cross-entropy as an example)
       loss_fn = nn.CrossEntropyLoss()
       # Optimizer (taking SGD as an example)
       optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
      

   2. Use (use after calling the forward propagation function):

      1
      2
      3
      4
      5
      6
       # Loss function
       loss = loss_fn(outputs, labels)
       # Optimizer
       optimizer.zero_grad()
       loss.backward()
       optimizer.step()
      

      Example (this example illustrates the model construction order; modules may vary for different models during training):

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
class ByolioModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # Use sequential model definition
        self.model = nn.Sequential(
        nn.Conv2d(1, 10, kernel_size=3, stride=1, padding=1),
        nn.MaxPool2d(2),
        nn.BatchNorm2d(10),
        nn.Flatten(),
        nn.Linear(10, 10),
        nn.Linear(10, 1),
        nn.Sigmoid()
        )
    def forward(self, x):
        # Define the forward propagation process
        x = self.model(x)
        return x

if __name__ == '__main__':  
    # Define the model
    model = ByolioModel()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Summary

The above are the modules required to build an AI model with PyTorch. I hope this article helps you.