I have a PyTorch architecture similar to this example
class MainModel(nn.Module):
def __init__(self):
super().__init__()
self.sub_model_1 = SubModel1()
self.sub_model_2 = SubModel2()
def forward(self, features1, features2):
target_1 = self.sub_model_1(features1)
target_2 = self.sub_model_2(target_1, features2)
return dict(
target_1=target_1,
target_2=target_2,
)
class SubModel1(nn.Module):
def __init__(self):
super().__init__()
self.layer = nn.Linear(2, 1)
def forward(self, features1):
return self.layer(features1)
class SubModel2(nn.Module):
def __init__(self):
super().__init__()
self.layer = nn.Linear(3, 1)
def forward(self, target_1, features2):
concatenated_features = torch.cat([target_1, features2], dim=1)
target_2 = self.layer(concatenated_features)
return target_2
I would like to train the first network (e.g. SubModel1) by providing target_1_hat and once that network os trained, I would like to fix to freeze the parameters of the model and try to optimizer the model SubModel2 by providing target_2_hat
Does PyTorch or a similar library has the option of training such models?