Method test_train_eval

test/test_backbone_utils.py:227–292

(self)

        out_agg.backward()

    def test_train_eval(self):
        class TestModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.dropout = torch.nn.Dropout(p=1.0)

            def forward(self, x):
                x = x.float().mean()
                x = self.dropout(x)  # dropout
                if self.training:
                    x += 100  # add
                else:
                    x *= 0  # mul
                x -= 0  # sub
                return x

        model = TestModel()

        train_return_nodes = ["dropout", "add", "sub"]
        eval_return_nodes = ["dropout", "mul", "sub"]

        def checks(model, mode):
            with torch.no_grad():
                out = model(torch.ones(10, 10))
            if mode == "train":
                # Check that dropout is respected
                assert out["dropout"].item() == 0
                # Check that control flow dependent on training_mode is respected
                assert out["sub"].item() == 100
                assert "add" in out
                assert "mul" not in out
            elif mode == "eval":
                # Check that dropout is respected
                assert out["dropout"].item() == 1
                # Check that control flow dependent on training_mode is respected
                assert out["sub"].item() == 0
                assert "mul" in out
                assert "add" not in out

        # Starting from train mode
        model.train()
        fx_model = self._create_feature_extractor(
            model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
        )
        # Check that the models stay in their original training state
        assert model.training
        assert fx_model.training
        # Check outputs
        checks(fx_model, "train")
        # Check outputs after switching to eval mode
        fx_model.eval()
        checks(fx_model, "eval")

        # Starting from eval mode
        model.eval()
        fx_model = self._create_feature_extractor(
            model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
        )