Get classification and next-to-last layer with 1 model


I am trying to use a pre-trained resnet model to classify images, and I also want to get the next-to-last layer features.

From this discussion I see that it’s pretty easy to construct a new classifier that returns the next-to-last layer weights.

Currently my code looks like this:

class_model = resnet152(pretrained=True)

lw_model = torch.nn.Sequential(*list(class_model.children())[:-1])

#im_tensor is a normalized JPEG
class_out = class_model(im_tensor.unsqueeze(0))
lw_out = lw_model(im_tensor.unsqueeze(0))

top5 = []
for idx in class_out[0].sort()[1][-5:]:
    top5.append((idx.item(), labels[(str(idx.item()))][1]))

_ = [print(c) for i,c in reversed(top5)]
# print(lw_out.squeeze_())

Is there any way to get a classification and the next to last layer weights using just one model instead of running two? Perhaps toggle something in the pretrained resnet to have it emit the last weights in addition to the classifications?

Any help would be greatly appreciated, thanks!

Why don’t you return two values from the forward function, the penultimate layer predictions as well as the final logits?

something like

lw_model = torch.nn.Sequential(*list(class_model.children())[:-1])
final_layer = list(class_model.children()))[-1]

and chain these two models to make the final prediction.

1 Like

Thanks! It hadn’t occurred to me that you could break up classifiers into chunks and compose them that way. Solved my issue.