Skip to content
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

Add TorchScript support for AttentiveFP #5868

Merged
merged 4 commits into from
Nov 4, 2022

Conversation

ftxj
Copy link
Contributor

@ftxj ftxj commented Nov 1, 2022

Add TorchScript JIT support on AttentiveFP.

Now AttentiveFP can use the following code to directly get TorchScript model.

model = AttentiveFP(in_channels=..., jittable=True)
jit = torch.jit.script(model)

To do this, we first add the propagate type in GATEConv and then reorganize the ModuleList logic to satisfy TorchScript language syntax.

@codecov
Copy link

codecov bot commented Nov 1, 2022

Codecov Report

Merging #5868 (eb91d0f) into master (bc83624) will decrease coverage by 2.00%.
The diff coverage is 5.88%.

❗ Current head eb91d0f differs from pull request most recent head 0369056. Consider uploading reports for the commit 0369056 to get more accurate results

@@            Coverage Diff             @@
##           master    #5868      +/-   ##
==========================================
- Coverage   86.06%   84.05%   -2.01%     
==========================================
  Files         349      349              
  Lines       19477    19482       +5     
==========================================
- Hits        16762    16376     -386     
- Misses       2715     3106     +391     
Impacted Files Coverage Δ
torch_geometric/nn/models/attentive_fp.py 19.35% <5.88%> (-2.34%) ⬇️
torch_geometric/nn/models/dimenet_utils.py 0.00% <0.00%> (-75.52%) ⬇️
torch_geometric/nn/models/dimenet.py 14.90% <0.00%> (-52.76%) ⬇️
torch_geometric/profile/profile.py 36.73% <0.00%> (-27.56%) ⬇️
torch_geometric/nn/conv/utils/typing.py 81.25% <0.00%> (-17.50%) ⬇️
torch_geometric/nn/pool/asap.py 92.10% <0.00%> (-7.90%) ⬇️
torch_geometric/nn/inits.py 67.85% <0.00%> (-7.15%) ⬇️
torch_geometric/nn/dense/linear.py 86.66% <0.00%> (-6.67%) ⬇️
torch_geometric/transforms/add_self_loops.py 94.44% <0.00%> (-5.56%) ⬇️
torch_geometric/nn/resolver.py 88.88% <0.00%> (-5.56%) ⬇️
... and 12 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

@rusty1s rusty1s changed the title add TorchScript support on AttentiveFP Add TorchScript support for AttentiveFP Nov 1, 2022
@rusty1s rusty1s enabled auto-merge (squash) November 4, 2022 07:14
@rusty1s rusty1s merged commit 254a3a3 into pyg-team:master Nov 4, 2022
JakubPietrakIntel pushed a commit to JakubPietrakIntel/pytorch_geometric that referenced this pull request Nov 25, 2022
Add TorchScript JIT support on AttentiveFP. 

Now AttentiveFP can use the following code to directly get TorchScript
model.

`model = AttentiveFP(in_channels=..., jittable=True)`
`jit = torch.jit.script(model)` 

To do this, we first add the propagate type in GATEConv and then
reorganize the ModuleList logic to satisfy TorchScript language syntax.

Co-authored-by: jxin <jxin@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants