How can I extract the final layer (vector) of the pre-trained DimeNet?

[qikuizhu]

I would like to use the final layer of the pre-trained DimeNet as the input vector for other prediction tasks; in other words, this is a transfer learning proposed in the following paper.

https://pubs.acs.org/doi/abs/10.1021/acs.jpca.0c06231

This approach is very useful but the above authors did not use the DimeNet. I think the DimeNet is more effective for such transfer learning task.

How can I extract the final layer of the pre-trained DimeNet? Specially, for example given the following molecule data input,

F 0.015 0.06 -0.02
C -0.02 1.39 0.01
F 1.24 1.84 -0.02
F -0.64 1.82 -1.08
C -0.70 1.86 1.22
C -1.26 2.25 2.20
H -1.76 2.59 3.08

how can I obtain the the final layer of this molecule from your code?

[gasteigerjo]

You can save the activations after any layer as an object variable (self.my_favorite_activation = x), e.g. before the final layer here. You can then access them via something like dimenet_pp.output_blocks[i].my_favorite_activation. You will probably want to concatenate the embeddings from all output layers, since they are all summed up to obtain the final prediction. Also, this will give you one representation per atom. You can then sum or average them to get one embedding per molecule (like we do for the prediction here).

Hope this helps!

[qikuizhu]

Thank you very much. I will try.