Published in

Understanding and using the class structure for creating models with PyTorch 🔥

Photo by Hitesh Choudhary on Unsplash

For creating the general models in PyTorch we use nn.Sequential API and create our neural network model.

model = nn.Sequential(

here the problem with this structure is that you can’t have hybrid models. Create complicated models and you are not having full control over your model and pass the data on the model.

For having whole control of the model and creating a model which will be helpful to us in real machine learning we can do the same using class like structure and inheriting the torch.nn.Module from PyTorch module class

import torch 
import torch.nn as nn
class MyModel(torch.nn.Module):
def __init__(self):
self.lay1 = nn.Conv2d(1,20,5)
self.lay2 = nn.Conv2d(20,64,5)
self.flatten = nn.Flatten()
self.relu = nn.ReLU()

def forward(self,inputs):
ins = self.relu(self.lay1(inputs))
ins = self.relu(self.lay2(ins))
output = self.flatten(ins)

here the advantage is you are having whole control over how does the data flow into the model you can have multiple layers and hybrid structures like this shown in the image below

While you are controlling the whole flow of the data which is inside your model let’s go through what does each part we have in our code previously shows

What is __init__ in class structure?

This method helps in object-oriented programming in python. this is a reserved method and it acts as a constructor to our class.

constructor is the first thing which initializes the contents in the class it’s the first method getting called when we create an object of our model and will get our Conv and Linear layers initialized with random weights

What is super(MyModel,self).__init__() doing ?

The super(MyModel, self).__init__() refers to the fact that this is a subclass of nn.Module and is inheriting all methods. With defining this we are ready to use all precooked code that is already implemented in nn.Module class of PyTorch.

In simple words the super method lets you use all the modules implemented in torch.nn.Module class.

What is self in class?

self represents the instance of the class this attribute can be used to access the attributes and the methods (functions ) defined inside the class. In the above example, we are making use of self to declare the layers so that we can use them in forward method.

What is the forward method?

this method is the driving method for our model this is where the weights of our layers get updated after calling our model in the training loop.

thanks for reading my blog :) follow for more say hi to me in comments it gives me encouragement for writing more blogs :) have a good day :)




Data Scientists must think like an artist when finding a solution when creating a piece of code. ⚪️ Artists enjoy working on interesting problems, even if there is no obvious answer ⚪️ 🔵 Follow to join our 18K+ Unique DAILY Readers 🟠

Recommended from Medium

Code Review: How to transfer images using PyTorch

Using VS Code To Build An Azure Function And SQL Database App With .Net Core On Azure Cloud

SAFe, get real with Scrum or get out!

Dev Update August 2019

Deprek8ion: Staying ahead of K8s API deprecations

7 Software Development Practices That Might Not Be Practical

A Beginner’s Guide to CSS Front End Frameworks

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store


More from Medium

Crowdsourcing — A Step Towards Advanced Deep Learning

4 Up-to-Date Techniques for Image Data Augmentation

ResNet — Understand and Implement from scratch

Automatic Liver Segmentation — Part 1/4: Introduction