如何在同一模块中用别名实例化两类Python数据集类，避免重复判断？

方案一：初始化阶段绑定统一别名 #

直接在DataModule的__init__方法里根据in_memory参数，把要使用的数据集类赋值给一个统一的别名（比如self.dataset_cls），后续所有实例化数据集的地方都用这个别名，彻底避免重复的if-else判断。

示例代码：

import pytorch_lightning as pl
from torch.utils.data import DataLoader

# 将两个数据集类放在当前文件内
class InMemoryDataSet:
    def __init__(self, data_path, split):
        # 全量加载逻辑实现
        self.data = ...

class IterativeDataSet:
    def __init__(self, data_path, split):
        # 逐段加载逻辑实现
        self.data = ...

class MyDataModule(pl.LightningDataModule):
    def __init__(self, data_path, in_memory=True, batch_size=32):
        super().__init__()
        self.data_path = data_path
        self.batch_size = batch_size
        
        # 核心：根据参数绑定统一的数据集类别名
        self.dataset_cls = InMemoryDataSet if in_memory else IterativeDataSet

    def setup(self, stage=None):
        # 所有stage都直接用统一别名实例化，无需重复判断
        if stage == "fit" or stage is None:
            self.train_dataset = self.dataset_cls(self.data_path, split="train")
            self.val_dataset = self.dataset_cls(self.data_path, split="val")
        if stage == "test" or stage is None:
            self.test_dataset = self.dataset_cls(self.data_path, split="test")
        if stage == "predict":
            self.predict_dataset = self.dataset_cls(self.data_path, split="predict")

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)

    def predict_dataloader(self):
        return DataLoader(self.predict_dataset, batch_size=self.batch_size)

方案二：字典映射类（适合多类型扩展） #

如果未来可能新增更多数据集加载类型，用字典做参数到类的映射会更易扩展，可读性也更强：

示例代码：

class MyDataModule(pl.LightningDataModule):
    # 定义数据集类型与对应类的映射表
    DATASET_MAP = {
        "in_memory": InMemoryDataSet,
        "iterative": IterativeDataSet
        # 新增类型时直接在这里添加键值对即可
    }

    def __init__(self, data_path, dataset_type="in_memory", batch_size=32):
        super().__init__()
        self.data_path = data_path
        self.batch_size = batch_size
        
        # 通过映射表获取目标数据集类
        self.dataset_cls = self.DATASET_MAP[dataset_type]

    # setup及dataloader方法和方案一完全一致，无需改动

优势说明 #

• 彻底消除setup方法中重复的嵌套判断，代码更简洁清爽
• 逻辑集中在初始化阶段，后续修改或扩展数据集类型时，只需改动一处
• 统一别名的方式让代码意图更清晰，降低维护成本

内容的提问来源于stack exchange，提问作者Jannik Kühn