Extract embeddings for clustering

[matr1x-1]

Hello everyone,

I am trying to extract the embeddings from the last layer (similar to what you did for the uncertainty quantification: https://pubs.aip.org/aip/jcp/article/158/16/164111/2886901) for clustering purposes. What would be the best way to do this?

I have tried using forward_hooks like this:

from pathlib import Path
from nequip.scripts.evaluate import _load_deployed_or_traindir
 
# Dictionary to store the outputs of layers
layer_outputs = {}
 
# Define the hook function that will store layer outputs
def hook_fn(module, input, output): # Save the output of the current layer in the dictionary
    layer_outputs[module] = output
 
model, loaded_deployed_model, model_r_max, type_names = _load_deployed_or_traindir(Path("best_model.pth"), "cuda")
 
for layer in model.modules(): # How to get the different layers (modules)
    layer.register_forward_hook(hook_fn)

However, I am not sure which layer is the correct one or if this approach makes sense at all.

Thank you for your help!