Should I initialize some submodules beforehand and pass them in as parameters to the final model's __init__ function?

1

Hello. I have a question as follows:

With a large model, if I initialize all the child modules in the init function of a final model I may end up having to pass in many parameters. Can I build a build_model() function like the example below?

import torch


class A(torch.nn.Module):
    def __init__(
        self,
        idim: int,
        odim: int,
        sub_1: torch.nn.Module,
        sub_2: torch.nn.Module,
    ):
        super().__init__()
        self.idim = idim
        self.odim = odim
        self.sub_1 = sub_1
        self.sub_2 = sub_2
        self.linear_out = torch.nn.Linear(self.idim, self.odim)
    

def build_model(
    idim: int,
    odim: int,
    sub1_param1: int,
    sub1_param2: int,
    sub2_param1: int,
    sub2_param2: int,
) -> A:
    sub_1 = torch.nn.Linear(sub1_param1, sub1_param2)
    sub_2 = torch.nn.Linear(sub2_param1, sub2_param2)
    
    return A(idim, odim, sub_1, sub_2)

Here is a simple descriptive example. For me there could be more submodules and each submodule has more parameters. Are there any disadvantages to this implementation? Thank you.

2

Your approach is fine and should not cause any issues. I would claim it comes down to your programming style and the use case.