如何为PyTorch DataLoader指定包含Batch格式的类型提示？

不需要继承DataLoader的抽象类，直接利用PyTorch对泛型的支持，就能给DataLoader标注明确的batch输出类型，以下是具体实现方式：

基础用法：指定单一Batch类型 #

如果预期batch是固定格式（比如Tuple[Tensor, Tensor]或Tuple[Tensor, Tensor, CustomObject]），直接在DataLoader后加泛型参数标注即可：

示例1：Batch为Tuple[Tensor, Tensor] #

from torch.utils.data import DataLoader
from torch import Tensor
from typing import Tuple

class SomeClass:
    def some_function(self, dataloader: DataLoader[Tuple[Tensor, Tensor]]):
        for idx, batch in enumerate(dataloader):
            # 类型检查工具会识别batch为Tuple[Tensor, Tensor]
            inputs, labels = batch
            # 后续操作可直接使用，无需额外类型断言

示例2：Batch为Tuple[Tensor, Tensor, CustomObject] #

from torch.utils.data import DataLoader
from torch import Tensor
from typing import Tuple

class CustomObject:
    # 自定义类的实现逻辑
    pass

class SomeClass:
    def some_function(self, dataloader: DataLoader[Tuple[Tensor, Tensor, CustomObject]]):
        for idx, batch in enumerate(dataloader):
            inputs, labels, custom_obj = batch
            # 此时custom_obj会被识别为CustomObject类型

进阶用法：支持多种Batch类型 #

如果需要兼容多种batch格式，可结合Union类型标注：

from torch.utils.data import DataLoader
from torch import Tensor
from typing import Tuple, Union

class CustomObject:
    pass

class SomeClass:
    def some_function(self, dataloader: DataLoader[Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, CustomObject]]]):
        for idx, batch in enumerate(dataloader):
            # 需通过长度判断或类型断言区分不同batch格式
            if len(batch) == 2:
                inputs, labels = batch
                # 处理双元素batch逻辑
            elif len(batch) == 3:
                inputs, labels, custom_obj = batch
                # 处理三元素batch逻辑

说明 #

PyTorch 1.10及以上版本的DataLoader原生支持泛型类型参数，标注后mypy、PyCharm等类型检查工具会自动识别batch类型，帮助提前发现类型错误，提升代码可读性和可靠性。

内容的提问来源于stack exchange，提问作者Robin van Hoorn