AI Pool
Log in
Discussions>How to create a custom layer in Pytorch?>

How to create a custom layer in Pytorch?

I used to use the Keras library for designing a network, but recently I've changed it to Pytorch and want to design a custom layer with its functionality.

Deep LearningAI Development/Engineering
pytorchcustom-layer
2 votesLP235.00
1 Answers
JW297.00
2

In Pytorch, basically, every part of the network is a Module, which can be customized.

import torch
import torch.nn as nn

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

    self.conv1 = nn.Sequential(nn.Conv2d(128, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU())
    self.conv2 = nn.Sequential(nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU())
    self.out = nn.Sigmoid()

  def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    out = self.out(x)
    return out

This is a simple example of how to create a custom layer in Pytorch. The functionality can be different and each time you need to define the instances in __init__ function. It lets to collect all trainable variables into a group called MyLayer. Those variables need to be visible for the optimizer you chose and nn.Module handles it automatically.

Reply
Couldn't find what you were looking for?and we will find an expert to answer.
How helpful was this page?