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.