code: import copy
  import numpy
  import string
  import time
  import torch
  import tqdm
  from draugr.numpy_utilities import Split
  from draugr.python_utilities import (
      rgb_drop_alpha_batch_nhwc,
      torch_vision_normalize_batch_nchw,
  )
  from draugr.torch_utilities import (
      TorchEvalSession,
      TorchTrainSession,
      global_torch_device,
      to_tensor,
      uint_nhwc_to_nchw_float_batch,
  )
  from draugr.visualisation import confusion_matrix_plot
  from matplotlib import pyplot
  from munin.generate_report import ReportEntry, generate_html, generate_pdf
  from munin.utilities.html_embeddings import generate_math_html, plt_html
  from pathlib import Path
  from sklearn.metrics import accuracy_score, precision_recall_fscore_support
  from warg import NOD
  
  __all__ = ["test_model", "pred_target_train_model"]
  
  
  def test_model(model, data_iterator, latest_model_path, num_columns: int = 2):
      model = model.eval().to(global_torch_device())
  
      inputs, labels = next(data_iterator)
  
      inputs = inputs.to(global_torch_device())
      labels = labels.to(global_torch_device())
      with torch.no_grad():
          pred = model(inputs)
  
      y_pred = pred.data.to("cpu").numpy()
      y_pred_max = numpy.argmax(y_pred, axis=-1)
      accuracy_w = accuracy_score(labels, y_pred_max)
      precision_a, recall_a, fscore_a, support_a = precision_recall_fscore_support(
          labels, y_pred_max
      )
      precision_w, recall_w, fscore_w, support_w = precision_recall_fscore_support(
          labels, y_pred_max, average="weighted"
      )
  
      _, predicted = torch.max(pred, 1)
  
      truth_labels = labels.data.to("cpu").numpy()
  
      input_images_rgb = [
          default_torch_retransform(x) for x in inputs.to(global_torch_device())
      ]
  
      cell_width = (800 / num_columns) - 6 - 6 * 2
  
      pyplot.plot(numpy.random.random((3, 3)))
  
      alphabet = string.ascii_lowercase
      class_names = numpy.array([*alphabet])
  
      samples = len(y_pred)
      predictions = [
          [None for _ in range(num_columns)] for _ in range(samples // num_columns)
      ]
      for i, a, b, c in zip(range(samples), input_images_rgb, y_pred_max, truth_labels):
          pyplot.imshow(a)
          if b == c:
              outcome = "tp"
          else:
              outcome = "fn"
  
          gd = ReportEntry(
              name=i,
              figure=plt_html(a, format="jpg", size=(cell_width, cell_width)),
              prediction=class_names[b],
              truth=class_names[c],
              outcome=outcome,
              explanation=None,
          )
  
          predictions[i // num_columns][i % num_columns] = gd
  
      cfmat = confusion_matrix_plot(y_pred_max, truth_labels, class_names)
  
      title = "Classification Report"
      model_name = latest_model_path
      confusion_matrix = plt_html(cfmat, format="png", size=(800, 800))
  
      accuracy = generate_math_html("\dfrac{tp+tn}{N}"), None, accuracy_w
      precision = generate_math_html("\dfrac{tp}{tp+fp}"), precision_a, precision_w
      recall = generate_math_html("\dfrac{tp}{tp+fn}"), recall_a, recall_w
      f1_score = (
          generate_math_html("2*\dfrac{precision*recall}{precision+recall}"),
          fscore_a,
          fscore_w,
      )
      support = generate_math_html("N_{class_truth}"), support_a, support_w
      metrics = NOD.nod_of(
          accuracy, precision, f1_score, recall, support
      ).as_flat_tuples()
  
      bundle = NOD.nod_of(title, model_name, confusion_matrix, metrics, predictions)
  
      file_name = Path(title.lower().replace(" ", "_"))
  
      generate_html(file_name.with_suffix(".html"), **bundle)
      generate_pdf(file_name.with_suffix(".html"), file_name.with_suffix(".pdf"))
  
      # plot_utilities.plot_prediction(input_images_rgb, truth_labels, predicted, pred)
      # pyplot.show()
  
  
  def pred_target_train_model(
          model,
          train_iterator,
          criterion,
          optimizer,
          scheduler,
          writer,
          interrupted_path,
          test_data_iterator=None,
          num_updates: int = 250000,
          early_stop=None,
  ) -> torch.nn.Module:
      """
  
      Args:
        model:
        train_iterator:
        criterion:
        optimizer:
        scheduler:
        writer:
        interrupted_path:
        test_data_iterator:
        num_updates:
        early_stop:
  
      Returns:
  
      """
      best_model_wts = copy.deepcopy(model.state_dict())
      best_val_loss = 1e10
      since = time.time()
  
      try:
          sess = tqdm.tqdm(range(num_updates), leave=False, disable=False)
          val_loss = 0
          update_loss = 0
          val_acc = 0
          last_val = None
          last_out = None
          with torch.autograd.detect_anomaly():
              for update_i in sess:
                  for phase in [Split.Training, Split.Validation]:
                      if phase == Split.Training:
                          with TorchTrainSession(model):
  
                              input, true_label = zip(*next(train_iterator))
  
                              rgb_imgs = torch_vision_normalize_batch_nchw(
                                  uint_nhwc_to_nchw_float_batch(
                                      rgb_drop_alpha_batch_nhwc(to_tensor(input))
                                  )
                              )
                              true_label = to_tensor(true_label, dtype=torch.long)
                              optimizer.zero_grad()
  
                              pred = model(rgb_imgs)
                              loss = criterion(pred, true_label)
                              loss.backward()
                              optimizer.step()
  
                              if last_out is None:
                                  last_out = pred
                              else:
                                  if not torch.dist(last_out, pred) > 0:
                                      print(f"Same output{last_out},{pred}")
                                  last_out = pred
  
                              update_loss = loss.data.cpu().numpy()
                              writer.scalar(f"loss/train", update_loss, update_i)
  
                              if scheduler:
                                  scheduler.step()
                      elif test_data_iterator:
                          with TorchEvalSession(model):
                              test_rgb_imgs, test_true_label = zip(*next(train_iterator))
                              test_rgb_imgs = torch_vision_normalize_batch_nchw(
                                  uint_nhwc_to_nchw_float_batch(
                                      rgb_drop_alpha_batch_nhwc(to_tensor(test_rgb_imgs))
                                  )
                              )
  
                              test_true_label = to_tensor(
                                  test_true_label, dtype=torch.long
                              )
  
                              with torch.no_grad():
                                  val_pred = model(test_rgb_imgs)
                                  val_loss = criterion(val_pred, test_true_label)
  
                              _, cat = torch.max(val_pred, -1)
                              val_acc = torch.sum(cat == test_true_label) / float(
                                  cat.size(0)
                              )
                              writer.scalar(f"loss/acc", val_acc, update_i)
                              writer.scalar(f"loss/val", val_loss, update_i)
  
                              if last_val is None:
                                  last_val = cat
                              else:
                                  if all(last_val == cat):
                                      print(f"Same val{last_val},{cat}")
                                  last_val = cat
  
                              if val_loss < best_val_loss:
                                  best_val_loss = val_loss
  
                                  best_model_wts = copy.deepcopy(model.state_dict())
                                  sess.write(
                                      f"New best validation model at update {update_i} with test_loss {best_val_loss}"
                                  )
                                  torch.save(model.state_dict(), interrupted_path)
  
                          if early_stop is not None and val_pred < early_stop:
                              break
                  sess.set_description_str(
                      f"Update {update_i} - {phase} "
                      f"update_loss:{update_loss:2f} "
                      f"test_loss:{val_loss}"
                      f"val_acc:{val_acc}"
                  )
  
      except KeyboardInterrupt:
          print("Interrupt")
      finally:
          pass
  
      model.load_state_dict(best_model_wts)  # load best model weights
  
      time_elapsed = time.time() - since
      print(f"{time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s")
      print(f"Best val loss: {best_val_loss:3f}")
  
      return model
  
apis: torch.sum sklearn.metrics.precision_recall_fscore_support torch.autograd.detect_anomaly torch.no_grad numpy.argmax numpy.random.random sklearn.metrics.accuracy_score matplotlib.pyplot.imshow torch.dist torch.max numpy.array