torch_em.trainer.default_trainer

  1from __future__ import annotations
  2
  3import os
  4import time
  5import inspect
  6import warnings
  7import contextlib
  8from tqdm import tqdm
  9from functools import partial
 10from datetime import datetime
 11from collections import OrderedDict
 12from importlib import import_module
 13from typing import Any, Callable, Dict, Optional, Union, Literal
 14
 15import numpy as np
 16
 17import torch
 18
 19from .wandb_logger import WandbLogger
 20from .tensorboard_logger import TensorboardLogger
 21from ..util import auto_compile, get_constructor_arguments, is_compiled
 22
 23
 24class DefaultTrainer:
 25    """Trainer class for training a segmentation network.
 26
 27    The trainer class implements the core logic for training a network with pytorch.
 28    It implements a training loop to run training and validation, which is started with `fit`.
 29    The checkpoints and logs from the training run will be saved in the current working directory,
 30    or in the directory specifified by `save_root`. Training can be continued from a checkpoint
 31    by passing it's location to the `load_from_checkpoint` argument of `fit`.
 32
 33    A pre-configured instance of the trainer can be obtained from `torch_em.default_segmentation_trainer`.
 34    Alternatively, the trainer class can also be instantiated as in this example:
 35    ```python
 36    import torch
 37    from torch_em.loss import DiceLoss
 38    from torch_em.model import UNet2d
 39    from torch_em.data.datasets.light_microscopy import get_dsb_loader
 40    from torch_em.trainer import DefaultTrainer
 41
 42    # The training data will be downloaded to this location.
 43    data_root = "/path/to/save/the/training/data"
 44    patch_shape = (256, 256)
 45
 46    # Create the model and optimizer.
 47    model = UNet2d(in_channels=1, out_channels=1)
 48    optimizer = torch.optim.AdamW(model.parameters())
 49
 50    trainer = DefaultTrainer(
 51        name="unet-training",
 52        train_loader=get_dsb_loader(path=data_root, patch_shape=patch_shape, split="train"),
 53        val_loader=get_dsb_loader(path=data_root, patch_shape=patch_shape, split="test"),
 54        model=model,
 55        loss=DiceLoss(),  # The loss function.
 56        optimizer=optimizer,
 57        metric=DiceLoss(),  # The metric. The trainer expects smaller values to represent better results.
 58        device="cuda",  # The device to use for training.
 59    )
 60    trainer.fit(iterations=int(2.5e4))  # Train for 25.000 iterations.
 61    ```
 62
 63    Args:
 64        name: The name of the checkpoint that will be created by the trainer.
 65        train_loader: The data loader containing the training data.
 66        val_loader: The data loader containing the validation data.
 67        model: The model to train.
 68        loss: The loss function for training.
 69        optimizer: The optimizer.
 70        metric: The metric for validation.
 71        device: The torch device to use for training. If None, will use a GPU if available.
 72        lr_scheduler: The learning rate scheduler.
 73        log_image_interval: The interval for saving images during logging, in training iterations.
 74        mixed_precision: Whether to train with mixed precision.
 75        early_stopping: The patience for early stopping in epochs. If None, early stopping will not be used.
 76        logger: The logger class. Will be instantiated for logging.
 77            By default uses `torch_em.training.tensorboard_logger.TensorboardLogger`.
 78        logger_kwargs: The keyword arguments for the logger class.
 79        id_: Unique identifier for the trainer. If None then `name` will be used.
 80        save_root: The root folder for saving the checkpoint and logs.
 81        compile_model: Whether to compile the model before training.
 82        rank: Rank argument for distributed training. See `torch_em.multi_gpu_training` for details.
 83        mixed_precision_dtype: The dtype for autocast in mixed precision training, 'float16' or 'bfloat16'.
 84            The default is 'float16' on the GPU and 'bfloat16' on the CPU. Use 'bfloat16' to avoid overflows.
 85    """
 86    def __init__(
 87        self,
 88        name: Optional[str],
 89        train_loader: torch.utils.data.DataLoader,
 90        val_loader: torch.utils.data.DataLoader,
 91        model: torch.nn.Module,
 92        loss: torch.nn.Module,
 93        optimizer: torch.optim.Optimizer,
 94        metric: Callable,
 95        device: Union[str, torch.device],
 96        lr_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
 97        log_image_interval: int = 100,
 98        mixed_precision: bool = True,
 99        early_stopping: Optional[int] = None,
100        logger=TensorboardLogger,
101        logger_kwargs: Optional[Dict[str, Any]] = None,
102        id_: Optional[str] = None,
103        save_root: Optional[str] = None,
104        compile_model: Optional[Union[bool, str]] = None,
105        rank: Optional[int] = None,
106        mixed_precision_dtype: Optional[str] = None,
107    ):
108        if name is None and not issubclass(logger, WandbLogger):
109            raise TypeError("Name cannot be None if not using the WandbLogger")
110
111        self._generate_name = name is None
112        self.name = name
113        self.id_ = id_ or name
114        self.train_loader = train_loader
115        self.val_loader = val_loader
116        self.model = model
117        self.loss = loss
118        self.optimizer = optimizer
119        self.metric = metric
120        self.device = torch.device(device)
121        self.lr_scheduler = lr_scheduler
122        self.log_image_interval = log_image_interval
123        self.save_root = save_root
124        self.compile_model = compile_model
125        self.rank = rank
126        self._device_type = "cpu" if self.device.type == "cpu" else "cuda"
127
128        self._iteration = 0
129        self._epoch = 0
130        self._best_epoch = 0
131
132        self.mixed_precision = mixed_precision
133        # These are the defaults of torch.autocast for each device type.
134        self.mixed_precision_dtype = mixed_precision_dtype or ("bfloat16" if self._device_type == "cpu" else "float16")
135        self.early_stopping = early_stopping
136        self.train_time = 0.0
137
138        if mixed_precision:
139            # Only float16 needs gradient scaling. bfloat16 has the same range as float32.
140            self.scaler = torch.GradScaler(self._device_type, enabled=self.mixed_precision_dtype == "float16")
141        else:
142            self.scaler = None
143
144        self.logger_class = logger
145        self.logger_kwargs = logger_kwargs
146        self.log_image_interval = log_image_interval
147
148    @property
149    def checkpoint_folder(self):
150        assert self.id_ is not None  # Because the logger may generate and set trainer.id on logger.__init__.
151        # Save_root enables saving the checkpoints somewhere else than in the local folder.
152        # This is handy for filesystems with limited space, where saving the checkpoints
153        # and log files can lead to running out of space.
154        save_root = getattr(self, "save_root", None)
155        return os.path.join("./checkpoints", self.id_) if save_root is None else\
156            os.path.join(save_root, "./checkpoints", self.id_)
157
158    @property
159    def iteration(self):
160        return self._iteration
161
162    @property
163    def epoch(self):
164        return self._epoch
165
166    class Deserializer:
167        """Determines how to deserialize the trainer kwargs from serialized 'init_data'.
168
169        Examples:
170            To extend the initialization process you can inherite from this Deserializer in an inherited Trainer class.
171            Note that `DefaultTrainer.Deserializer.load_generic()` covers most cases already.
172
173            This example adds `the_answer` kwarg, which requires 'calculations' upon initialization:
174            >>> class MyTrainer(DefaultTrainer):
175            >>>     def __init__(self, *args, the_answer: int, **kwargs):
176            >>>         super().__init__(*args, **kwargs)
177            >>>         self.the_answer = the_answer  # this allows the default Serializer to save the new kwarg,
178            >>>                                       # see DefaultTrainer.Serializer
179            >>>
180            >>>     class Deserializer(DefaultTrainer.Deserializer):
181            >>>         def load_the_answer(self):
182            >>>             generic_answer = self.init_data["the_answer"]
183            >>>             # (device dependent) special deserialization
184            >>>             if self.trainer_kwargs["device"].type == "cpu":  # accessing previously deserialized kwarg
185            >>>                 self.trainer_kwargs["the_answer"] = generic_answer + 1
186            >>>             else:
187            >>>                 self.trainer_kwargs["the_answer"] = generic_answer * 2
188
189        Args:
190            init_data: The initialization data of the trainer.
191            save_path: The path where the checkpoint was saved.
192            device: The device.
193        """
194
195        def __init__(self, init_data: Dict, save_path: str, device: Union[str, torch.device]):
196            self.init_data = init_data
197            self.save_path = save_path
198            # Populate with deserialized trainer kwargs during deserialization; possibly overwrite 'device'.
199            self.trainer_kwargs: Dict[str, Any] = dict(
200                device=torch.device(self.init_data["device"]) if device is None else torch.device(device)
201            )
202
203        def load(self, kwarg_name: str, optional):
204            """@private
205            """
206            # `optional` is True if self.trainer.__class__.__init__ specifies a default value for 'kwarg_name'
207            if kwarg_name == "device":
208                pass  # deserialized in __init__
209            elif kwarg_name.endswith("_loader"):
210                self.load_data_loader(kwarg_name, optional)
211            else:
212                load = getattr(self, f"load_{kwarg_name}", self.load_generic)
213                load(kwarg_name, optional=optional)
214
215        def load_data_loader(self, loader_name, optional) -> None:
216            """@private
217            """
218            ds = self.init_data.get(loader_name.replace("_loader", "_dataset"))
219            if ds is None and optional:
220                return
221
222            loader_kwargs = self.init_data[f"{loader_name}_kwargs"]
223            loader = torch.utils.data.DataLoader(ds, **loader_kwargs)
224            # monkey patch shuffle loader_name to the loader
225            loader.shuffle = loader_kwargs.get("shuffle", False)
226            self.trainer_kwargs[loader_name] = loader
227
228        def load_generic(
229            self,
230            kwarg_name: str,
231            *dynamic_args: Dict,
232            optional: bool,
233            only_class: bool = False,
234            dynamic_kwargs: Optional[Dict[str, Any]] = None,
235        ) -> None:
236            """@private
237            """
238            if kwarg_name in self.init_data:
239                self.trainer_kwargs[kwarg_name] = self.init_data[kwarg_name]
240                return
241
242            this_cls = self.init_data.get(f"{kwarg_name}_class", None)
243            if this_cls is None:
244                if optional:
245                    return
246                else:
247                    raise RuntimeError(f"Could not find init data for {kwarg_name} in {self.save_path}")
248
249            assert isinstance(this_cls, str), this_cls
250            assert "." in this_cls, this_cls
251            cls_p, cls_m = this_cls.rsplit(".", 1)
252            this_cls = getattr(import_module(cls_p), cls_m)
253            if only_class:
254                self.trainer_kwargs[kwarg_name] = this_cls
255            else:
256                self.trainer_kwargs[kwarg_name] = this_cls(
257                    *dynamic_args, **self.init_data.get(f"{kwarg_name}_kwargs", {}), **(dynamic_kwargs or {})
258                )
259
260        def load_name(self, kwarg_name: str, optional: bool):
261            """@private
262            """
263            self.trainer_kwargs[kwarg_name] = os.path.split(os.path.dirname(self.save_path))[1]
264
265        def load_optimizer(self, kwarg_name: str, optional: bool):
266            """@private
267            """
268            self.load_generic(kwarg_name, self.trainer_kwargs["model"].parameters(), optional=optional)
269
270        def load_lr_scheduler(self, kwarg_name: str, optional: bool):
271            """@private
272            """
273            self.load_generic(kwarg_name, self.trainer_kwargs["optimizer"], optional=optional)
274
275        # todo: remove and rename kwarg 'logger' to 'logger_class'
276        def load_logger(self, kwarg_name: str, optional: bool):
277            """@private
278            """
279            assert kwarg_name == "logger"
280            self.load_generic("logger", optional=optional, only_class=True)
281
282    @staticmethod
283    def _get_save_dict(save_path, device):
284        if not os.path.exists(save_path):
285            raise ValueError(f"Cannot find checkpoint {save_path}")
286        return torch.load(save_path, map_location=device, weights_only=False)
287
288    @classmethod
289    def from_checkpoint(
290        cls,
291        checkpoint_folder: Union[os.PathLike, str],
292        name: Literal["best", "latest"] = "best",
293        device: Optional[Union[str, torch.device]] = None,
294    ):
295        """@private
296        """
297        save_path = os.path.join(checkpoint_folder, f"{name}.pt")
298        # make sure the correct device is set if we don't have access to CUDA
299        if not torch.cuda.is_available():
300            device = "cpu"
301        save_dict = cls._get_save_dict(save_path, device)
302        deserializer = cls.Deserializer(save_dict["init"], save_path, device)
303
304        has_kwargs = False
305        deserialized = []
306        for name, parameter in inspect.signature(cls).parameters.items():
307            if name == "kwargs":
308                has_kwargs = True
309                continue
310            deserializer.load(name, optional=parameter.default is not inspect.Parameter.empty)
311            deserialized.append(name)
312
313        # to deserialze kwargs we can't rely on inspecting the signature, so we
314        # go through the remaning kwarg names in init data instead
315        if has_kwargs:
316            kwarg_names = list(set(deserializer.init_data.keys()) - set(deserialized))
317            for name in kwarg_names:
318                if name.endswith("_kwargs"):
319                    continue
320                elif name.endswith("_dataset"):
321                    deserializer.load(name.replace("dataset", "loader"), optional=False)
322                elif name.endswith("_class"):
323                    deserializer.load(name.replace("_class", ""), optional=False)
324                else:
325                    deserializer.load(name, optional=False)
326
327        trainer = cls(**deserializer.trainer_kwargs)
328        trainer._initialize(0, save_dict)
329        trainer._is_initialized = True
330        return trainer
331
332    class Serializer:
333        """Implements how to serialize trainer kwargs from a trainer instance.
334
335        Examples:
336            To extend the serialization process you can inherite from this Serializer in a derived Trainer class.
337            Note that the methods `dump_generic_builtin()`, `dump_generic_class()` and `dump_generic_instance()`
338            called by the `dump()` method when appropriate cover most cases already.
339
340            This example adds `the_answer` kwarg, which requires extra steps on dumping only because we don't keep a
341            'the_answer' attribute:
342            >>> class MyTrainer(DefaultTrainer):
343            >>>     def __init__(self, *args, the_answer: int, **kwargs):
344            >>>         super().__init__(*args, **kwargs)
345            >>>         # self.the_answer = the_answer  # this would allow the default Serializer to save the new kwarg,
346            >>>         # but let's make things more interesting...
347            >>>         self.the = the_answer // 10
348            >>>         self.answer = the_answer % 10
349            >>>
350            >>>     class Serializer(DefaultTrainer.Serializer):
351            >>>         trainer: MyTrainer
352            >>>         def dump_the_answer(self, kwarg_name: str) -> None:  # custom dump method for 'the_answer' kwarg
353            >>>             assert kwarg_name == "the_answer"
354            >>>             # populate self.init_data with the serialized data required by Deserializer
355            >>>             # to restore the trainer kwargs
356            >>>             self.init_data["the_answer"] = self.trainer.the * 10 + self.trainer.answer
357
358            This example with both Serializer and Deserializer adds `the_answer` kwarg,
359            while saving it in two separate entries 'the' and 'answer'
360            >>> class MyTrainer(DefaultTrainer):
361            >>>     def __init__(self, *args, the_answer: int, **kwargs):
362            >>>         super().__init__(*args, **kwargs)
363            >>>         self.the_answer = the_answer
364            >>>
365            >>>     class Serializer(DefaultTrainer.Serializer):
366            >>>         trainer: MyTrainer
367            >>>         def dump_the_answer(self, kwarg_name: str):
368            >>>             assert kwarg_name == "the_answer"
369            >>>             self.init_data.update({
370            >>>                 "the": self.trainer.the_answer // 10,
371            >>>                 "answer": self.trainer.the_answer % 10
372            >>>             })
373            >>>
374            >>>     class Deserializer(DefaultTrainer.Deserializer):
375            >>>         def load_the_answer(self, kwarg_name: str, optional: bool):
376            >>>             assert kwarg_name == "the_answer"
377            >>>             # 'optional' is True if MyTrainer.__init__ specifies a default value for 'kwarg_name'
378            >>>             self.trainer_kwargs[kwarg_name] = self.init_data["the"] * 10 + self.init_data["answer"]
379
380        Args:
381            trainer: The trainer instance.
382        """
383
384        def __init__(self, trainer: DefaultTrainer):
385            self.trainer = trainer
386            self.init_data = {}  # to be populated during serialization process
387
388        def dump(self, kwarg_name: str) -> None:
389            """@private
390            """
391            dumper = getattr(self, f"dump_{kwarg_name}", None)
392            if dumper is not None:
393                dumper(kwarg_name)
394            elif kwarg_name.endswith("_loader"):
395                self.dump_data_loader(kwarg_name)
396            elif kwarg_name.endswith("_class"):
397                self.dump_generic_class(kwarg_name)
398            elif not hasattr(self.trainer, kwarg_name):
399                raise AttributeError(
400                    f"{self.trainer.__class__} missing attribute '{kwarg_name}' "
401                    f"or special dump method {self.trainer.__class__}.Serializer.dump_{kwarg_name}()"
402                )
403            else:
404                assert hasattr(self.trainer, kwarg_name)
405                obj = getattr(self.trainer, kwarg_name)
406                if obj is None or type(obj) in (
407                    bool,
408                    bytearray,
409                    bytes,
410                    dict,
411                    float,
412                    frozenset,
413                    int,
414                    list,
415                    set,
416                    str,
417                    tuple,
418                ):
419                    self.dump_generic_builtin(kwarg_name)
420                else:
421                    self.dump_generic_instance(kwarg_name)
422
423        def dump_generic_builtin(self, kwarg_name: str) -> None:
424            """@private
425            """
426            assert hasattr(self.trainer, kwarg_name)
427            self.init_data[kwarg_name] = getattr(self.trainer, kwarg_name)
428
429        def dump_generic_class(self, kwarg_name: str) -> None:
430            """@private
431            """
432            assert hasattr(self.trainer, kwarg_name)
433            assert kwarg_name.endswith("_class")
434            obj = getattr(self.trainer, kwarg_name)
435            self.init_data[kwarg_name] = None if obj is None else f"{obj.__module__}.{obj.__name__}"
436
437        def dump_generic_instance(self, kwarg_name: str) -> None:
438            """@private
439            """
440            assert hasattr(self.trainer, kwarg_name)
441            instance = getattr(self.trainer, kwarg_name)
442            self.init_data.update(
443                {
444                    f"{kwarg_name}_class": f"{instance.__class__.__module__}.{instance.__class__.__name__}",
445                    f"{kwarg_name}_kwargs": get_constructor_arguments(instance),
446                }
447            )
448
449        def dump_device(self, kwarg_name: str):
450            """@private
451            """
452            assert hasattr(self.trainer, kwarg_name)
453            self.init_data[kwarg_name] = str(getattr(self.trainer, kwarg_name))
454
455        def dump_data_loader(self, kwarg_name: str) -> None:
456            """@private
457            """
458            assert hasattr(self.trainer, kwarg_name)
459            loader = getattr(self.trainer, kwarg_name)
460            if loader is None:
461                return
462            self.init_data.update(
463                {
464                    f"{kwarg_name.replace('_loader', '_dataset')}": loader.dataset,
465                    f"{kwarg_name}_kwargs": get_constructor_arguments(loader),
466                }
467            )
468
469        def dump_logger(self, kwarg_name: str):  # todo: remove and rename kwarg 'logger' to 'logger_class'
470            """@private
471            """
472            self.dump_generic_class(f"{kwarg_name}_class")
473
474        def dump_model(self, kwarg_name: str):
475            """@private
476            """
477            if is_compiled(self.trainer.model):
478                self.init_data.update(
479                    {"model_class": self.trainer._model_class, "model_kwargs": self.trainer._model_kwargs}
480                )
481            else:
482                self.dump_generic_instance("model")
483
484    def _build_init(self) -> Dict[str, Any]:
485        serializer = self.Serializer(self)
486        for name in inspect.signature(self.__class__).parameters:
487            # special rules to serialize kwargs
488            # if a trainer class inherits from DefaultTrainer and has **kwargs
489            # they need to be saved in self._kwargs
490            if name == "kwargs":
491                if not hasattr(self, "_kwargs"):
492                    msg = "The trainer class has **kwargs in its signature, but is missing the _kwargs attribute. " +\
493                          "Please add self._kwargs to its __init__ function"
494                    raise RuntimeError(msg)
495                kwargs = getattr(self, "_kwargs")
496                for kwarg_name in kwargs:
497                    serializer.dump(kwarg_name)
498                continue
499            serializer.dump(name)
500
501        return serializer.init_data
502
503    def _initialize(self, iterations, load_from_checkpoint, epochs=None):
504        assert self.train_loader is not None
505        assert self.val_loader is not None
506        assert self.model is not None
507        assert self.loss is not None
508        assert self.optimizer is not None
509        assert self.metric is not None
510        assert self.device is not None
511
512        if load_from_checkpoint is not None:
513            self.load_checkpoint(load_from_checkpoint)
514
515        if sum((iterations is not None, epochs is not None)) != 1:
516            raise ValueError(
517                "Exactly one of 'iterations' or 'epochs' has to be specified to initialize the trainer."
518                f"You have passed 'iterations'={iterations} and 'epochs'={epochs}"
519            )
520
521        if epochs is None:
522            epochs = int(np.ceil(float(iterations) / len(self.train_loader)))
523        else:
524            iterations = epochs * len(self.train_loader)
525
526        self.max_iteration = self._iteration + iterations
527        self.max_epoch = self._epoch + epochs
528
529        if not getattr(self, "_is_initialized", False):
530            # check if we compile the model (only supported by pytorch 2)
531            # to enable (de)serialization of compiled models, we keep track of the model class and kwargs
532            if is_compiled(self.model):
533                warnings.warn(
534                    "You have passed a compiled model to the trainer."
535                    "It will not be possible to (de)serialize the trainer with it."
536                    "If you want to be able to do this please pass the normal model."
537                    "It can be automatically compiled by setting 'compile_model' to True"
538                )
539            self._model_class = f"{self.model.__class__.__module__}.{self.model.__class__.__name__}"
540            self._model_kwargs = get_constructor_arguments(self.model)
541            self.model = auto_compile(self.model, self.compile_model)
542
543            self.model.to(self.device)
544            self.loss.to(self.device)
545
546            # this saves all the information that is necessary
547            # to fully load the trainer from the checkpoint
548            self.init_data = self._build_init()
549
550            if self.logger_class is None:
551                self.logger = None
552            else:
553                # may set self.name if self.name is None
554                save_root = getattr(self, "save_root", None)
555                try:
556                    self.logger = self.logger_class(self, save_root, **(self.logger_kwargs or {}))
557                except (PermissionError, RuntimeError):
558                    warnings.warn(
559                        f"The checkpoint folder at {self.checkpoint_folder} could not be created."
560                        "The most likely reason for this is that you copied the checkpoint somewhere else,"
561                        "so we skip this error to enable loading the model from this checkpoint."
562                    )
563
564            try:
565                os.makedirs(self.checkpoint_folder, exist_ok=True)
566            except PermissionError:
567                warnings.warn(
568                    f"The checkpoint folder at {self.checkpoint_folder} could not be created."
569                    "The most likely reason for this is that you copied the checkpoint somewhere else,"
570                    "so we skip this error to enable loading the model from this checkpoint."
571                )
572                pass
573
574        best_metric = np.inf
575        return best_metric
576
577    def save_checkpoint(self, name, current_metric, best_metric, train_time=0.0, **extra_save_dict):
578        """@private
579        """
580        save_path = os.path.join(self.checkpoint_folder, f"{name}.pt")
581        extra_init_dict = extra_save_dict.pop("init", {})
582        save_dict = {
583            "iteration": self._iteration,
584            "epoch": self._epoch,
585            "best_epoch": self._best_epoch,
586            "best_metric": best_metric,
587            "current_metric": current_metric,
588            "model_state": self.model.state_dict(),
589            "optimizer_state": self.optimizer.state_dict(),
590            "init": self.init_data | extra_init_dict,
591            "train_time": train_time,
592            "timestamp": datetime.now().strftime("%d-%m-%Y (%H:%M:%S)"),
593        }
594        save_dict.update(**extra_save_dict)
595        if self.scaler is not None:
596            save_dict.update({"scaler_state": self.scaler.state_dict()})
597        if self.lr_scheduler is not None:
598            save_dict.update({"scheduler_state": self.lr_scheduler.state_dict()})
599
600        rank = getattr(self, "rank", None)
601        if rank is None or rank == 0:
602            torch.save(save_dict, save_path)
603
604    def load_checkpoint(self, checkpoint="best"):
605        """@private
606        """
607        if isinstance(checkpoint, str):
608            save_path = os.path.join(self.checkpoint_folder, f"{checkpoint}.pt")
609            if not os.path.exists(save_path):
610                warnings.warn(f"Cannot load checkpoint. {save_path} does not exist.")
611                return
612            save_dict = torch.load(save_path, weights_only=False)
613        elif isinstance(checkpoint, dict):
614            save_dict = checkpoint
615        else:
616            raise RuntimeError
617
618        self._iteration = save_dict["iteration"]
619        self._epoch = save_dict["epoch"]
620        self._best_epoch = save_dict["best_epoch"]
621        self.best_metric = save_dict["best_metric"]
622        self.current_metric = save_dict["current_metric"]
623        self.train_time = save_dict.get("train_time", 0.0)
624
625        model_state = save_dict["model_state"]
626        # to enable loading compiled models
627        compiled_prefix = "_orig_mod."
628        model_state = OrderedDict(
629            [(k[len(compiled_prefix):] if k.startswith(compiled_prefix) else k, v) for k, v in model_state.items()]
630        )
631        self.model.load_state_dict(model_state)
632        # we need to send the network to the device before loading the optimizer state!
633        self.model.to(self.device)
634
635        self.optimizer.load_state_dict(save_dict["optimizer_state"])
636        scaler_state = save_dict.get("scaler_state")
637        if self.scaler is not None and scaler_state:
638            self.scaler.load_state_dict(scaler_state)
639        if self.lr_scheduler is not None:
640            self.lr_scheduler.load_state_dict(save_dict["scheduler_state"])
641
642        return save_dict
643
644    def _verify_if_training_completed(self, checkpoint="latest"):
645        save_path = os.path.join(self.checkpoint_folder, f"{checkpoint}.pt")
646        save_dict = torch.load(save_path, weights_only=False) if os.path.exists(save_path) else None
647        if save_dict and self.max_iteration == save_dict.get("iteration"):
648            return True
649        return False
650
651    def fit(
652        self,
653        iterations: Optional[int] = None,
654        load_from_checkpoint: Optional[Union[os.PathLike, str]] = None,
655        epochs: Optional[int] = None,
656        save_every_kth_epoch: Optional[int] = None,
657        progress=None,
658        overwrite_training: bool = True,
659    ):
660        """Run neural network training.
661
662        Exactly one of 'iterations' or 'epochs' has to be passed.
663
664        Args:
665            iterations: How long to train, specified in iterations.
666            load_from_checkpoint: Path to a checkpoint from where training should be continued .
667            epochs: How long to train, specified in epochs.
668            save_every_kth_epoch: Save checkpoints after every kth epoch in a separate file.
669                The corresponding checkpoints will be saved with the naming scheme 'epoch-{epoch}.pt'.
670            progress: Optional progress bar for integration with external tools. Expected to follow the tqdm interface.
671            overwrite_training: Whether to overwrite existing checkpoints in the save directory.
672        """
673        best_metric = self._initialize(iterations, load_from_checkpoint, epochs)
674
675        if not overwrite_training:
676            if load_from_checkpoint is not None:
677                raise ValueError(
678                    "We do not support 'overwrite_training=False' and 'load_from_checkpoint' at the same time."
679                )
680
681            if self._verify_if_training_completed():
682                print(
683                    f"The model is trained for {self.max_iteration} iterations / {self.max_epoch} epochs "
684                    "and 'overwrite_training' is set to 'False'."
685                )
686                print(f"The checkpoints are located at '{os.path.abspath(self.checkpoint_folder)}'.")
687                return
688
689        print(
690            "Start fitting for",
691            self.max_iteration - self._iteration,
692            "iterations / ",
693            self.max_epoch - self._epoch,
694            "epochs",
695        )
696        print("with", len(self.train_loader), "iterations per epoch")
697
698        if self.mixed_precision:
699            train_epoch = self._train_epoch_mixed
700            validate = self._validate_mixed
701            print("Training with mixed precision")
702        else:
703            train_epoch = self._train_epoch
704            validate = self._validate
705            print("Training with single precision")
706
707        total_iterations = epochs * len(self.train_loader) if iterations is None else iterations
708        if progress is None:
709            progress = tqdm(total=total_iterations, desc=f"Epoch {self._epoch}", leave=True)
710        else:
711            progress.total = total_iterations
712            progress.set_description(f"Epoch {self._epoch}")
713
714        msg = "Epoch %i: average [s/it]: %f, current metric: %f, best metric: %f"
715        train_epochs = self.max_epoch - self._epoch
716        t_start = time.time()
717        for epoch in range(train_epochs):
718
719            # Ensure data is shuffled differently at each epoch.
720            try:
721                self.train_loader.sampler.set_epoch(epoch)
722            except AttributeError:
723                pass
724
725            # Run training and validation for this epoch
726            t_per_iter = train_epoch(progress)
727            current_metric = validate()
728
729            # perform all the post-epoch steps:
730
731            # apply the learning rate scheduler
732            if self.lr_scheduler is not None:
733                self.lr_scheduler.step(current_metric)
734
735            # how long did we train in total?
736            total_train_time = (time.time() - t_start) + self.train_time
737
738            # save this checkpoint as the new best checkpoint if
739            # it has the best overall validation metric
740            if current_metric < best_metric:
741                best_metric = current_metric
742                self._best_epoch = self._epoch
743                self.save_checkpoint("best", current_metric, best_metric, train_time=total_train_time)
744
745            # save this checkpoint as the latest checkpoint
746            self.save_checkpoint("latest", current_metric, best_metric, train_time=total_train_time)
747
748            # if we save after every k-th epoch then check if we need to save now
749            if save_every_kth_epoch is not None and (self._epoch + 1) % save_every_kth_epoch == 0:
750                self.save_checkpoint(
751                    f"epoch-{self._epoch + 1}", current_metric, best_metric, train_time=total_train_time
752                )
753
754            # if early stopping has been specified then check if the stopping condition is met
755            if self.early_stopping is not None:
756                epochs_since_best = self._epoch - self._best_epoch
757                if epochs_since_best > self.early_stopping:
758                    print("Stopping training because there has been no improvement for", self.early_stopping, "epochs")
759                    break
760
761            self._epoch += 1
762            progress.set_description(msg % (self._epoch, t_per_iter, current_metric, best_metric), refresh=True)
763
764        print(f"Finished training after {self._epoch} epochs / {self._iteration} iterations.")
765        print(f"The best epoch is number {self._best_epoch}.")
766
767        if self._generate_name:
768            self.name = None
769
770        # Update the train time
771        self.train_time = total_train_time
772
773        # TODO save the model to wandb if we have the wandb logger
774        if isinstance(self.logger, WandbLogger):
775            self.logger.get_wandb().finish()
776
777    def _backprop(self, loss):
778        loss.backward()
779        self.optimizer.step()
780
781    def _backprop_mixed(self, loss):
782        self.scaler.scale(loss).backward()
783        self.scaler.step(self.optimizer)
784        self.scaler.update()
785
786    def _train_epoch(self, progress):
787        return self._train_epoch_impl(progress, contextlib.nullcontext, self._backprop)
788
789    def _train_epoch_mixed(self, progress):
790        return self._train_epoch_impl(
791            progress,
792            partial(torch.autocast, device_type=self._device_type, dtype=getattr(torch, self.mixed_precision_dtype)),
793            self._backprop_mixed
794        )
795
796    def _forward_and_loss(self, x, y):
797        pred = self.model(x)
798        if self._iteration % self.log_image_interval == 0:
799            if pred.requires_grad:
800                pred.retain_grad()
801
802        loss = self.loss(pred, y)
803        return pred, loss
804
805    def _train_epoch_impl(self, progress, forward_context, backprop: Callable[[torch.Tensor], None]):
806        self.model.train()
807
808        n_iter = 0
809        t_per_iter = time.time()
810        for x, y in self.train_loader:
811            x, y = x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
812
813            self.optimizer.zero_grad()
814
815            with forward_context():
816                pred, loss = self._forward_and_loss(x, y)
817
818            backprop(loss)
819
820            lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
821            if self.logger is not None:
822                self.logger.log_train(self._iteration, loss, lr, x, y, pred, log_gradients=True)
823
824            self._iteration += 1
825            n_iter += 1
826            if self._iteration >= self.max_iteration:
827                break
828            progress.update(1)
829
830        t_per_iter = (time.time() - t_per_iter) / n_iter
831        return t_per_iter
832
833    def _validate(self):
834        return self._validate_impl(contextlib.nullcontext)
835
836    def _validate_mixed(self):
837        return self._validate_impl(
838            partial(torch.autocast, device_type=self._device_type, dtype=getattr(torch, self.mixed_precision_dtype))
839        )
840
841    def _validate_impl(self, forward_context):
842        self.model.eval()
843
844        metric_val = 0.0
845        loss_val = 0.0
846
847        with torch.no_grad():
848            for x, y in self.val_loader:
849                x, y = x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
850                with forward_context():
851                    pred, loss = self._forward_and_loss(x, y)
852                    metric = self.metric(pred, y)
853
854                loss_val += loss.item()
855                metric_val += metric.item()
856
857        metric_val /= len(self.val_loader)
858        loss_val /= len(self.val_loader)
859        if self.logger is not None:
860            self.logger.log_validation(self._iteration, metric_val, loss_val, x, y, pred)
861        return metric_val