How to load pre-trained model and train with new untrained one

Now what I want to achieve is, to load a pre-trained model A and then combine it with an untrained model B to form a large new model (A + B) and then train the large new model.

How can this be achieved?

Have you tried something like

import torch
import torch.nn as nn
from torchvision import models

class MyNewModel(nn.Module):
  def __init__(self):
     super(MyNewModel, self).__init__()
     self.resnet = models.resnet18(pretrained=True)
     self.out = nn.Sequential([nn.Conv2d(...) , nn.Sigmoid()])
  def forward(self,x):
     x = self.resnet(x)
     x = self.out(x)
     return  x


def train():
    model = MyNewModel()
  #model training

In this case using a pretrained resnet18

1 Like

Such a good idea. :grin:
So, we can define class Model() like:

class MyNewModel(nn.Module):
  def __init__(self):
    super(MyNewModel, self).__init__()
    self.pre_trained_model = pre_trained_model
    self.untrained_part = untrained_part
  def forward(self, x):
    x = self.pre_trained_model(x)
    x = self.untrained_part(x)

Thank you so much for your timely help. :rofl: :rofl:

1 Like