torch_em.trainer.default_trainer

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