How can I concatenate two model then save it separately?

1

I defining two model which are encoder and classfication respectively. I want use encoder’s ouput as the input of the classfication then train the two model. After training I want to save the two model separately.

here’s my two model

# encoder model 
class source(nn.Module):
    def __init__(self, initial_channel):
        super(source, self).__init__()
        self.initial_channel = initial_channel


        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True)
            )
        self.model = self.model = nn.Sequential(conv_bn(initial_channel, 64, 2),
                                   conv_bn(64, 128, 2),
                                   conv_bn(128, 128, 1),
                                   conv_bn(128, 256, 2),
                                   conv_bn(256, 256, 1),
                                   nn.AdaptiveAvgPool2d([1,1]))
    def forward(self, x):
        x = self.model(x)
        x = x.view(-1,256)
        return x

# classfication model
class classfication(nn.Module):
    def __init__(self, class_num):
        super(classfication, self).__init__()
        self.class_num = class_num

        self.fc1 = nn.Sequential(nn.Linear(256, 64),
                                 nn.ReLU(inplace=True),
                                 nn.Dropout(0.5))
        self.fc2 = nn.Sequential(nn.Linear(64, class_num), )

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

And following my main model to describe what I want to do in detail.

encoder = source(initial_channel)
cls = classfication(classnum)

model = torch.nn.Sequential(encoder, cls)
train(model)
save(encoder)
save(cls)
2

You could store the state_dicts of both models separately.
Note that it’s recommended to save the state_dict rather than the complete model as described here.

3

You can do as as @ptrblck mentioned. Also you can just save their states_dict in a single file as-

torch.save({'encoder': encoder.state_dict(), 'cls': cls.state_dict()}, 'both_models_state_dict.pt')

And then load it later as-

both_models_state_dict = torch.load('both_models_state_dict.pt')
encoder.load_state_dict(both_models_state_dict['encoder'])
cls.load_state_dict(both_models_state_dict['cls'])