How to convert this pytorch neural net into nn.Sequential format?

1

Hi how’s it going? I have a simple neural net that takes in a convolutional layer, flattens the output, concats another observation to the output, and then passes it through two linear layers. I know how to do everything in nn.Sequential except for the concatenation. How would I structure this?

Here is the original neural net:

        self.conv1 = nn.Conv2d(in_channels=in_dim, out_channels=1, kernel_size=(1, 1), stride=1)
        self.fc1 = nn.Linear(flattened_dim+obs_dim,20)
        self.fc2 = nn.Linear(20,output_dim)

        def network(x,obs):
            x = self.conv1(x)
            x = torch.flatten(x)
            x = torch.cat((x,obs))
            x = self.fc1(x)
            x = self.fc2(x)

            return x

And I’d like to convert it to something like this:

        self.network = nn.Sequential(
                        nn.Conv2d(in_channels=in_dim, out_channels=1, kernel_size=(1, 1), stride=1),
                        nn.Flatten(start_dim=0),

                       '''Here I need to concatenate an observation with 
                         the output of nn.Flatten before the next layer'''

                        nn.Linear(flattened_dim+obs_dim,20),
                        nn.Linear(20,output_dim),
                    )

How would you go about implementing this?

2

While there might be a way to write a custom nn.Module accepting two inputs and somehow make sure the nn.Sequential contain won’t break, I think the best approach would still be to use your original model, which uses torch.cat in its forward.
nn.Sequential containers are used for simple sequential models, while custom modules allow more flexibility, which seems to perfectly fit your use case.
What’s the reason you want to transform it into an nn.Sequential?

1 Like