pytorch - 💡(How to fix) Fix Improve static type hints for dataloader with known `collate_fn` [1 comments, 1 participants]

🚀 The feature, motivation and pitch

Feature Requested

I want to improve the type hints that are provided to developers when they are training their models, by enabling the IDE to know the types of the data that the Dataloader will yield.

In particular, in code bases where multiple models share a single dataset type, or inversely where multiple datasets must be compatible with one model, this intermediate step would enable one to automatically validate a training loop statically, which is much better than the sometimes grueling work of booting a training server, witnessing it fail because the yielded type is wrong, and then having to modify it and try again (especially with the large model size we nowadays trend towards)

While the code states that there is no way to set a default value if the user doesn't pass in a custom collate_fn, this is actually possible with overloads in a .pyi file.

Alternatives

Had Python had a type intersection, this would've been my first choice for my own projects. Alas, that is not the case.

Otherwise, users might decide to wrap the dataloader initialization and hide the fact that they are handling a dataloader, for example:

DataType = TypeVar("DataType", covariant=True)
Collated = TypeVar("Collated", covariant=True)

def init_dataloader(dataset: Dataset[DataType], collate_fn: Callable[[list[DataType]], Collated], **kwargs: Any) -> Iterable[Collated]:
    return DataLoader(dataset, collate_fn=collate_fn, **kwargs)

Which I hope is of a much lesser interest (since we completely loose the information of the dataloader being a dataloader.

Additional context

Implementation-wise, this feature is quite easy to make: it would involve adding a TypeVar (_T_co_res in the examples below), and using it:

_collate_fn_t = Callable[[list[_T]], _T_co_res]
class DataLoader(Generic[_T_co, _T_co_res]):
    def __init__(self, ..., collate_fn: _collate_fn_t[_T_co, _T_co_res] | None = None, ...):
        ...
    def __iter__(self) -> _BaseDataLoaderIter[_T_co_res]:
        ...

class _BaseDataLoaderIter(Generic[_T_co_res]):
    def __next__(self) -> _T_co_res:
        ....

In the dataloader.pyi file:

...
class DataLoader(Generic[_T_co, _T_co_res]):
    @overload
    def __init__(self, ..., collate_fn: _collate_fn_t[_T_co, _T_co_res], ...): ...
    @overload
    def __init__(self: "DataLoader[_T_co, Any]", ..., collate_fn: _collate_fn_t[_T_co, Any] | None = None, ...): ...
...

Note that you could define some more default overloads in the .pyi for common dataset types (tuple[Tensor, Tensor], etc.), but I think this is out of scope