gnn-tracking / gnn_tracking

Reconstruct billions of particle trajectories with graph neural networks
https://gnn-tracking.rtfd.io/
MIT License
31 stars 16 forks source link

Improve memory consumption of condensation losses #391

Closed jnpark3 closed 1 year ago

jnpark3 commented 1 year ago

Pull request in resolving issue #293.

Background loss calculation was updated to require less memory usage while computing the indices of alpha particles. Potential loss calculation was updated by adding the '_fast_condensation_loss' function, which uses the updated alpha index search methods and radius graphs to significantly decrease the memory usage of computing condensation loss. This function uses an approximate method to calculate repulsive loss, as torch_cluster's 'radius_graph' requires an upper limit 'max_neighbors' to the number of edges constructed originating from each point. This limitation should not significantly change the characteristic of the loss function for sufficiently large values of 'max_neighbors'. When 'max_neighbors' is set to <=0 while setting up 'PotentialLoss' object, the original condensation loss function is used.

Comments added to explain in detail the workings of this function.

codecov[bot] commented 1 year ago

Codecov Report

Patch coverage: 28.94% and project coverage change: +2.89% :tada:

Comparison is base (fb2c5ff) 71.56% compared to head (f7a23ab) 74.46%. Report is 46 commits behind head on main.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #391 +/- ## ========================================== + Coverage 71.56% 74.46% +2.89% ========================================== Files 48 50 +2 Lines 3327 3340 +13 ========================================== + Hits 2381 2487 +106 + Misses 946 853 -93 ``` | [Files Changed](https://app.codecov.io/gh/gnn-tracking/gnn_tracking/pull/391?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gnn-tracking) | Coverage Δ | | |---|---|---| | [src/gnn\_tracking/metrics/losses.py](https://app.codecov.io/gh/gnn-tracking/gnn_tracking/pull/391?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gnn-tracking#diff-c3JjL2dubl90cmFja2luZy9tZXRyaWNzL2xvc3Nlcy5weQ==) | `72.07% <28.94%> (-4.65%)` | :arrow_down: | ... and [24 files with indirect coverage changes](https://app.codecov.io/gh/gnn-tracking/gnn_tracking/pull/391/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gnn-tracking)

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

jnpark3 commented 1 year ago

pre-commit.ci autofix