Initializing a Neural Network Model in PyTorch

Lesson Overview

Hey there, budding data scientist! In this lesson, we continue our journey into the world of PyTorch. More specifically, we are going to explore how to initialize a basic Neural Network in PyTorch.

In order to do this, you will learn how to:

  • Utilize PyTorch modules
  • Build a simple neural network
  • Define forward pass
  • Instantiate the neural network model
  • Print the model's architecture

Let's dive in!

Introduction to PyTorch Modules

Before building a neural network, let's understand what a PyTorch module is.

PyTorch’s modules are encapsulated as Python classes and serve as building blocks for designing models. The base to these is the nn.Module class. Any model you create in PyTorch is a subclass of the nn.Module. Let's check the initial lines in our code:

import torch
import torch.nn as nn

Here we imported PyTorch’s nn module, a base class for all neural network modules.

Build a Simple Neural Network

Let's dive into defining our basic neural network. The code snippet below might seem a bit daunting at first, but don’t worry—we’ll break it down step-by-step:

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.layer1 = nn.Linear(in_features=2, out_features=10)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(in_features=10, out_features=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        x = self.sigmoid(x)
        return x

This code lays out the structure of our neural network. Let's explain each part in detail:

Class Definition and Inheritance

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()

We start by defining a class named SimpleNN which inherits from PyTorch's nn.Module. This inheritance allows SimpleNN to build on the robust features and functionalities provided by PyTorch.

In the __init__ method, calling super(SimpleNN, self).__init__() ensures that nn.Module's base properties are correctly initialized. This is crucial for setting up our neural network layers and activation functions properly.

Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal