47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
"""Data-related utility functions used by the trainer main loop."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
|
|
def move_batch_to_device(
|
|
batch: Any,
|
|
device: torch.device,
|
|
*,
|
|
non_blocking: bool = True,
|
|
) -> Any:
|
|
"""Move every tensor in ``batch`` to ``device``, recursing into dict/list.
|
|
|
|
Returns a new structure with the same shape; the input ``batch`` is not
|
|
mutated. dict / list containers are rebuilt; tensors are moved via
|
|
``.to(device, non_blocking=...)``; everything else is passed through
|
|
by reference.
|
|
|
|
Parameters
|
|
----------
|
|
batch : Any
|
|
Typically a dict produced by the dataloader, but recursion accepts
|
|
dict / list / tensor / arbitrary leaves.
|
|
device : torch.device
|
|
Target device, typically ``self.device`` on the trainer.
|
|
non_blocking : bool
|
|
Whether to use pinned-memory async copy. Default True because the
|
|
trainer uses pinned loaders; pass False if the dataloader hasn't
|
|
pinned memory.
|
|
"""
|
|
if isinstance(batch, dict):
|
|
return {
|
|
k: move_batch_to_device(v, device, non_blocking=non_blocking)
|
|
for k, v in batch.items()
|
|
}
|
|
if isinstance(batch, list):
|
|
return [
|
|
move_batch_to_device(v, device, non_blocking=non_blocking) for v in batch
|
|
]
|
|
if isinstance(batch, torch.Tensor):
|
|
return batch.to(device, non_blocking=non_blocking)
|
|
return batch
|