如何将PyTorch子模块保持在评估模式?

一个解决方案是覆盖 train 这样地:

from torch import nn

class FixedModule(nn.Module):
    pass

class TrainableModule(nn.Module):
    def __init__(self, fixed_module):
        super().__init__()
        self.fixed_module = fixed_module

    def train(self):
        super().train()
        self.fixed_module.eval()

fixed = FixedModule().eval()
assert not fixed.training

trainable = TrainableModule(fixed)
assert trainable.training and not trainable.fixed_module.training

trainable.train()
assert trainable.fixed_module.training  # This gives an error now

你可以覆盖 train 在里面 FixedModule 防止它改变模式。注意 eval 只是打电话 train(False) ,所以你也不需要覆盖它。但是打电话 FixedModule.eval 我现在什么都做不了,所以你得准备好 training = False 在init中。

from torch import nn

class FixedModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.training = False

        # add any other nn.Module attributes here before calling self.children

        # you could override `train` in each child too if you really wanted,
        # but that seems like overkill unless there are external references
        # to any submodules of FixedModule
        for module in self.children():
            module.eval()

    def train(self, mode):
        return self

class TrainableModule(nn.Module):
    def __init__(self, fixed_module):
        super().__init__()
        self.fixed_module = fixed_module    

fixed = FixedModule().eval()
assert not fixed.training

trainable = TrainableModule(fixed)
assert trainable.training and not trainable.fixed_module.training

trainable.train()
assert not trainable.fixed_module.training # passes