def __init__(self, in_channels, out_channels, batch_norm=False):
"""Constructor method."""
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.batch_norm = batch_norm
self.conv = conv3x3(self.in_channels, self.out_channels)
if self.batch_norm:
self.bn = nn.BatchNorm2d(self.out_channels)
self.relu = nn.ReLU()
def __init__(self, num_filters, num_blocks=4, batch_norm=False):
"""Constructor method."""
super().__init__()
self.num_filters = num_filters
self.num_blocks = num_blocks
self.batch_norm = batch_norm
for i in range(num_blocks):
self.add_module(
f"convblock_{i+1}",
ConvBlock(self.num_filters * 2**i, batch_norm=self.batch_norm))
self.add_module(
f"conv2d_proj_{i+1}",
conv3x3(self.num_filters * 2**i,
self.num_filters * 2**(i + 1),
stride=(2, 2)))
self.add_module(
"bottleneck",
ConvBlock(self.num_filters * 2**self.num_blocks,
batch_norm=self.batch_norm))
def test_conv_3x3_bloc(self):
"""Test conv block"""
block = conv3x3(3, 64)
tensor = torch.randn(8, 3, 224, 224)
self.assertSequenceEqual(block(tensor).shape, (8, 64, 224, 224))