-
Notifications
You must be signed in to change notification settings - Fork 1.7k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
GNNModular #3339
GNNModular #3339
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good start! Mostly requests for more documentation below
deepchem/models/torch_models/gnn.py
Outdated
The number of negative edges added is equal to half the number of edges in the original graph. This is useful for tasks like edge prediction, where the model needs to learn to differentiate between existing and non-existing edges. | ||
""" | ||
|
||
def __call__(self, data): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you add a docstring?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
done
@@ -84,7 +84,7 @@ def __init__(self, model: nn.Module, components: dict, **kwargs): | |||
super().__init__(self.model, self.loss_func, **kwargs) | |||
self.model.to(self.device) | |||
self.components = { | |||
k: v.to(self.device) for k, v in self.components.items() | |||
k: v.to(self.device) for k, v in self.components.items() if isinstance(v, nn.Module) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When would we expect to have a non nn.module
in self.components?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In this GNNModular class. The pooling components can be either a function or parameterized nn.module.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One minor refactoring request below (self.JK
-> self.jump_knowledge
). Once that's done and a requested comment below is added this should be ready to merge in
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, feel free to merge once CI is clear
This PR adds GNNModular to torch_models. GNNModular allows users to define graph convolutional encoders, prediction heads, and specific tasks such as node prediction, edge prediction, regression, classification, etc. This PR implements edge prediction only, using the GINEConv layer from torch geometric. as the GCN
Type of change
Please check the option that is related to your PR.
Checklist
yapf -i <modified file>
and check no errors (yapf version must be 0.32.0)mypy -p deepchem
and check no errorsflake8 <modified file> --count
and check no errorspython -m doctest <modified file>
and check no errors