Memory efficient KRR - #30
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR introduces a memory-efficient implementation of the Kernel Ridge Regression model for the LCOT pipeline. The new implementation reduces memory complexity from O(n_spheres * n_subjects^2) to O(n_subjects^2) through several optimizations including embedding aggregation, chunked computation, and CPU/GPU memory management.
Key changes:
- New
EfficientKernelRidgeRegressionmodel that aggregates sphere embeddings and uses chunked distance computation to avoid OOM errors - Replacement of the original
KernelRidgeRegressionwith the efficient version for "lcot" models - DataLoader configuration changes: increased
num_workersfrom 0 to 10 and disabledpin_memory
Reviewed Changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 20 comments.
| File | Description |
|---|---|
| src/diff_benchmark/models/efficient_lcot_model.py | New memory-efficient Kernel Ridge Regression implementation with chunked computation and optimized CPU/GPU memory management |
| src/diff_benchmark/models/model_configurations.py | Updated model factory to return EfficientKernelRidgeRegression for "lcot" models instead of the original implementation |
| src/diff_benchmark/dataloaders/dataloaders.py | Modified DataLoader settings to use 10 workers and disable pinned memory |
Comments suppressed due to low confidence (1)
src/diff_benchmark/models/efficient_lcot_model.py:3
- Import of 'tqdm' is not used.
from tqdm import tqdm
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from tqdm import tqdm | ||
|
|
||
|
|
There was a problem hiding this comment.
The tqdm import on line 3 is unused in this file. Consider removing it to keep imports clean.
| from tqdm import tqdm |
| if name == "lcot": | ||
| return KernelRidgeRegression(**config) | ||
|
|
||
| # return KernelRidgeRegression(**config) |
There was a problem hiding this comment.
[nitpick] The old implementation is commented out rather than removed. Since the new EfficientKernelRidgeRegression is intended to replace KernelRidgeRegression, consider removing the commented-out line to keep the code clean. If there's a need to preserve this for reference or rollback purposes, rely on version control history instead.
| # return KernelRidgeRegression(**config) |
| torch.cuda.empty_cache() | ||
|
|
||
| print("Prediction complete!") | ||
| return torch.cat(all_predictions, dim=0).cpu() |
There was a problem hiding this comment.
Redundant .cpu() call on line 331. The predictions are already moved to CPU on line 321 before being appended to all_predictions, so calling .cpu() again after concatenation has no effect and can be removed.
| return torch.cat(all_predictions, dim=0).cpu() | |
| return torch.cat(all_predictions, dim=0) |
| if embeddings.dim() == 5: | ||
| embeddings = embeddings.squeeze(1) |
There was a problem hiding this comment.
Using .squeeze(1) without checking which dimension equals 1 could remove the wrong dimension if the batch size is 1. For example, if embeddings has shape (1, n_subjects, n_spheres, n_bvals, d) with a batch size of 1, .squeeze(1) would remove dimension 1 (n_subjects) instead of dimension 0 (batch). Consider using .squeeze(0) to specifically remove the batch dimension, or use more explicit indexing like embeddings[0] if you expect the batch dimension to always be at position 0.
| if embeddings.dim() == 5: | ||
| embeddings = embeddings.squeeze(1) | ||
| if power.dim() == 4: | ||
| power = power.squeeze(1) |
There was a problem hiding this comment.
Using .squeeze(1) without checking which dimension equals 1 could remove the wrong dimension if the batch size is 1. Consider using .squeeze(0) to specifically remove the batch dimension, or use more explicit indexing like power[0] if you expect the batch dimension to always be at position 0.
| power = power.squeeze(1) | |
| power = power.squeeze(0) |
| Returns: | ||
| aggregated: shape (n_subjects, n_bvals, d) - one embedding per subject per bval | ||
| """ | ||
| n_subjects, n_spheres, n_bvals, d = embeddings.shape |
There was a problem hiding this comment.
Variable n_subjects is not used.
| n_subjects, n_spheres, n_bvals, d = embeddings.shape | |
| _, n_spheres, n_bvals, d = embeddings.shape |
| Returns: | ||
| aggregated: shape (n_subjects, n_bvals, d) - one embedding per subject per bval | ||
| """ | ||
| n_subjects, n_spheres, n_bvals, d = embeddings.shape |
There was a problem hiding this comment.
Variable n_spheres is not used.
| n_subjects, n_spheres, n_bvals, d = embeddings.shape | |
| n_subjects, _, n_bvals, d = embeddings.shape |
| Returns: | ||
| aggregated: shape (n_subjects, n_bvals, d) - one embedding per subject per bval | ||
| """ | ||
| n_subjects, n_spheres, n_bvals, d = embeddings.shape |
There was a problem hiding this comment.
Variable n_bvals is not used.
| n_subjects, n_spheres, n_bvals, d = embeddings.shape | |
| n_subjects, n_spheres, _, d = embeddings.shape |
| Returns: | ||
| aggregated: shape (n_subjects, n_bvals, d) - one embedding per subject per bval | ||
| """ | ||
| n_subjects, n_spheres, n_bvals, d = embeddings.shape |
There was a problem hiding this comment.
Variable d is not used.
| n_subjects, n_spheres, n_bvals, d = embeddings.shape | |
| n_subjects, n_spheres, n_bvals, _ = embeddings.shape |
|
|
||
| n = emb1.shape[0] | ||
| m = emb2.shape[0] | ||
| n_bvals = emb1.shape[1] |
There was a problem hiding this comment.
Variable n_bvals is not used.
| n_bvals = emb1.shape[1] |
…o wip_efficient_krr
…debug Wip efficient krr merge debug
No description provided.