Class ThreadContainer

monai/utils/jupyter_utils.py:262–372

Contains a running `Engine` object within a separate thread from main thread in a Jupyter notebook. This allows an engine to begin a run in the background and allow the starting notebook cell to complete. A user can thus start a run and then navigate away from the notebook without conce

class ThreadContainer(Thread):
    """
    Contains a running `Engine` object within a separate thread from main thread in a Jupyter notebook. This
    allows an engine to begin a run in the background and allow the starting notebook cell to complete. A
    user can thus start a run and then navigate away from the notebook without concern for loosing connection
    with the running cell. All output is acquired through methods which synchronize with the running engine
    using an internal `lock` member, acquiring this lock allows the engine to be inspected while it's prevented
    from starting the next iteration.

    Args:
        engine: wrapped `Engine` object, when the container is started its `run` method is called
        loss_transform: callable to convert an output dict into a single numeric value
        metric_transform: callable to convert a named metric value into a single numeric value
        status_format: format string for status key-value pairs.
    """

    def __init__(
        self,
        engine: Engine,
        loss_transform: Callable = _get_loss_from_output,
        metric_transform: Callable = lambda name, value: value,
        status_format: str = "{}: {:.4}",
    ):
        super().__init__()
        self.lock = RLock()
        self.engine = engine
        self._status_dict: dict[str, Any] = {}
        self.loss_transform = loss_transform
        self.metric_transform = metric_transform
        self.fig: plt.Figure | None = None
        self.status_format = status_format

        self.engine.add_event_handler(Events.ITERATION_COMPLETED, self._update_status)

    def run(self):
        """Calls the `run` method of the wrapped engine."""
        self.engine.run()

    def stop(self):
        """Stop the engine and join the thread."""
        self.engine.terminate()
        self.join()

    def _update_status(self):
        """Called as an event, updates the internal status dict at the end of iterations."""
        with self.lock:
            state = self.engine.state
            stats: dict[str, Any] = {
                StatusMembers.EPOCHS.value: 0,
                StatusMembers.ITERS.value: 0,
                StatusMembers.LOSS.value: float("nan"),
            }

            if state is not None:
                if state.max_epochs is not None and state.max_epochs >= 1:
                    epoch = f"{state.epoch}/{state.max_epochs}"
                else:
                    epoch = str(state.epoch)