search

divelab / DIG

A library for graph deep learning research
https://diveintographs.readthedocs.io/
GNU General Public License v3.0
1.84k stars 281 forks source link

SchNet: How to get graph embedding instead of prediction #122

Closed LanceKnight closed 2 years ago

LanceKnight commented 2 years ago

Hello, it seems the output of SchNet is of dimension 1, which is the prediction dimension. Is it possible to change to output the graph embedding instead of the predicton? This way it's more flexible for the users to put whatever downstream classifier after the embedding and have a control of how many prediction dimension there can be (multi-task potentially)

limei0307 commented 2 years ago

Hi @LanceKnight,

Thanks for your suggestion. We exactly followed the SchNet paper and fixed the output dimension as 1. I updated our code here so that users can define the output dimension.

Thanks, Limei

LanceKnight commented 2 years ago

Thank you! @limei0307