Add ```dataclass``` support in ```move_data_to_device```

[dalek-who]

🚀 Feature

Motivation

NamedTuple is often used to wrap batch data, but if a class inherit NamedTuple, it can't be inherited again, so it's difficult to add some new fields. A better choice is dataclass, but it's not supported by move_data_to_device

Pitch

Alternatives

Additional context

from pytorch_lightning.utilities import move_data_to_device
import torch
from dataclasses import dataclass
from typing import NamedTuple, List

class Data_1(NamedTuple):
    example_id: List[str]
    x: torch.Tensor
    y: torch.Tensor

@dataclass
class Data_2:
    example_id: List[str]
    x: torch.Tensor
    y: torch.Tensor

batch_size = 5

batch_1 = Data_1(
    example_id=[f"e-{i}" for i in range(batch_size)],
    x = torch.rand(batch_size),
    y = torch.rand(batch_size),
)

batch_2 = Data_2(
    example_id=[f"e-{i}" for i in range(batch_size)],
    x = torch.rand(batch_size),
    y = torch.rand(batch_size),
)

device = torch.device("cuda:0")

move_data_to_device(batch=batch_1, device=device)
# Data_1(
# example_id=['e-0', 'e-1', 'e-2', 'e-3', 'e-4'], 
# x=tensor([0.3385, 0.6415, 0.8117, 0.6030, 0.2551], device='cuda:0'), 
# y=tensor([0.2586, 0.8260, 0.0066, 0.0321, 0.3881], device='cuda:0')
# )

move_data_to_device(batch=batch_2, device=device)
# Data_2(
# example_id=['e-0', 'e-1', 'e-2', 'e-3', 'e-4'], 
# x=tensor([0.9459, 0.0063, 0.3763, 0.4537, 0.6941]), 
# y=tensor([0.8080, 0.8041, 0.2999, 0.5154, 0.3706])
# )

[tchaton]

Dear @dalek-who,

Great suggestion. Would you mind making a PR to add support for it ?

Best,
T.C

[dalek-who]

@tchaton
I'm note sure whether I should change apply_to_collection to support it, here's my implementation:

def apply_to_collection(
    data: Any,
    dtype: Union[type, tuple],
    function: Callable,
    *args,
    wrong_dtype: Optional[Union[type, tuple]] = None,
    **kwargs
) -> Any:
    """
    Recursively applies a function to all elements of a certain dtype.

    Args:
        data: the collection to apply the function to
        dtype: the given function will be applied to all elements of this dtype
        function: the function to apply
        *args: positional arguments (will be forwarded to calls of ``function``)
        wrong_dtype: the given function won't be applied if this type is specified and the given collections is of
            the :attr:`wrong_type` even if it is of type :attr`dtype`
        **kwargs: keyword arguments (will be forwarded to calls of ``function``)

    Returns:
        the resulting collection
    """
    elem_type = type(data)

    # Breaking condition
    if isinstance(data, dtype) and (wrong_dtype is None or not isinstance(data, wrong_dtype)):
        return function(data, *args, **kwargs)

    # Recursively apply to collection items
    if isinstance(data, Mapping):
        return elem_type({
            k: apply_to_collection(v, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs)
            for k, v in data.items()
        })

    if isinstance(data, tuple) and hasattr(data, '_fields'):  # named tuple
        return elem_type(
            *(apply_to_collection(d, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs) for d in data)
        )

    if isinstance(data, Sequence) and not isinstance(data, str):
        return elem_type([
            apply_to_collection(d, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs) for d in data
        ])

    #############################
    # my pitch:
    if dataclasses.is_dataclass(data) and not isinstance(data, type):
        return elem_type(**{
            field: apply_to_collection(getattr(data, field), dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs)
            for field in data.__dataclass_fields__
        })
    #############################

    # data is neither of dtype, nor a collection
    return data

[tchaton]

Yes, looks great ! Please open a PR and add a test.

[dalek-who]

OK, I'll have a try!

[dalek-who]

@tchaton I've finished the change and ready to make a pr, but how can I update the changelog to include the pr #xxx id itself?

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!