Ah. Okay. How do we implement that? Something like the following?
Custom module.
class CustomModule(nn.Module):
def __init__(self, input_channels_to_this_block, output_channels_from_this_block):
super(CustomModule, self).__init__()
self.cnn2 = nn.Conv2d()
self.cnn3 = nn.Conv2d()
self.cnn4 = nn.Conv2d()
def forward(self, img_rep):
x = cnn2(img_rep)
y = cnn3(img_rep)
z = cnn4(img_rep)
# perform custom operation on x, y, and z and merge them into a single tensor called p. p is the same shape as x, y, and z.
o = torch.concat(x, y, z) # or average the 3 tensors so that we have a new tensor that is the same shape as x,y, and z.
return o
Classifier which uses the custom module.
class SomeClassifier(nn.Module):
def __init__(self):
super(SomeClassifier, self).__init__()
# all three conv blocks maintain the original image size.
self.cnn1 = nn.Conv2d()
self.custom_block = CustomModule(input_channels_to_this_block, output_channels_from_this_block)
self.fc = nn.Linear()
def forward(self, img):
img_rep = self.cnn1(img)
p = self.custom_block(img)
p = torch.flatten()
p = self.fc(p)
return p