Creating model from pretrained model and keep its weights

vision
1

Hi, I’m new to PyTorch so excuse me beforehand for my rookie mistakes.

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(256, 1, 1)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(14*14, 1)

    def forward(self, x):
        output1 = F.sigmoid(self.conv1(x))
        output2 = F.sigmoid(self.fc1(output1))
        return output1, output2


class Final(nn.Module):
    def __init__(self):
        super(Final, self).__init__()
        pretrained_model = torch.hub.load('pytorch/vision:v0.6.0', 'densenet201', pretrained=True)
        model = MyModel()
        list_of_layers = list(pretrained_model.features)[:8]
        list_of_layers.extend(list(model.children()))
        self.final_model = nn.Sequential(*list_of_layers)

    def forward(self, x):
        outputs = list()
        for ii, model in enumerate(self.final_model):
            x = model(x)
            if ii == 8 or ii == 10:
                outputs.append(x)
        return outputs

final = Final()
optimizer = optim.Adam(final.parameters())

This is the way I managed to create my desired model at last. Pay attention that I couldn’t simply create a sequential model because I want to have two outputs and it’s not possible in sequential models as far as I know.

  1. Any comment on how to make it more simpler if possible?
  2. (main question) If I train this model do I start with random weights for every layer or my dense section will have Its own pretrained weights as the starting point of training?