How do I customize my module's behavior for train and eval?

1

Dropout and Batch-Norm already have this kind of behavior. In other words what a Dropout or Batch-Norm module outputs when it is in eval mode is different from when it is in train mode. How can I achieve something similar with my custom module?

1 Like
2

You could use the internal self.training attribute.
Here is a dummy example:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 10)
        
    def forward(self, x):
        x = self.fc(x)
        if self.training:
            x = x * 1000
        return x

model = MyModel()
x = torch.randn(1, 10)

output = model(x)
print(output)
> tensor([[ -151.6117,    20.6451,  -589.1161,  -120.6478,   395.1652,  -950.3046,
         -1062.1073,   973.9295,    61.4954,  -412.5521]],
       grad_fn=<MulBackward0>)

model.eval()
output = model(x)
print(output)
> tensor([[-0.1516,  0.0206, -0.5891, -0.1206,  0.3952, -0.9503, -1.0621,  0.9739,
          0.0615, -0.4126]], grad_fn=<AddmmBackward>)
15 Likes
3

Thank you. That seems very nice and simple.