Python module PyTorch introduction and basic syntax

Cover image is the official PyTorch logo, captured from Github on February 22, 2025.

Introduction

PyTorch is an open-source machine learning library developed based on the Python language and is one of the most popular deep learning frameworks today. PyTorch’s origins trace back to the technical accumulation of several research institutions and projects, mainly the Torch library based on the Lua language and Facebook’s Artificial Intelligence Research team (Facebook AI Research, FAIR, now Meta AI), which developed the cross-platform high-performance computing Caffe2 deep learning framework. In 2017, companies like Microsoft, Facebook, and IBM collaborated to develop the Open Neural Network Exchange (ONNX) format, which is hosted as open-source on GitHub, enabling PyTorch to be interoperable with other deep learning frameworks.

Features

Compared to other deep learning frameworks, PyTorch offers the following advantages and features:

Integration with Python

PyTorch uses Python as its main API and has excellent support for other Python modules, such as NumPy and SciPy. Additionally, due to PyTorch’s modular design for neural networks, we can easily and quickly build models by creating layers such as linear layers, convolutional neural network layers (CNN), recurrent neural network layers (RNN), and long short-term memory layers (LSTM).

Automatic Differentiation and Gradient Calculation

Differentiating a model or calculating gradients is often a tedious task. However, PyTorch’s submodule torch.autograd can automatically compute gradients. For programmers, this makes solving various deep learning and optimization problems in a simple, intuitive way possible. We only need to write .backward() in Python to automatically compute gradients, achieving efficient and accurate backpropagation.

GPU Acceleration

PyTorch defines a class called Tensor, which is used to store or compute multi-dimensional arrays of numbers. Tensors are similar to NumPy arrays, making it easy to convert between PyTorch tensors and NumPy arrays. Unlike NumPy, PyTorch can run on Nvidia GPUs with CUDA support. By simply writing .to(device), we can move Tensors or models to the GPU for faster computations compared to the CPU.

Compatibility

PyTorch supports ONNX, allowing it to interoperate with other deep learning frameworks and easily transfer and deploy across different environments and machines.

Basic Syntax

Installing PyTorch

We need to install Python in order to use PyTorch. Below is an example of installation using Python and the pip package management system. If you’re using a different Python distribution, such as Anaconda, please refer to the relevant installation methods.

You can install PyTorch by entering the following command in the terminal or refer to the official documentation https://pytorch.org/get-started/locally/ for further instructions.

1

pip install torch torchvision torchaudio

Once the installation is complete, we can load the PyTorch module in the Python terminal.

1

import torch

And check the PyTorch version to verify if the installation was successful.

1

print(torch.__version__)

1

2.6.0+cpu

If you have an Nvidia GPU and your GPU supports CUDA, please update the GPU drivers to the latest version. After that, obtain the installation instructions from the official documentation https://pytorch.org/get-started/locally/ and run the installation command in the terminal to enable GPU computation.

1

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

We can use the following code to check if CUDA is successfully installed and accessible by PyTorch.

1
2

print(torch.__version__)
print(torch.cuda.is_available())

If CUDA is available, it will return True.

1
2

2.6.0+cu126
True

Creating a Tensor

When using PyTorch, we can create a Tensor and perform operations on it.

1
2
3
4
5
6
7
8
9

# Import PyTorch module
import torch

# Create a 2x2 Tensor
x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
# View the Tensor array stored in the variable x
print(x)
# Check which device it is on (CPU or GPU)
print(x.device)

1
2
3

tensor([[1., 2.],
        [3., 4.]])
cpu

We can also convert a NumPy array to a Tensor.

1
2
3
4
5
6
7
8

import torch
import numpy as np

# This is a NumPy array
y = np.array([[1, 2], [3, 4]])

# Convert it to a Tensor
x = torch.tensor(y, dtype=torch.float32)

Similarly, we can convert a Tensor back into a NumPy array.

1
2

# Convert the Tensor back to a NumPy array
y = x.cpu().numpy()

The .cpu() function ensures that during the conversion, the data in the Tensor can be copied from the GPU to the CPU before conversion, preventing any conversion errors.

GPU Acceleration

We can use .to("cuda") to move a variable from the CPU to the GPU.

1
2

x = x.to("cuda")
print(x.device) # Verify that the Tensor has moved to the GPU

1

cuda:0

However, having to modify the code every time it runs in different environments can be a bit troublesome. Therefore, we can rewrite it as follows to automatically determine which processor should be used.

1
2
3

x = x.to("cuda")
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)

Creating a Simple Neural Network Model

In Python, we can create a neural network model using the class function through PyTorch. Below is an example.

 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

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleNN(nn.Module): # Define the neural network and inherit from PyTorch's nn.Module
    def __init__(self): # Initialize the neural network
        # Initialize the parent module nn.Module
        super(SimpleNN, self).__init__() 

        # Define the first layer
        self.layer1 = nn.Sequential( # Define a container to execute multiple steps
            nn.Linear(2, 3), # Fully connected layer, input 2 dimensions, output 3 dimensions
            nn.ReLU(), # Apply ReLU activation function
        )

        # Define the second layer
        self.layer2 = nn.Sequential( # Define a container to execute multiple steps
            nn.Linear(3, 1), # Fully connected layer, input 3 dimensions, output 1 dimension
            nn.Sigmoid(), # Apply Sigmoid activation function
        )

    def forward(self, x): # Define forward propagation
        x = self.layer1(x) # First layer
        x = self.layer2(x) # Second layer
        return x

The example above creates a model with 2-dimensional input and 1-dimensional output, suitable for a simple binary classification problem.

After successfully creating the model, if the execution environment supports CPU acceleration, we can move the model to the GPU.

1
2
3

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SimpleNN().to(device)
print(model) # Display the model architecture

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

SimpleNN(
  (layer1): Sequential(
    (0): Linear(in_features=2, out_features=3, bias=True)
    (1): ReLU()
  )
  (layer2): Sequential(
    (0): Linear(in_features=3, out_features=1, bias=True)
    (1): Sigmoid()
  )
)

Training the Model

Next, we define the criterion and optimizer, which are the loss function and gradient descent method, respectively. Since the model defined earlier is for binary classification, we use Binary Cross Entropy Loss as the loss function. The learning rate (lr) for gradient descent is set to 0.01.

1
2

criterion = nn.BCELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

Prepare the data for training.

1
2
3

# Assuming 4 input features, each with 2 values
inputs = torch.tensor([[0.1, 0.2], [0.4, 0.5], [0.7, 0.8], [0.9, 0.1]], dtype=torch.float32).to(device)
labels = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32).to(device)  # Labels of 0 or 1

Now, let’s train the model with a simple loop.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

num_epochs = 100 # Training epochs
for epoch in range(num_epochs):
    optimizer.zero_grad() # Clear the previous gradients
    outputs = model(inputs) # Forward propagation
    loss = criterion(outputs, labels) # Calculate the loss
    loss.backward() # Compute gradients and perform backpropagation
    optimizer.step() # Update the model weights

    if (epoch + 1) % 10 == 0:  # Output loss every 10 epochs
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

Epoch [10/100], Loss: 0.7178
Epoch [20/100], Loss: 0.7162
Epoch [30/100], Loss: 0.7146
Epoch [40/100], Loss: 0.7132
Epoch [50/100], Loss: 0.7117
Epoch [60/100], Loss: 0.7103
Epoch [70/100], Loss: 0.7090
Epoch [80/100], Loss: 0.7077
Epoch [90/100], Loss: 0.7064
Epoch [100/100], Loss: 0.7052

• Forward Propagation refers to the process where data is passed through the network layers from the input to the output.
• Backpropagation is a key algorithm in training neural networks, aimed at calculating and updating the model’s weights and biases to minimize the loss function.

Testing the Model

Assume we have a set of data for testing, the prediction result would look like this:

1
2
3

test_input = torch.tensor([[0.5, 0.6]], dtype=torch.float32).to(device) # Generate test data
test_output = model(test_input) # Make prediction
print("Test output:", test_output.item()) # Output the prediction result

1

Test output: 0.41304945945739746

Since the training data is a binary classification model with 0 or 1 labels, we can classify the neural network output as label 0 if it’s less than 0.5 and label 1 if it’s greater than 0.5.

This is how we can establish a simple neural network model.

Of course, larger training datasets and more training iterations may improve the model’s prediction accuracy, but it will also increase the training time. The relationship between dataset size, training iterations, and prediction accuracy is not linear. In fact, too much data or too many training iterations can cause the prediction accuracy to decline, so finding a balance is important.

Conclusion

Neural networks are one of the core technologies in modern machine learning and artificial intelligence. With the development of deep learning, neural networks have become the foundation for many advanced technologies, including speech recognition, image recognition, time-series forecasting, and data imputation. With an in-depth understanding of neural networks, we can try to build more complex model structures and apply them in more complex fields, allowing neural networks to solve various complicated problems for us.

References

• Get Started: Select preferences and run the command to install PyTorch locally, or get started quickly with one of the supported cloud platforms.(n.d.) PyTorch. Retrieved on February 24, 2025 from https://pytorch.org/get-started/locally/
• Open Neural Network Exchange. (October 7, 2024). Wikipedia, the free encyclopedia. Retrieved on February 24, 2025 from https://en.wikipedia.org/wiki/Open_Neural_Network_Exchange
• PyTorch. (October 7, 2024). Wikipedia, the free encyclopedia. Retrieved on February 24, 2025 from https://en.wikipedia.org/wiki/PyTorch
• PyTorch. (February 24, 2025). Github. Retrieved on February 24, 2025 from https://github.com/pytorch/pytorch