Function eval_mode

monai/networks/utils.py:458–489

Set network(s) to eval mode and then return to original state at the end. Args: nets: Input network(s) Examples .. code-block:: python t=torch.rand(1,1,16,16) p=torch.nn.Conv2d(1,1,3) print(p.training) # True with eval_mode(p):

(*nets: nn.Module)

@contextmanager
def eval_mode(*nets: nn.Module):
    """
    Set network(s) to eval mode and then return to original state at the end.

    Args:
        nets: Input network(s)

    Examples

    .. code-block:: python

        t=torch.rand(1,1,16,16)
        p=torch.nn.Conv2d(1,1,3)
        print(p.training)  # True
        with eval_mode(p):
            print(p.training)  # False
            print(p(t).sum().backward())  # will correctly raise an exception as gradients are calculated
    """

    # Get original state of network(s).
    # Check the training attribute in case it's TensorRT based models which don't have this attribute.
    training = [n for n in nets if hasattr(n, "training") and n.training]

    try:
        # set to eval mode
        with torch.no_grad():
            yield [n.eval() if hasattr(n, "eval") else n for n in nets]
    finally:
        # Return required networks to training
        for n in training:
            if hasattr(n, "train"):
                n.train()


@contextmanager