Skip to content

Memory-mapped caching for image translation training #218

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.

Sign up for GitHub

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

Jump to bottom
Merged
merged 86 commits into from
Mar 27, 2025
Merged

Memory-mapped caching for image translation training #218

merged 86 commits into from
Mar 27, 2025

Conversation

ziw-liu
Copy link
Collaborator

@ziw-liu ziw-liu commented Jan 6, 2025

#195 implemented in-RAM caching for image translation trainings. However it does not scale to datasets larger than available system memory. This PR implements a node-local disk cache via tensordict's memory-mapped tensor.

edyoshikun and others added 30 commits October 21, 2024 16:28
@edyoshikun @ziw-liu
caching dataloader
e19ee14
@edyoshikun @ziw-liu
caching data module
d31978d
@edyoshikun @ziw-liu
black
041d738
@edyoshikun @ziw-liu
ruff
7f76174
@edyoshikun @ziw-liu
Bump torch to 2.4.1 (#174)
85ea791 
* update torch >2.4.1

* black

* ruff
@edyoshikun @ziw-liu
adding timeout to ram_dataloader
1838581
@edyoshikun @ziw-liu
bandaid to cached dataloader
f5c01a3
@edyoshikun @ziw-liu
fixing the dataloader using torch collate_fn
26a06b8
@edyoshikun @ziw-liu
replacing dictionary with single array
f2ff43c
@edyoshikun @ziw-liu
loading prior to epoch 0
5fb96d7
@edyoshikun @ziw-liu
Revert "replacing dictionary with single array"
848cd63 
This reverts commit 8c13f49.
@edyoshikun @ziw-liu
using multiprocessing manager
f7e57ae
@ziw-liu
add sharded distributed sampler
c4797b4
@ziw-liu
add example script for ddp caching
2c31e7d
@ziw-liu
format and lint
5300b4a
@edyoshikun
addding the custom distrb sampler to hcs_ram.py
8a8b4b0
@edyoshikun
adding sampler to val train dataloader
49764fa
@ziw-liu
fix divisibility of the last shard
1fe5491
@ziw-liu
hcs_ram format and lint
0b005cf
@ziw-liu
data module that only crops and does not collate
023ca88
@ziw-liu
wip: execute transforms on the GPU
f7ce0ba
@edyoshikun
path for if not ddp
daa6860
@ziw-liu
fix randomness in inversion transform
55499de
@ziw-liu
add option to pop the normalization metadata
4280677
@ziw-liu
move gpu transform definition back to data module
1561802
@ziw-liu
add tiled crop transform for validation
2e37217
@ziw-liu
add stack channel transform for gpu augmentation
7edf36e
@ziw-liu
fix typing
eda5d1b
@ziw-liu
collate before sending to gpu
550101d
@ziw-liu
inherit gpu transforms for livecell dataset
92e3722
@edyoshikun edyoshikun self-requested a review March 25, 2025 00:44
@edyoshikun
Copy link
Contributor

Marking this ready for review. Taking a look now! @ziw-liu

@edyoshikun edyoshikun marked this pull request as ready for review March 25, 2025 00:46
@edyoshikun edyoshikun mentioned this pull request Mar 25, 2025
Merged
Base automatically changed from segmentation-module to main March 25, 2025 04:16
@ziw-liu ziw-liu added this to the v0.3.0 milestone Mar 25, 2025
@ziw-liu ziw-liu requested a review from edyoshikun March 26, 2025 17:21
edyoshikun
edyoshikun approved these changes Mar 27, 2025
View reviewed changes
Copy link
Contributor

@edyoshikun edyoshikun left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@ziw-liu ziw-liu merged commit 89b2917 into main Mar 27, 2025
4 checks passed
@ziw-liu ziw-liu deleted the mmap-cached branch March 27, 2025 17:46
This was referenced Mar 27, 2025
Open
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@edyoshikun edyoshikun edyoshikun approved these changes

Assignees
No one assigned
Labels
enhancement New feature or request translation Image translation (VS)
Projects
None yet
Milestone
v0.3.0
Development

Successfully merging this pull request may close these issues.
2 participants
@ziw-liu @edyoshikun