Access TwoMLPHead and FastRCNNPredictor separately from saved model after training

vision
1

I am training a FasterRCNN model using VGG16 as backbone network.

I save the model with torch.save(‘model.pt’)

Then I can load the model model = torch.load(‘model.pt’)

Now, I want to access some functions I added in the TwoMLPhead class and Fast RCNNPredictor class. But I am not able to access them, as I get the error. 'FasterRCNN ’ has no attribute ‘function_name’

How can I use the function from the class?

class TwoMLPHead(nn.Module):

"""

Standard heads for FPN-based models

Arguments:

    in_channels (int): number of input channels

    representation_size (int): size of the intermediate representation

"""

def __init__(self, in_channels, representation_size):

    super(TwoMLPHead, self).__init__()

    self.fc6 = nn.Linear(in_channels, representation_size)

    self.fc7 = nn.Linear(representation_size, representation_size)

def forward(self, x):

    x = x.flatten(start_dim=1)

    x = F.relu(self.fc6(x))

    x = F.relu(self.fc7(x))

    return x

def extract(self, x, verbose=False):

    out1 = x.flatten(start_dim=1)

    out2 = F.relu(self.fc6(out1))

    out3 = F.relu(self.fc7(out2))

    if verbose == True:

        print(str(out1.shape) + ' ' + str(np.prod(out0.shape[1:])))

        print(str(out2.shape) + ' ' + str(np.prod(out1.shape[1:])))

        print(str(out3.shape) + ' ' + str(np.prod(out2.shape[1:])))

    return [out1, out2, out3]

extract() is added to classes. I wan to access the function from saved model.