code: from __future__ import print_function
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from tqdm import tqdm
  
  
  class Trainer():
  
      def __init__(self, model, device, train_loader, test_loader, optimizer, loss_func, lr_scheduler):
          self.is_last_epoch = False
  #         self.train_losses = []  # detailed training loss
          self.test_losses = []
  #         self.train_acc = []  # detailed  training accuracy
          self.train_acc_total = []  # per epoch training accuracy
          self.train_loss_total = []  # per epoch train loss
          self.test_acc = []
          self.model = model
          self.device = device
          self.train_loader = train_loader
          self.test_loader = test_loader
          self.optimizer = optimizer
          self.loss_func = loss_func
          self.lr_scheduler = lr_scheduler
  
      def train_model(self, lambda_l1, epochs=5):
          for epoch in range(epochs):
              print("\nCurrent EPOCH:", epoch)
  
              if self.lr_scheduler is not None:
                  if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                      print("Current EPOCH:", epoch, "last LR=", self.lr_scheduler.get_last_lr(), "LR = ", self.lr_scheduler.get_lr())
                  else:
                      print("Learning Rate(ReduceLROnPlateau) = ", self.optimizer.param_groups[0]['lr'])
                      # https://discuss.pytorch.org/t/how-to-retrieve-learning-rate-from-reducelronplateau-scheduler/54234
  
              self.train(epoch, lambda_l1)
              self.is_last_epoch = epoch == epochs
              tst_metric = self.test()
  
              if self.lr_scheduler is not None:
                  if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                      val_loss = tst_metric # test_losses[-1]
                      print("ReduceLROnPlateau, ReduceLROnPlateau::step(), val_loss", val_loss)
                      self.lr_scheduler.step(val_loss)
  
                  if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.StepLR):
                      self.lr_scheduler.step()
  
  
          return (self.train_loss_total, self.train_acc_total, self.test_losses, self.test_acc)
  
  #    def get_detailed_train_stats(self):
  #        return (self.train_losses, self.train_acc)
  
      def train(self, epoch, lambda_l1):
          self.model.train()
          pbar = tqdm(self.train_loader)
          correct = 0
          processed = 0
          loss = 0
          for batch_idx, (data, target) in enumerate(pbar):
              # get samples
              data, target = data.to(self.device), target.to(self.device)
  
              # Init
              self.optimizer.zero_grad()
              # In PyTorch, we need to set the gradients to zero before starting to do backpropragation because PyTorch accumulates the gradients on subsequent backward passes.
              # Because of this, when you start your training loop, ideally you should zero out the gradients so that you do the parameter update correctly.
  
              # Predict
              y_pred = self.model(data)
  
              # Calculate loss
              loss = self.loss_func(y_pred, target)
  
              # L2 loss
  
              # L1 loss
              loss_l1 = 0
              # lambda_l1 = 0.05
              if lambda_l1 > 0:
                  for p in self.model.parameters():
                      loss_l1 = loss_l1 + p.abs().sum()
                  loss = loss + lambda_l1 * loss_l1
  
  #             self.train_losses.append(loss)
  
              # Backpropagation
              loss.backward()
              self.optimizer.step()
  
              pred = y_pred.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
              correct += pred.eq(target.view_as(pred)).sum().item()
              processed += len(data)
  
              pbar.set_description(
                  desc=f'Train set: Loss={loss.item()} Batch_id={batch_idx} Accuracy={100 * correct / processed:0.2f}')
  #             self.train_acc.append(100 * correct / processed)
  
              if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.OneCycleLR):
                  self.lr_scheduler.step()
  
  
          training_accuracy_perepoch = 100 * correct / processed
          self.train_acc_total.append(training_accuracy_perepoch)
          self.train_loss_total.append(loss)
  
      def test(self):
          self.model.eval()
          test_loss = 0
          correct = 0
          with torch.no_grad():
              for data, target in self.test_loader:
                  data, target = data.to(self.device), target.to(self.device)
                  output = self.model(data)
                  test_loss += self.loss_func(output, target).item()  # sum up batch loss
                  pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
                  is_correct = pred.eq(target.view_as(pred))
                  correct += is_correct.sum().item()
  
          test_loss /= len(self.test_loader.dataset)
          self.test_losses.append(test_loss)
  
          print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
              test_loss, correct, len(self.test_loader.dataset),
              100. * correct / len(self.test_loader.dataset)))
  
          self.test_acc.append(100. * correct / len(self.test_loader.dataset))
          return test_loss
  
      def getValues(self):
          return (self.train_loss_total, self.train_acc_total, self.test_losses, self.test_acc)
  
      def get_misclassified(self):
          self.model.eval()
          misclassified_imgs = []
          test_loss = 0
          correct = 0
          with torch.no_grad():
              for data, target in self.test_loader:
                  data, target = data.to(self.device), target.to(self.device)
                  output = self.model(data)
                  pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
                  is_correct = pred.eq(target.view_as(pred))
                  if True:
                      misclassified_inds = (is_correct == 0).nonzero()[:, 0]
                      for mis_ind in misclassified_inds:
                          if len(misclassified_imgs) == 25:
                              break
                          misclassified_imgs.append({
                              "target": target[mis_ind].cpu().numpy(),
                              "pred": pred[mis_ind][0].cpu().numpy(),
                              "img": data[mis_ind].cpu().numpy()
                          })
                  correct += is_correct.sum().item()
  
          return misclassified_imgs
  
      def classwise_acc(self, classes):
          class_correct = list(0. for i in range(10))
          class_total = list(0. for i in range(10))
          with torch.no_grad():
              for images, labels in self.test_loader:
                  images, labels = images.to(self.device), labels.to(self.device)
                  outputs = self.model(images)
                  _, predicted = torch.max(outputs, 1)
                  c = (predicted == labels).squeeze()
                  for i in range(4):
                      label = labels[i]
                      class_correct[label] += c[i].item()
                      class_total[label] += 1
  
          # print class-wise test accuracies
          print()
          for i in range(10):
              print('Accuracy of %5s : %2d %%' % (
                  classes[i], 100 * class_correct[i] / class_total[i]))
          print()
  
apis: torch.no_grad torch.max