I am trying to create a general feedforward net with CNN and FNN layers. What I am trying to do is: if some conv_layer_params are passed in the class, then I want to create conv layers according to how config is specified.
Would this be a good way of doing it inside a class EncodingNetwork(nn.Module):?
layers = []
if conv_layer_params:
conv_layer_type = nn.Conv2d
for config in conv_layer_params:
if len(config) == 5:
(in_channels, out_channels, kernel_size, stride, dilation_rate) = config
elif len(config) == 4:
(in_channels, out_channels, kernel_size, stride) = config
dilation_rate = (1, 1)
layers.append(
conv_layer_type(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
dilation_rate=dilation_rate,
activation=activation_fn,
kernel_initializer=kernel_initializer,
dtype=dtype))
layers.append(torch.flatten())