text: import os
  
  import numpy as np
  import torch
  import time
  from .loss import LossCounter
  
  
  class ModelTrainer:
      """
      Class to train object detection models
      """
  
      def __init__(self, model,
                   optimizer,
                   scheduler,
                   base_dir,
                   device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
                   num_epochs=100,
                   model_file=None,
                   log_path=None,
                   verbose=True,
                   verbose_step=100,
                   keep_models=3):
          """
  
          :param model: Model to train.
          :param device: Device to train the model on. CPU or CUDA.
          :param optimizer: Model optimizer.
          :param scheduler: Learning rate scheduler for training.
          :param base_dir: Default directory for saving output.
          :param num_epochs: Number of epochs to train. Defaults to 100
          :param model_file: File for loading partially trained models.
          :param log_path: Optional value for specifying log file location
          :param verbose: Boolean for printing output. Defaults to True
          :param verbose_step: Interval on which to print status. Defaults to 100.
          :param keep_models: Maximum number of models to save. Defaults to 3
          """
          self.model = model
          self.device = device
          self.model.to(self.device)
          self.optimizer = optimizer
          self.scheduler = scheduler
          self.num_epochs = num_epochs
          self.base_dir = base_dir
          self.verbose = verbose
          self.verbose_step = verbose_step
          self.keep_models = keep_models
          self.zfill = np.ceil(np.log10(self.num_epochs)).astype(int)
          if model_file:
              self.load(model_file)
          else:
              self.start_epoch = 0
              self.best_summary_loss = []
          if not os.path.exists(self.base_dir):
              os.makedirs(self.base_dir)
  
          self.log_path = log_path if log_path else f"{self.base_dir}/log.csv"
          print(f'Fitter prepared. Device is {self.device}')
  
      def _csv_header(self):
          """
          Writes log file header.
          """
          self.log("Epoch, Stage, Summary Loss, Class Loss, Box Loss", print_line=False)
  
      @staticmethod
      def _csv_line(summary_loss, epoch, stage):
          """
          Builds the string for one line of logging to csv file.
  
          :param summary_loss: Total current loss
          :param epoch: Current epoch
          :param stage: Training or Validation
          :return: String to write to log file.
          """
          return f"{epoch}, {stage}, {summary_loss.avg:.5f}," + \
                 f"{summary_loss.class_avg:.5f}," + \
                 f"{summary_loss.box_avg:.5f}"
  
      @staticmethod
      def _output_dict(output):
          return {"loss": output["loss"].detach().item(),
                  "class_loss": output["class_loss"].detach().item(),
                  "box_loss": output["box_loss"].detach().item()}
  
      @staticmethod
      def _print_line(summary_loss, step, total_steps, stage, t):
          """
          Prints a line of output showing the current status of training.
  
          :param summary_loss: Total current loss.
          :param step: Current step in epoch.
          :param total_steps: Total number of steps in epoch.
          :param stage: Training or Validation.
          :param t: Start time of training or validation epoch.
          """
          print(
              f"{stage} Step {step}/{total_steps}, " +
              f"summary_loss: {summary_loss.avg:.5f}, " +
              f"class_loss: {summary_loss.class_avg:.5f}, " +
              f"box_loss: {summary_loss.box_avg:.5f}, " +
              f"time: {(time.time() - t):.5f}\n")
  
      def fit(self, train_loader, validation_loader):
          """
          Trains the model. Saves best models after each epoch and updates learning rate.
  
          :param train_loader: Data Loader for training data.
          :param validation_loader: Data Loader for validation data.
          """
          if self.start_epoch == 0:
              self._csv_header()
          for epoch in range(self.start_epoch, self.num_epochs):
              # Train one epoch
              summary_loss = self.train(train_loader)
  
              self.log(self._csv_line(summary_loss, epoch, "Train"))
              self.save(f'{self.base_dir}/last-checkpoint.bin', epoch)
  
              # Evaluate on validation data.
              summary_loss = self.validate(validation_loader)
  
              self.log(self._csv_line(summary_loss, epoch, "Val"))
  
              # Update list of saved models
              if len(self.best_summary_loss) < self.keep_models:
                  self.best_summary_loss.append((summary_loss.avg, epoch))
                  self.save(f'{self.base_dir}/best-checkpoint-{str(epoch).zfill(self.zfill)}epoch.bin', epoch)
                  self.best_summary_loss.sort()
              elif summary_loss.avg < self.best_summary_loss[-1][0]:
                  _, old_epoch = self.best_summary_loss.pop()
                  self.best_summary_loss.append((summary_loss.avg, epoch))
                  self.save(f'{self.base_dir}/best-checkpoint-{str(epoch).zfill(self.zfill)}epoch.bin', epoch)
                  os.remove(f'{self.base_dir}/best-checkpoint-{str(old_epoch).zfill(self.zfill)}epoch.bin')
                  self.best_summary_loss.sort()
  
              # Update learning rate.
              if self.scheduler:
                  self.scheduler.step()
  
      def load(self, path):
          """
          Loads saved model, including optimizer, scheduler history, and previous best models.
  
          :param path: Path of model to load.
          """
          checkpoint = torch.load(path)
          self.model.model.load_state_dict(checkpoint["state_dict"])
          self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
          self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
          self.best_summary_loss = checkpoint["best_summary_loss"]
          self.start_epoch = checkpoint["epoch"] + 1
  
      def log(self, message, print_line=True):
          """
          Write data to log.
          :param message: Message to write.
          :param print_line: Boolean, print message to stdout. Default True.
          """
          if self.verbose and print_line:
              print(message)
          with open(self.log_path, 'a+') as logger:
              logger.write(f'{message}\n')
  
      def save(self, path, epoch):
          """
          Saves model after training epoch.
  
          :param path: Path of file for saved model
          :param epoch: Epoch number.
          """
          self.model.eval()
          torch.save({
              "state_dict": self.model.model.state_dict(),
              "optimizer_state_dict": self.optimizer.state_dict(),
              "scheduler_state_dict": self.scheduler.state_dict(),
              "best_summary_loss": self.best_summary_loss,
              "epoch": epoch,
          }, path)
  
      def train(self, train_loader):
          """
          Trains model for one epoch.
          :param train_loader: Data Loader for training data
          :type train_loader: torch.utils.data.DataLoader
          :return: Total Loss on training data for epoch.
          """
          self.model.train()
          summary_loss = LossCounter()
          t = time.time()
          for step, (images, targets) in enumerate(train_loader):
              if self.verbose and step and step % self.verbose_step == 0:
                  self._print_line(summary_loss, step, len(train_loader), "Train", t)
  
              images = torch.stack(images)
              images = images.to(self.device).float()
              batch_size = images.shape[0]
              boxes = [target['bboxes'].to(self.device).float() for target in targets]
              labels = [target['labels'].to(self.device).float() for target in targets]
  
              self.optimizer.zero_grad()
  
              output = self.model(images, {'bbox': boxes, 'cls': labels})
  
              output['loss'].backward()
  
              summary_loss.update_avg(self._output_dict(output), batch_size)
  
              self.optimizer.step()
  
          return summary_loss
  
      def validate(self, val_loader):
          """
          Evaluates model on validation data after one epoch of training.
  
          :param val_loader: Data Loader for validation data.
          :type val_loader: torch.utils.data.DataLoader
          :return Total Loss on validation data.
          """
          self.model.eval()
          summary_loss = LossCounter()
          t = time.time()
          for step, (images, targets) in enumerate(val_loader):
              if self.verbose and step and step % self.verbose_step == 0:
                  self._print_line(summary_loss, step, len(val_loader), "Val", t)
              with torch.no_grad():
                  images = torch.stack(images)
                  batch_size = images.shape[0]
                  images = images.to(self.device).float()
                  boxes = [target['bboxes'].to(self.device).float() for target in targets]
                  labels = [target['labels'].to(self.device).float() for target in targets]
  
                  output = self.model(images, {'bbox': boxes, 'cls': labels,
                                               "img_scale": None,
                                               "img_size": None})
                  summary_loss.update_avg(self._output_dict(output), batch_size)
  
          return summary_loss
  