Add softmax layer to Resnet50

Hi there,

I would like to add a Softmax layer to the end of a custom Resnet50 that i have trained. I wana add it for validation. My ResNet50 is:

    model = torchvision.models.resnet50(pretrained=True)
    model.fc = nn.Sequential(
        nn.Linear(model.fc.in_features, 1000),
        nn.Linear(1000, 128),
        nn.Linear(128, 64),
        nn.Linear(64, 8)
    )

How could I do it?

Thanks

1 Like

You want to get your trained model without softmax and feed it an image and get its output and apply softmax on the result, right? if so, you can use sth like: nn.Softmax()(model(image)) or F.softmax(model(image)). it will give you the probability for 8 classes.

1 Like

You can use nn.Sequential to add your layer.

model.fc = nn.Sequential(
    *model.fc,
    nn.Softmax(),
)
1 Like