class_docstr: This wrapper handles torch models for image classification.
  
  It can either support models from torchvision.models (resnet18, alexnet, ...) or a
  simple CNN for MNIST (simple-cnn, see SimpleCnn class).
class_docstr_tok: ["This","wrapper","handles","torch","models","for","image","classification",".","It","can","either","support","models","from","torchvision.models","(","resnet18",",","alexnet",",","...",")","or","a","simple","CNN","for","MNIST","(","simple-cnn",",","see","SimpleCnn","class",")","."]
code: def _train(
      self,
      train_data,
      val_data,
      test_data,
      writer,
      experiment,
      dry_run: bool = False,
  ) -> None:
      use_cuda = torch.cuda.is_available()
  
      # Preprocess all datasets.
      logger.info("Preprocessing datasets...")
      train_loader = self._preprocess_for_training(
          "train", train_data, use_cuda
      )
      val_loader = self._preprocess_for_training(
          "val", val_data, use_cuda
      )
      test_loader = self._preprocess_for_training(
          "test", test_data, use_cuda
      )
      logger.info("")
  
      # Set up model and move it to device.
      logger.info("Creating model...")
      self._create_model(self.num_classes)
      device = torch.device("cuda" if use_cuda else "cpu")
      logger.info(f"    device: {device}")
      self.model = self.model.to(device)
  
      # Set up optimizer and loss.
      optimizer = self._create_optimizer()
      loss_func = nn.CrossEntropyLoss()
      logger.info(f"    loss function: cross-entropy")
      logger.info("")
  
      # Dedicate a few images that will be plotted as samples to tensorboard.
      num_samples_to_plot = self.config.get("num_samples_to_plot", 5)
  
      def get_samples(loader):
          if loader is None:
              return None, None
          else:
              return next(
                  iter(
                      DataLoader(
                          loader.dataset, batch_size=num_samples_to_plot
                      )
                  )
              )
  
      train_sample_images, train_sample_labels = get_samples(
          train_loader
      )
      val_sample_images, val_sample_labels = get_samples(val_loader)
      test_sample_images, test_sample_labels = get_samples(test_loader)
  
      # Configure trainer and metrics.
      accumulate_train_metrics = self.config.get(
          "accumulate_train_metrics", True
      )
  
      # We need to transform the output of the trainer and metrics here to accumulate
      # metrics during training (otherwise, we have to re-evaluate on the complete
      # train set which takes a long time). By default, the trainer outputs
      # `loss.item()` and the metrics expect `y_pred, y` (which is what the evaluator
      # outputs). We are now outputting `y_pred, y, loss` from the trainer and then
      # slicing off the `loss` before it goes into the metric.
      # See also the footnote here but note that it's a bit wrong:
      # https://pytorch.org/ignite/quickstart.html#
      def trainer_output_transform(x, y, y_pred, loss):
          return y_pred, y, loss.item()
  
      def metrics_output_transform(output):
          return output[:2]  # use only y_pred, y
  
      trainer = create_supervised_trainer(
          self.model,
          optimizer,
          loss_func,
          device=device,
          output_transform=trainer_output_transform,
      )
      if accumulate_train_metrics:
          # TODO: Maybe put train_metrics and val_metrics into one dict.
          train_metrics = {
              "accuracy": Accuracy(
                  output_transform=metrics_output_transform
              ),
              "loss": Loss(
                  loss_func, output_transform=metrics_output_transform
              ),
              # "confusion_matrix": ConfusionMatrix(num_classes),
          }
          for name, metric in train_metrics.items():
              # Attach metrics to trainer to accumulate them during training.
              metric.attach(trainer, name)
      val_metrics = {
          "accuracy": Accuracy(),
          "loss": Loss(loss_func),
          # "confusion_matrix": ConfusionMatrix(num_classes),
      }
      evaluator = create_supervised_evaluator(
          self.model, metrics=val_metrics, device=device
      )
  
      @trainer.on(
          Events.ITERATION_COMPLETED(
              every=self.config.get("print_every", 100)
          )
      )
      def log_batch(trainer):
          batch = (
              trainer.state.iteration - 1
          ) % trainer.state.epoch_length + 1
          logger.info(
              f"Epoch {trainer.state.epoch} / {num_epochs}, "
              f"batch {batch} / {trainer.state.epoch_length}: "
              f"Loss: {trainer.state.output[2]:.3f}"
              # f"Loss: {trainer.state.output:.3f}"
          )
  
      def log_results(name, metrics, epoch):
          """Log results of an epoch to stdout, tensorboard and comet."""
          logger.info(
              f"{name}: Average loss: {metrics['loss']:.3f}, "
              f"Average accuracy: {metrics['accuracy']:.3f}"
          )
          experiment.log_metric(f"{name}_loss", metrics["loss"])
          experiment.log_metric(f"{name}_accuracy", metrics["accuracy"])
          writer.add_scalar(f"{name}_loss", metrics["loss"], epoch)
          writer.add_scalar(
              f"{name}_accuracy", metrics["accuracy"], epoch
          )
  
      # TODO: This iterates over complete train set again, maybe accumulate as in the
      #   example in the footnote here: https://pytorch.org/ignite/quickstart.html#
      @trainer.on(Events.EPOCH_COMPLETED)
      def log_epoch(trainer):
          logger.info("")
          logger.info(
              f"Epoch {trainer.state.epoch} / {num_epochs} results: "
          )
  
          # Train data.
          if accumulate_train_metrics:
              log_results(
                  "train", trainer.state.metrics, trainer.state.epoch
              )
              logger.info(
                  "(train metrics are accumulated during training; "
                  "to re-evaluate on the complete train set after training, "
                  "use config parameter 'accumulate_train_metrics': False)"
              )
          else:
              evaluator.run(train_loader)
              log_results(
                  "train", evaluator.state.metrics, trainer.state.epoch
              )
  
          # Val data.
          if val_loader:
              evaluator.run(val_loader)
              log_results(
                  "val", evaluator.state.metrics, trainer.state.epoch
              )
  
          # Test data.
          if test_loader:
              evaluator.run(test_loader)
              log_results(
                  "test", evaluator.state.metrics, trainer.state.epoch
              )
  
          logger.info("")
  
      @trainer.on(Events.EPOCH_COMPLETED)
      def checkpoint_model(trainer):
          # TODO: Do not checkpoint at every step.
          checkpoint_dir = (
              self.out_dir
              / "checkpoints"
              / f"epoch{trainer.state.epoch}"
          )
          checkpoint_dir.mkdir(parents=True, exist_ok=True)
          torch.save(self.model, checkpoint_dir / "model.pt")
  
      @trainer.on(Events.EPOCH_COMPLETED)
      def plot_samples(trainer):
          """Plot a few sample images and probabilites to tensorboard."""
  
          def write_samples_plot(name, sample_images, sample_labels):
              # TODO: This can be improved by just using the outputs already
              #   calculated in evaluator.state.output in the functions above.
              #   Problem: At least in the train evaluator, the batches are not equal,
              #   so the plotted images will differ from run to run.
              if sample_images is None:
                  return
  
              with torch.no_grad():
                  sample_output = self.model(sample_images.to(device))
                  sample_pred = torch.softmax(sample_output, dim=1)
  
              visualization.plot_samples(
                  writer,
                  f"{name}-samples",
                  trainer.state.epoch,
                  sample_images.to("cpu").numpy(),
                  sample_labels.to("cpu").numpy(),
                  sample_pred.to("cpu").numpy(),
              )
  
          write_samples_plot(
              "train", train_sample_images, train_sample_labels
          )
          write_samples_plot(
              "val", val_sample_images, val_sample_labels
          )
          write_samples_plot(
              "test", test_sample_images, test_sample_labels
          )
  
      # Start training.
      num_epochs = 1 if dry_run else self.config.get("num_epochs", 5)
      if dry_run:
          num_batches = 1
          logger.info(
              f"Training model on device {device}... (DRY RUN, only 1 batch)"
          )
      elif "num_samples" in self.config:
          # TODO: Make sure batch_size doesn't differ from the value extracted during
          #   preprocessing.
          batch_size = self.config.get("batch_size", 128)
          # TODO: This always uses a few more samples than num_samples. Maybe get it
          #   to the correct value.
          num_batches = int(self.config["num_samples"] / batch_size) + 1
          logger.info(
              f"Training model on device {device}... (using "
              f"{self.config['num_samples']} of {len(train_loader.dataset)} samples)"
          )
      else:
          num_batches = None  # all batches
          logger.info(f"Training model on device {device}...")
          logger.info(
              "(if this takes too long, train on less data with the config "
              "parameter 'num_samples')"
          )
      logger.info(
          "(show more steps by setting the config parameter 'print_every')"
      )
      logger.info("")
      trainer.run(
          train_loader, max_epochs=num_epochs, epoch_length=num_batches
      )
      logger.info("Training finished!")
  
      # Save the trained model.
      torch.save(self.model, self.out_dir / "model.pt")
  
imports: import comet_ml
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torch.utils.data import DataLoader
  import torchvision
  from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
  from ignite.metrics import Accuracy, Loss
  import numpy
  from loguru import logger
  from ..model_wrapper import ModelWrapper
  from .None import preprocessing, visualization
  from ..None import utils
label_desc: Use image node_id 4 for calling the TorchImageClassificationWrapper obj's underlying member method code with example usage: obj._train(train_data, val_data, test_data, writer, experiment, dry_run) and returns: y_pred,  y,  loss, output, None,  None, next