Hi everyone,
i am trying to implement a model that consists of multiple encoders and one classifier.
Therefore I already implemented an Encoder as a PyTorch Model (a Class that inherits from nn.Module). I now want to implement my “Main-Model”, i.e. a model that consists of multiple Encoders and a classifier. In order to achieve this, I implemented a second class that inherits from nn.Module and passed the encoders to it as member-variables. However, if I call my_model.to(device), the encoders won’t be passed to the GPU. I think there must be a better way of assembling multiple “submodels” to one big model.
Can you tell me, what the best practice is for my problem?
Thank you!
I think I found the solution by myself.
For everyone struggling with the same problem:
You can use ModuleList. I my example, I can just append each encoder and the classifier to the ModuleList. Using this class, my “Main-Model” is aware of its submodels and for example the number of parameters is calculated correctly. I think there is a pretty good explanation of the concept here.