class_docstr: Class for sound class embedding training"
      
class_docstr_tok: ["Class","for","sound","class","embedding","training","''"]
docstr: Gets called at the end of an epoch.
  
  Arguments
  ---------
  stage : sb.Stage
      One of sb.Stage.TRAIN, sb.Stage.VALID, sb.Stage.TEST
  stage_loss : float
      The average loss for all of the data processed in this stage.
  epoch : int
      The currently-starting epoch. This is passed
      `None` during the test stage.
docstr_tok: ["Gets","called","at","the","end","of","an","epoch",".","Arguments","--","--","--","--","-","stage",":","sb.Stage","One","of","sb.Stage.TRAIN",",","sb.Stage.VALID",",","sb.Stage.TEST","stage_loss",":","float","The","average","loss","for","all","of","the","data","processed","in","this","stage",".","epoch",":","int","The","currently-starting","epoch",".","This","is","passed","`","None","`","during","the","test","stage","."]
code: def on_stage_end(self, stage, stage_loss, epoch=None):
      """Gets called at the end of an epoch.
  
      Arguments
      ---------
      stage : sb.Stage
          One of sb.Stage.TRAIN, sb.Stage.VALID, sb.Stage.TEST
      stage_loss : float
          The average loss for all of the data processed in this stage.
      epoch : int
          The currently-starting epoch. This is passed
          `None` during the test stage.
      """
      # Compute/store important stats
      if stage == sb.Stage.TRAIN:
          self.train_loss = stage_loss
          self.train_stats = {
              "loss": self.train_loss,
              "acc": self.acc_metric.summarize(
                  "average"
              ),  # "acc": self.train_acc_metric.summarize(),
          }
      # Summarize Valid statistics from the stage for record-keeping.
      elif stage == sb.Stage.VALID:
          valid_stats = {
              "loss": stage_loss,
              "acc": self.acc_metric.summarize(
                  "average"
              ),  # "acc": self.valid_acc_metric.summarize(),
              "error": self.error_metrics.summarize("average"),
          }
      # Summarize Test statistics from the stage for record-keeping.
      else:
          test_stats = {
              "loss": stage_loss,
              "acc": self.acc_metric.summarize(
                  "average"
              ),  # "acc": self.test_acc_metric.summarize(),
              "error": self.error_metrics.summarize("average"),
          }
  
      # Perform end-of-iteration things, like annealing, logging, etc.
      if stage == sb.Stage.VALID:
          old_lr, new_lr = self.hparams.lr_annealing(epoch)
          sb.nnet.schedulers.update_learning_rate(
              self.optimizer, new_lr
          )
  
          # Tensorboard logging
          if self.hparams.use_tensorboard:
              self.hparams.tensorboard_train_logger.log_stats(
                  stats_meta={"Epoch": epoch},
                  train_stats=self.train_stats,
                  valid_stats=valid_stats,
              )
              # Log confusion matrix fig to tensorboard
              cm_fig = create_cm_fig(
                  self.valid_confusion_matrix,
                  display_labels=list(
                      self.hparams.label_encoder.ind2lab.values()
                  ),
              )
              self.hparams.tensorboard_train_logger.writer.add_figure(
                  "Validation Confusion Matrix", cm_fig, epoch
              )  # TODO use global_step from writer
  
          # Per class accuracy from Validation confusion matrix
          per_class_acc_arr = np.diag(
              self.valid_confusion_matrix
          ) / np.sum(self.valid_confusion_matrix, axis=1)
          per_class_acc_arr_str = "\n" + "\n".join(
              "{:}: {:.3f}".format(
                  self.hparams.label_encoder.decode_ndim(class_id),
                  class_acc,
              )
              for class_id, class_acc in enumerate(per_class_acc_arr)
          )
  
          # The train_logger writes a summary to stdout and to the logfile.
          self.hparams.train_logger.log_stats(
              stats_meta={"epoch": epoch, "lr": old_lr},
              train_stats=self.train_stats,
              valid_stats=valid_stats,
          )
          # Save the current checkpoint and delete previous checkpoints,
          self.checkpointer.save_and_keep_only(
              meta=valid_stats, min_keys=["error"]
          )
  
      # We also write statistics about test data to stdout and to the logfile.
      if stage == sb.Stage.TEST:
          # Per class accuracy from Test confusion matrix
          per_class_acc_arr = np.diag(
              self.test_confusion_matrix
          ) / np.sum(self.test_confusion_matrix, axis=1)
          per_class_acc_arr_str = "\n" + "\n".join(
              "{:}: {:.3f}".format(class_id, class_acc)
              for class_id, class_acc in enumerate(per_class_acc_arr)
          )
  
          self.hparams.train_logger.log_stats(
              {
                  "Epoch loaded": self.hparams.epoch_counter.current,
                  "\n Per Class Accuracy": per_class_acc_arr_str,
                  "\n Confusion Matrix": "\n{:}\n".format(
                      self.test_confusion_matrix
                  ),
              },
              test_stats=test_stats,
          )
  
file_path: speechbrain/recipes/UrbanSound8k/SoundClassification/train.py
imports: import os
  import sys
  import torch
  import torchaudio
  import speechbrain
  from hyperpyyaml import load_hyperpyyaml
  from speechbrain.utils.distributed import run_on_main
  from urbansound8k_prepare import prepare_urban_sound_8k
  from sklearn.metrics import confusion_matrix
  import numpy
  from confusion_matrix_fig import create_cm_fig
label_desc: Use image node_id 4 for calling the UrbanSound8kBrain obj's underlying member method code with example usage: obj.on_stage_end(stage, stage_loss, epoch) without return types