| |
|
|
| from collections import defaultdict |
| from dataclasses import fields, is_dataclass |
| from typing import Any, Mapping, Protocol, runtime_checkable |
|
|
| import torch |
|
|
|
|
| def _is_named_tuple(x) -> bool: |
| return isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields") |
|
|
|
|
| @runtime_checkable |
| class _CopyableData(Protocol): |
| def to(self, device: torch.device, *args: Any, **kwargs: Any): |
| """Copy data to the specified device""" |
| ... |
|
|
|
|
| def copy_data_to_device(data, device: torch.device, *args: Any, **kwargs: Any): |
| """Function that recursively copies data to a torch.device. |
| |
| Args: |
| data: The data to copy to device |
| device: The device to which the data should be copied |
| args: positional arguments that will be passed to the `to` call |
| kwargs: keyword arguments that will be passed to the `to` call |
| |
| Returns: |
| The data on the correct device |
| """ |
|
|
| if _is_named_tuple(data): |
| return type(data)( |
| **copy_data_to_device(data._asdict(), device, *args, **kwargs) |
| ) |
| elif isinstance(data, (list, tuple)): |
| return type(data)(copy_data_to_device(e, device, *args, **kwargs) for e in data) |
| elif isinstance(data, defaultdict): |
| return type(data)( |
| data.default_factory, |
| { |
| k: copy_data_to_device(v, device, *args, **kwargs) |
| for k, v in data.items() |
| }, |
| ) |
| elif isinstance(data, Mapping): |
| return type(data)( |
| { |
| k: copy_data_to_device(v, device, *args, **kwargs) |
| for k, v in data.items() |
| } |
| ) |
| elif is_dataclass(data) and not isinstance(data, type): |
| new_data_class = type(data)( |
| **{ |
| field.name: copy_data_to_device( |
| getattr(data, field.name), device, *args, **kwargs |
| ) |
| for field in fields(data) |
| if field.init |
| } |
| ) |
| for field in fields(data): |
| if not field.init: |
| setattr( |
| new_data_class, |
| field.name, |
| copy_data_to_device( |
| getattr(data, field.name), device, *args, **kwargs |
| ), |
| ) |
| return new_data_class |
| elif isinstance(data, _CopyableData): |
| return data.to(device, *args, **kwargs) |
| return data |
|
|
