Function train_mode

monai/networks/utils.py:493–525

Set network(s) to train 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) p.eval() print(p.training) # False with train

(*nets: nn.Module)

@contextmanager
def train_mode(*nets: nn.Module):
    """
    Set network(s) to train 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)
        p.eval()
        print(p.training)  # False
        with train_mode(p):
            print(p.training)  # True
            print(p(t).sum().backward())  # No exception
    """

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

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


def get_state_dict(obj: torch.nn.Module | Mapping):