PyTorch Lightning中stage参数何时取值为None？

在PyTorch Lightning的DataModule示例代码中，setup方法的stage参数被定义为Optional[str] = None，示例代码如下：

import pytorch_lightning as pl
from torch.utils.data import random_split, DataLoader
from typing import Optional  # 注意：原示例遗漏该导入，实际使用需添加

# Note - you must have torchvision installed for this example
from torchvision.datasets import MNIST
from torchvision import transforms


class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir: str = "./"):
        super().__init__()
        self.data_dir = data_dir
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage: Optional[str] = None):
        # Assign train/val datasets for use in dataloaders
        if stage == "fit" or stage is None:
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        # Assign test dataset for use in dataloader(s)
        if stage == "test" or stage is None:
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

        if stage == "predict" or stage is None:
            self.mnist_predict = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32)

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

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

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

请问该stage参数何时会取值为None？

stage参数取None的场景主要有以下几种：

• 手动调用setup方法时不传参：如果在代码中直接调用datamodule.setup()，不指定stage参数，就会触发默认值None。这种情况下，setup方法里所有不同阶段的数据集初始化逻辑都会执行，一次性完成训练、验证、测试、预测数据集的准备。
• 多阶段连续执行场景：当使用Trainer连续执行多个不同阶段（比如在同一个脚本里先执行fit训练，紧接着执行test测试），Lightning会在首次调用setup时传入None，确保所有必要的数据集都被初始化，避免重复执行setup逻辑。
• 自定义扩展场景：如果基于Lightning实现了自定义运行逻辑，且没有明确指定stage参数调用setup，也会出现stage为None的情况。

内容的提问来源于stack exchange，提问作者Gulzar