CopyToDevice transform

[rijobro]

Will enable (dictionary example given, array is the same):

transforms = Compose([
    LoadImaged(keys=["image","label"]),
    ToTensord(keys=["image","label"]),
    CopyToDeviced(keys=["image","label"], device=device),
])

for batch_data in train_loader:
    im, label = batch_data["image"], batch_data["label"]

removing the need for:

for batch_data in train_loader:
    im, label = batch_data["image"].to(device), batch_data["label"].to(device)

Status

Ready

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh --codeformat --coverage.
• Quick tests passed locally by running ./runtests.sh --quick.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[Nic-Ma]

Hi @rijobro ,

Thanks for your good idea.
I have a question In my simple mind: if moving data to GPU in a transform in every process, it will batch together in DataLoader later and the batch data is not contiguous in GPU memory, maybe the performance is worse than "moving to GPU" after batching? Just curious, maybe you can provide some test results to compare later.

Thanks.

[rijobro]

Thanks @Nic-Ma I hadn't thought of that. I'll have a look!

[rijobro]

Hi @Nic-Ma was looking into this, but it seems it's a bad idea, since each of the subprocesses from the DataLoader would have to share GPU memory across processes. See here and here.

Hence, I'll close this PR.

[wyli]

but I think copy_to_device is still a useful utility, isn't it?

[rijobro]

that functionality still exists in the lr_finder PR, I could split it into its own PR if you'd like. But I think if we're about to merge the lr_finder PR, there's not much point.