Differences between torch.nn and torch.nn.functional

A neural network is a subset of machine learning that uses the interconnected layers of nodes to process the data and find patterns. These patterns or meaningful insights help us in strategic decision-making for various use cases. PyTorch is a Deep-learning framework that allows us to do this.

It includes various modules for creating and training the neural networks. Among these, torch.nn and torch.nn.functional are popular. Let us discuss them in more detail in this article.

What is PyTorch?

Facebook developed PyTorch in 2016 for building Machine Learning Applications. You can use this framework to perform the Deep Learning operations such as Natural Language Processing. It basically includes the Graphical Processing Unit and Deep Neural Networks for Machine Learning-based Tasks. It is an open-source and hence, it has modules for various tasks like research prototyping and production deployment. It follows the dynamic graph computation rather than the static graph approach. Due to this, it can immediately execute operations. Now, let us discuss two important modules that are used to train the layers of the neural networks.

What is torch.nn?

The torch.nn module is the collection that includes various pre-defined layers, activation functions, loss functions, and utilities for building and training the Neural Networks. All the components include various mathematical functions and operations for training the Deep Learning Models. Its base class includes the parameters, functions, and layers.nn module.

It mainly includes four classes namely the Parameters, Containers, Layers, and Functions which are discussed briefly as follows:

• Parameters: The torch.nn.Parameter() subclass includes the Tensors, which is a multi-dimensional array containing the learnable parameters like weights/biases.

• Containers: It allows us to create complex neural networks. For example, torch.nn.Sequential() combines the different layers and torch.nn.ParameterDict() stores the parameters.

• Layers: These layers perform specific operations like convolution, activation, etc. on data.

• Functions: These include the loss functions, similarity functions, etc. to apply to the data. For example, we can use the torch.nn.CrossEntropyLoss() loss function to evaluate the difference between actual value and predicted value.

Now, let us see how these things differ from the torch.nn.functional module.

What is torch.nn.functional?

The torch.nn.functional includes a functional approach to work on the input data. It means that the functions of the torch.nn.functional module work directly on the input data, without creating an instance of a neural network layer. Also, the functions in this module are stateless. Hence, they do not include any learnable parameters like weights and biases which are modified as the model gets trained. They perform direct operations like convolution, activation, pooling, etc.

Now, let us see the difference between torch.nn and torch.nn.functional module.

What are Stateless and Stateful Models?

• Stateless models are models that do not retain information between calls. Each inference or prediction is made independently, without any memory of past inputs or outputs. These models are typically used for simple, independent calculations where the context of previous inputs is not important.
• Stateful models, on the other hand, retain information between calls. They maintain a state that can include weights, hidden states, or other parameters that are updated and used in subsequent calls. Stateful models are often used for sequential data, such as time series, where the context of past inputs is important for making predictions.

Differences Between torch.nn and torch.nn.functional

How to choose between torch.nn and torch.nn.functional?

Both torch.nn and functional have methods such as Conv2d, Max Pooling, ReLU, etc. But when it comes to the implementation, there is a slight difference between them. Let us now discuss when to choose the torch.nn module and when we should opt for the torch.nn.functional.

• You should use the ‘torch.nn’ when you want to train the layers with learnable parameters. But if you want to make operations simple, ‘torch.nn.functional’ is suitable as it has stateless operations without any parameters.

• The ‘torch.nn’ module is less flexible than the ‘torch.nn.functional’ module. This is because we can use the torch.nn.functional to define custom operations and use them in the neural network. So, if you want flexibility, the torch.nn.functional module is used.

• If you want to create and train the neural network using the pre-defined layers, torch.nn is suitable. But if you want to customize some parts of the neural network, you can use the torch.nn.functional within the custom modules.

Conclusion

The torch.nn and torch.nn.functional module allows us to use various operations to develop the Deep Learning Neural Network. They include the layers, functions, and components that can process the data. However, they differ in terms of use cases. The ‘torch. nn’ module is less flexible with predefined layers.

But the ‘torch.nn.functional’ module provides the options to customize the Network Layers. In addition to this, their efficiency depends on the use and applications of both layers. After gaining a clear understanding of their difference, we can easily choose the right tool for our business.