Get input of fully connected layer of ResNet model during runtime
Found a Solution, left it as an answer to this question down below :)
Info about the project: Classification task with 2 classes.
I am trying to get the output of the fully connected layer of my model for each image I put into the model during runtime. I plan to use them after the model is done training or testing all images to visualize with UMAP.
The model:
#Load resnet
def get_model():
    model = torchvision.models.resnet50(pretrained=True)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    return model

The relevant part of pl module:
class classifierModel(pl.LightningModule):
   def __init__(self, model):
     super().__init__()
     self.model = model
     self.learning_rate = 0.0001

def training_step(self, batch, batch_idx):
        x= batch['image']
        y = batch['targets']
        x_hat = self.model(x)
        output = nn.CrossEntropyLoss()
        loss= output(x_hat,y)
        return loss
 
def test_step(self, batch, batch_idx):
        x= batch['image']
        y = batch['targets']
        x_hat = self.model(x)

Is it possible to do this by adding a empty list to the init of the pl module and then add the output after x_hat = model(x) is executed?
How would i know if after x_hat = model(x) is executed, the out_features aren't immediatly deleted/discarded ?

I was able to do this using a forward hook on the avgpool layer and saving the output on each test_step like described here :
#Define Hook: 
def get_features(name):
    def hook(model, input, output):
        features[name] = output.detach()
    return hook

Now when I load my model, I register the hook:
#Load resnet model:
def get_model():
    model = models.resnet50(pretrained=True)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    model.avgpool.register_forward_hook(get_features('feats')) #register the hook
    return model

I did not need to change the init of the pytorch lightning model but the test step function:
FEATS = []
# placeholder for batch features
features = {}

class classifierModel(pl.LightningModule):
   def __init__(self, model):
     super().__init__()
     self.model = model
     self.learning_rate = 0.0001

def test_step(self, batch,batch_idx):
        x= batch['image']
        y = batch['targets']
        x_hat = self.model(x)
        FEATS.append(features['feats'].cpu().numpy()) #added this line to save output

now we have the output FEATS[0].shape --&gt; (16, 2048, 1, 1) which is what i wanted to get(16 is the batchsize is use).
