text: import numpy as np
  import torch
  from torch.autograd import Variable
  
  from base import BaseTrainer
  from utils.util import freeze_network, unfreeze_network
  
  
  class Trainer(BaseTrainer):
      """
      Trainer class
  
      Note:
          Inherited from BaseTrainer.
          self.optimizer is by default handled by BaseTrainer based on config.
      """
  
      def __init__(self, model, loss, metrics, resume, config,
                   data_loader, valid_data_loader=None, train_logger=None):
          super(Trainer, self).__init__(model, loss, metrics, resume, config, train_logger)
          self.config = config
          self.batch_size = data_loader.batch_size
          self.data_loader = data_loader
          self.valid_data_loader = valid_data_loader
          self.valid = True if self.valid_data_loader is not None else False
          self.log_step = int(np.sqrt(self.batch_size))
  
      def _to_tensor(self, data, target):
          data, target = torch.FloatTensor(data), torch.FloatTensor(target)
          if self.with_cuda:
              data, target = data.to(self.gpu), target.to(self.gpu)
          return data, target
  
      def _eval_metrics(self, output, target):
          acc_metrics = np.zeros(len(self.metrics))
          output = output.cpu().data.numpy()
          target = target.cpu().data.numpy()
          output = np.argmax(output, axis=1)
          for i, metric in enumerate(self.metrics):
              acc_metrics[i] += metric(output, target)
          return acc_metrics
  
      def _train_epoch(self, epoch):
          """
          Training logic for an epoch
  
          :param epoch: Current training epoch.
          :return: A log that contains all information you want to save.
  
          Note:
              If you have additional information to record, for example:
                  > additional_log = {"x": x, "y": y}
              merge it with log before return. i.e.
                  > log = {**log, **additional_log}
                  > return log
  
              The metrics in log must have the key 'metrics'.
          """
          self.model.train()
  
          # if epoch < 50:
          #     freeze_network(self.model.net)
          # else:
          #     unfreeze_network(self.model.net)
  
          total_loss = 0
          total_metrics = np.zeros(len(self.metrics))
          for batch_idx, (data, target) in enumerate(self.data_loader):
              data, target = self._to_tensor(data, target)
  
              data = Variable(data).cuda()
  
              self.optimizer.zero_grad()
              output = self.model(data)
  
              loss = self.loss(output, target)
              loss.backward()
              self.optimizer.step()
  
              total_loss += loss.item()
              total_metrics += self._eval_metrics(output, target)
  
              self.print_log(batch_idx, epoch, loss)
  
          log = {
              'loss': total_loss / len(self.data_loader),
              'metrics': (total_metrics / len(self.data_loader)).tolist()
          }
  
          if self.valid:
              val_log = self._valid_epoch()
              log = {**log, **val_log}
  
          return log
  
      def print_log(self, batch_idx, epoch, loss):
          if self.verbosity >= 2 and batch_idx % self.log_step == 0:
              self.logger.info('Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(
                  epoch,
                  batch_idx * self.data_loader.batch_size,
                  len(self.data_loader) * self.data_loader.batch_size,
                  100.0 * batch_idx / len(self.data_loader),
                  loss.item()))
  
      def _valid_epoch(self):
          """
          Validate after training an epoch
  
          :return: A log that contains information about validation
  
          Note:
              The validation metrics in log must have the key 'val_metrics'.
          """
          self.model.eval()
          total_val_loss = 0
          total_val_metrics = np.zeros(len(self.metrics))
          with torch.no_grad():
              for batch_idx, (data, target) in enumerate(self.valid_data_loader):
                  data, target = self._to_tensor(data, target)
  
                  output = self.model(data)
                  loss = self.loss(output, target)
  
                  total_val_loss += loss.item()
                  total_val_metrics += self._eval_metrics(output, target)
  
          return {
              'val_loss': total_val_loss / len(self.valid_data_loader),
              'val_metrics': (total_val_metrics / len(self.valid_data_loader)).tolist()
          }
  