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

Question

我有一个预训练的 model 与正在训练的 model 一起使用。 我希望预训练的 model 始终处于 eval 模式，但另一个 model 将在 eval 和 train 模式之间来回移动。 不过，我仍然希望预训练的 model 成为另一个的子模块（例如，以便所有参数都保留在同一设备上）。 有没有办法做到这一点？ 这是一个最小的例子：

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

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  # I'd like this to give an error

我知道我可以解决这个问题，例如，总是做

trainable.train()
trainable.fixed_module.eval()

但这很容易出错，并且不适用于现有代码。

Answer 1

一种解决方案是像这样覆盖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

Answer 2

您可以覆盖FixedModule中的train以防止其更改模式。 请注意， eval只调用train(False) ，因此您也不需要覆盖它。 但是现在调用FixedModule.eval不会做任何事情，所以你必须在 init 中设置training = False 。

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