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.
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.
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.