如何全局设置PyTorch的train/eval模式而非针对单个nn.Module

PyTorch本身没有原生的会话级全局开关，能直接让所有nn.Module实例自动切换train/eval模式，但可以通过几种变通方法实现类似批量控制的效果：

方法1：自定义全局模式基类 #

创建一个继承自nn.Module的基类，内置全局状态同步逻辑，所有自定义模型都继承该基类：

import torch.nn as nn

_global_train_state = True

class GlobalSyncModule(nn.Module):
    def train(self, mode=True):
        global _global_train_state
        _global_train_state = mode
        # 递归同步所有子模块状态
        for module in self.children():
            if isinstance(module, GlobalSyncModule):
                module.train(mode)
        return self
    
    def eval(self):
        return self.train(False)
    
    def forward(self, x):
        # 子模型需自行实现forward逻辑
        pass

后续所有自定义模型都继承GlobalSyncModule，调用任意一个模型的train()/eval()，所有同基类的模型都会同步切换状态。

方法2：维护模型列表统一管理 #

在脚本中维护所有模型的列表，切换模式时遍历列表批量操作：

# 初始化所有模型并加入统一列表
models = []
model1 = nn.Linear(10, 2)
models.append(model1)
model2 = nn.Conv2d(3, 16, 3)
models.append(model2)

# 全局切换到评估模式
for model in models:
    model.eval()

# 全局切换到训练模式
for model in models:
    model.train()

这种方法简单直接，适合脚本内模型数量不多的场景，只需确保所有模型都被加入列表即可。

方法3：上下文管理器封装临时切换 #

编写上下文管理器，在指定代码块内自动切换所有模型的模式，退出后恢复原状态：

from contextlib import contextmanager

@contextmanager
def batch_model_mode(models, target_mode='eval'):
    # 保存所有模型的原始状态
    original_states = [model.training for model in models]
    # 批量设置目标模式
    for model in models:
        model.train(target_mode == 'train')
    try:
        yield
    finally:
        # 恢复所有模型的原始状态
        for model, state in zip(models, original_states):
            model.train(state)

# 使用示例：临时切换到评估模式执行推理
with batch_model_mode(models, target_mode='eval'):
    output = model1(torch.randn(1, 10))
    # 此处所有模型均为eval模式

需要说明的是，PyTorch的training属性是每个nn.Module实例独立维护的，不存在真正的“会话级全局开关”，上述方法本质是通过批量操作或自定义逻辑来模拟全局控制的效果。

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