How to write different layers setting in Sequential?

1

Hello all, I have a sequential layer like

self.conv_norm_relu = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True))

I want to change nn.BatchNorm2d to GroupNorm, so the new one is

self.conv_norm_relu = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.GroupNorm(16, out_channels),
            nn.ReLU(inplace=True))

So, I want to make flexible way that allows to select between GroupNorm and BatchNorm2d. A list of normalized method is normalize_method =['GroupNorm'. 'BatchNorm2d']. If I select normalize_method[0] then self.conv_norm_relu will use GroupNorm, and If I select normalize_method[1] then self.conv_norm_relu will use BatchNorm2d

normalize_method =['GroupNorm'. 'BatchNorm2d']
self.conv_norm_relu = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            #something change here,
            nn.ReLU(inplace=True))

Do we have any option to do it?

2

Does “A if expression else B” work?

3

Not good option. if make the code too long

4

What about defining a method which returns the desired layer based on your keyword:

def get_norm_layer(out_channels, layer_type='BatchNorm'):
    if layer_type == 'BatchNorm':
        return nn.BatchNorm2d(out_channels)
    elif layer_type == 'GroupNorm':
        return nn.GroupNorm(16, out_channels)

conv_norm_relu = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            get_norm_layer(out_channels, 'GroupNorm'),
            nn.ReLU(inplace=True))
3 Likes