Combine Feature Extractor Model with main Prediction Model

Hey guys,
A noob in pytorch here. My model planning for a task includes combining a feature extractor model which is a conv 1d model with multiple layers with a prediction model which is a stacked lstm layers. I am having a hard time understanding how to combine both these models while the initialization stages.
One of the sample models I checked initialized the feature extraction model in the main prediction model itself and then picked the inputs for the prediction model from the feature extractor.
Here is a noob attempt of how I am trying to do and probably not understanding if it is even right way to approach.


class PredictionModel(nn.Module):
    def __init__(self, features, targets, nlayers,convlayers, hsize_pred,hsize_fe,dropout):
        super().__init__()
        FE = FeatureExtractor(features, convlayers, hsize_fe,dropout) # can I connect my FE model from # here in the main prediction model??
        new_features = FE.size()[2] # this is of course throwing an error but I just planned to pick new features like this. is it even right approach??
        layers = []
        for _ in range(nlayers):
            if len(layers) ==0:
                layers.append(nn.LSTM(new_features,hsize_pred))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.Tanh())
                layers.append(nn.LSTM(hsize_pred,hsize_pred))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.Tanh())
                layers.append(nn.LSTM(hsize_pred,hsize_pred))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.Tanh())
            else:
                layers.append(nn.LSTM(hsize_pred,hsize_pred))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.Tanh())
                layers.append(nn.LSTM(hsize_pred,hsize_pred))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.Tanh())
                layers.append(nn.LSTM(hsize_pred,hsize_pred))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.Tanh())
        layers.append(nn.Linear(445,targets))
        self.model = nn.Sequential(*layers)

    def forward(self,x):
        self.x = x
        return self.model(self.x)

class FeatureExtractor(nn.Module):
    def __init__(self, features, convlayers, hsize_fe,dropout):
        super().__init__()
        # self.features = features
        # self.targets = targets
        # self.layers = nlayers
        layers = []
        for _ in range(convlayers):
            if len(layers) == 0:
                layers.append(nn.Conv1d(features, hsize_fe, 3))
                layers.append(nn.BatchNorm1d(hsize_fe))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.ReLU())
            else:
                layers.append(nn.Conv1d(hsize_fe, hsize_fe, 3))
                layers.append(nn.BatchNorm1d(hsize_fe))
                layers.append(nn.Dropout(dropout))
                layers.append(nn.ReLU())
        layers.append(nn.Linear(26,features))
        self.model = nn.Sequential(*layers)
        
    def forward(self,x):
        self.x = x
        return self.model(self.x)

Some questions are in the code comment section.
How to proceed further in this quest.
Any help is greatly helpful for me.

Thanks

FE = FeatureExtractor(features, convlayers, hsize_fe,dropout) # can I connect my FE model from # here in the main prediction model??

The usual approach would be to use the feature extractor in the forward and pass its outputs to the classifier afterwards.

new_features = FE.size()[2] # this is of course throwing an error but I just planned to pick new features like this. is it even right approach??

Your LSTM modules expect an input of [seq_len, batch_size, features] so you would have to check how these dimensions map to the output of your feature extractor and especially how the temporal dimension should be created.
E.g. if your FeatureExtractor module returns an output of [batch_size, out_channels, height, width] you could try to use the spatial dimensions (height and width) as the temporal dimension and the out_channels as the features. Of course other approaches are possible and in the end it depends on your use case and your idea.

I would also assume that self.model will raise an error as you are wrapping LSTM modules in it which will return multiple outputs while nn.Sequential expects a single in- and output.
There are ways to expand it, but I would probably just use an nn.ModuleList and call the layers in the forward separately in a loop.

Yes, thanks you for the feature extraction approach. I planned to initiate both the models separately in the engine.py file and eventually supply inputs sequentially yo each model

And yes that list of layers did throw an error aa I missed out the two outputs given by lstm layers.
Thank you again for second approach.

Will surely implement it from here.

Regards
Nidhi