how to convert nn.Sequential code to nn.Linear

Question

I am new to deep learning, and I came across this while learning.

Is there a way to convert the nn.Sequential() functions here to nn.Linear(), because of how flexible it is to use nn.Linear() functions.

class FashionMNISTModelV2(nn.Module):
def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
    super().__init__()
    self.block_1 = nn.Sequential(
        nn.Conv2d(in_channels=input_shape, 
                  out_channels=hidden_units, 
                  kernel_size=3,
                  stride=1,
                  padding=1),
        nn.ReLU(),
        nn.Conv2d(in_channels=hidden_units, 
                  out_channels=hidden_units,
                  kernel_size=3,
                  stride=1,
                  padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2,
                     stride=2)
    )
    self.block_2 = nn.Sequential(
        nn.Conv2d(hidden_units, hidden_units, 3, padding=1),
        nn.ReLU(),
        nn.Conv2d(hidden_units, hidden_units, 3, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2)
    )
    self.classifier = nn.Sequential(
        nn.Flatten(),
        nn.Linear(in_features=hidden_units*7*7, 
                  out_features=output_shape)
    )

def forward(self, x: torch.Tensor):
    x = self.block_1(x)
    x = self.block_2(x)
    x = self.classifier(x)
    return x

Answer 1

You can define your linear layer separately from the classifier as a standalone layer:

self.linear = nn.Linear(in_features=hidden_units*7*7, 
                        out_features=output_shape))

Then in the forward function, the equivalent implementation would be:

def forward(self, x: torch.Tensor):
    x = self.block_1(x)
    x = self.block_2(x)
    x = x.flatten(1)
    x = self.linear(x)
    return x

Answer 2

Try this:

self.flatten = nn.Flatten()
self.linear1 = nn.Linear(in_features=hidden_units*7*7, out_features=output_shape))

def forward(self, x:torch.Tensor): 
x = self.block_1(x)
x = self.block_2(x)
x = self.flatten(x)
x = self.linear1(x)
return x

If you don't want to add flatten layer to your model, you can simply do it in forward function:

def forward(self, x:torch.Tensor):
x = self.block_1(x)
x = self.block_2(x)
x = x.view(x.size(0), -1)
x = self.linear1(x)
return x