Pool crystals with scatter_reduce over a flat crystal index - #164
Open
jcwang587 wants to merge 1 commit into
Open
Pool crystals with scatter_reduce over a flat crystal index#164jcwang587 wants to merge 1 commit into
jcwang587 wants to merge 1 commit into
Conversation
collate_pool now emits crystal_atom_idx as a single (N,) LongTensor mapping each atom to its crystal, and CrystalGraphConvNet.pooling aggregates with one scatter_reduce_(reduce="mean") call instead of a Python loop over crystals (one kernel launch per crystal, 256 per batch at the default batch size). Call sites move the index to the device with a single copy instead of one transfer per crystal. This is a breaking change to the model's forward interface for code that builds batches manually; code using collate_pool with a DataLoader is unaffected. Pooled features match the previous implementation to 1e-6. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
collate_poolnow emitscrystal_atom_idxas a single(N,)LongTensor mapping each atom to its crystal, replacing the list of per-crystal index tensors.CrystalGraphConvNet.poolingaggregates with onescatter_reduce_(reduce="mean")call instead of a Python loop over crystals (previously one kernel launch per crystal - 256 per batch at the default batch size). This also removes the main obstacle totorch.compile.cgcnn_test,cgcnn_descriptor, both training loops) move the index to the device with a single copy instead of one transfer per crystal.Breaking change
The model
forwardinterface changes for code that builds batches manually; anything usingcollate_poolwith aDataLoader(the documented pattern) is unaffected.Verification
cgcnn-trsmoke run passes.Part 2/4, stacked on #163 (
perf/data-fast-graph).🤖 Generated with Claude Code