class_docstr: TRA Model
  
  Args:
      model_config (dict): model config (will be used by RNN or Transformer)
      tra_config (dict): TRA config (will be used by TRA)
      model_type (str): which backbone model to use (RNN/Transformer)
      lr (float): learning rate
      n_epochs (int): number of total epochs
      early_stop (int): early stop when performance not improved at this step
      update_freq (int): gradient update frequency
      max_steps_per_epoch (int): maximum number of steps in one epoch
      lamb (float): regularization parameter
      rho (float): exponential decay rate for `lamb`
      alpha (float): fusion parameter for calculating transport loss matrix
      seed (int): random seed
      logdir (str): local log directory
      eval_train (bool): whether evaluate train set between epochs
      eval_test (bool): whether evaluate test set between epochs
      pretrain (bool): whether pretrain the backbone model before training TRA.
          Note that only TRA will be optimized after pretraining
      init_state (str): model init state path
      freeze_model (bool): whether freeze backbone model parameters
      freeze_predictors (bool): whether freeze predictors parameters
      transport_method (str): transport method, can be none/router/oracle
      memory_mode (str): memory mode, the same argument for MTSDatasetH
class_docstr_tok: ["TRA","Model","Args",":","model_config","(","dict",")",":","model","config","(","will","be","used","by","RNN","or","Transformer",")","tra_config","(","dict",")",":","TRA","config","(","will","be","used","by","TRA",")","model_type","(","str",")",":","which","backbone","model","to","use","(","RNN\/Transformer",")","lr","(","float",")",":","learning","rate","n_epochs","(","int",")",":","number","of","total","epochs","early_stop","(","int",")",":","early","stop","when","performance","not","improved","at","this","step","update_freq","(","int",")",":","gradient","update","frequency","max_steps_per_epoch","(","int",")",":","maximum","number","of","steps","in","one","epoch","lamb","(","float",")",":","regularization","parameter","rho","(","float",")",":","exponential","decay","rate","for","`","lamb","`","alpha","(","float",")",":","fusion","parameter","for","calculating","transport","loss","matrix","seed","(","int",")",":","random","seed","logdir","(","str",")",":","local","log","directory","eval_train","(","bool",")",":","whether","evaluate","train","set","between","epochs","eval_test","(","bool",")",":","whether","evaluate","test","set","between","epochs","pretrain","(","bool",")",":","whether","pretrain","the","backbone","model","before","training","TRA",".","Note","that","only","TRA","will","be","optimized","after","pretraining","init_state","(","str",")",":","model","init","state","path","freeze_model","(","bool",")",":","whether","freeze","backbone","model","parameters","freeze_predictors","(","bool",")",":","whether","freeze","predictors","parameters","transport_method","(","str",")",":","transport","method",",","can","be","none\/router\/oracle","memory_mode","(","str",")",":","memory","mode",",","the","same","argument","for","MTSDatasetH"]
code: def train_epoch(self, epoch, data_set, is_pretrain=False):
      self.model.train()
      self.tra.train()
      data_set.train()
      self.optimizer.zero_grad()
  
      P_all = []
      prob_all = []
      choice_all = []
      max_steps = len(data_set)
      if self.max_steps_per_epoch is not None:
          if epoch == 0 and self.max_steps_per_epoch < max_steps:
              self.logger.info(
                  f"max steps updated from {max_steps} to {self.max_steps_per_epoch}"
              )
          max_steps = min(self.max_steps_per_epoch, max_steps)
  
      cur_step = 0
      total_loss = 0
      total_count = 0
      for batch in tqdm(data_set, total=max_steps):
          cur_step += 1
          if cur_step > max_steps:
              break
  
          if not is_pretrain:
              self.global_step += 1
  
          data, state, label, count = (
              batch["data"],
              batch["state"],
              batch["label"],
              batch["daily_count"],
          )
          index = (
              batch["daily_index"]
              if self.use_daily_transport
              else batch["index"]
          )
  
          with torch.set_grad_enabled(not self.freeze_model):
              hidden = self.model(data)
  
          all_preds, choice, prob = self.tra(hidden, state)
  
          if is_pretrain or self.transport_method != "none":
              # NOTE: use oracle transport for pre-training
              loss, pred, L, P = self.transport_fn(
                  all_preds,
                  label,
                  choice,
                  prob,
                  state.mean(dim=1),
                  count,
                  self.transport_method
                  if not is_pretrain
                  else "oracle",
                  self.alpha,
                  training=True,
              )
              data_set.assign_data(index, L)  # save loss to memory
              if (
                  self.use_daily_transport
              ):  # only save for daily transport
                  P_all.append(
                      pd.DataFrame(
                          P.detach().cpu().numpy(), index=index
                      )
                  )
                  prob_all.append(
                      pd.DataFrame(
                          prob.detach().cpu().numpy(), index=index
                      )
                  )
                  choice_all.append(
                      pd.DataFrame(
                          choice.detach().cpu().numpy(), index=index
                      )
                  )
              decay = self.rho ** (
                  self.global_step // 100
              )  # decay every 100 steps
              lamb = 0 if is_pretrain else self.lamb * decay
              reg = (
                  prob.log().mul(P).sum(dim=1).mean()
              )  # train router to predict TO assignment
              if self._writer is not None and not is_pretrain:
                  self._writer.add_scalar(
                      "training/router_loss",
                      -reg.item(),
                      self.global_step,
                  )
                  self._writer.add_scalar(
                      "training/reg_loss", loss.item(), self.global_step
                  )
                  self._writer.add_scalar(
                      "training/lamb", lamb, self.global_step
                  )
                  if not self.use_daily_transport:
                      P_mean = P.mean(axis=0).detach()
                      self._writer.add_scalar(
                          "training/P",
                          P_mean.max() / P_mean.min(),
                          self.global_step,
                      )
              loss = loss - lamb * reg
          else:
              pred = all_preds.mean(dim=1)
              loss = loss_fn(pred, label)
  
          (loss / self.update_freq).backward()
          if cur_step % self.update_freq == 0:
              self.optimizer.step()
              self.optimizer.zero_grad()
  
          if self._writer is not None and not is_pretrain:
              self._writer.add_scalar(
                  "training/total_loss", loss.item(), self.global_step
              )
  
          total_loss += loss.item()
          total_count += 1
  
      if self.use_daily_transport and len(P_all) > 0:
          P_all = pd.concat(P_all, axis=0)
          prob_all = pd.concat(prob_all, axis=0)
          choice_all = pd.concat(choice_all, axis=0)
          P_all.index = data_set.restore_daily_index(P_all.index)
          prob_all.index = P_all.index
          choice_all.index = P_all.index
          if not is_pretrain:
              self._writer.add_image(
                  "P", plot(P_all), epoch, dataformats="HWC"
              )
              self._writer.add_image(
                  "prob", plot(prob_all), epoch, dataformats="HWC"
              )
              self._writer.add_image(
                  "choice", plot(choice_all), epoch, dataformats="HWC"
              )
  
      total_loss /= total_count
  
      if self._writer is not None and not is_pretrain:
          self._writer.add_scalar("training/loss", total_loss, epoch)
  
      return total_loss
  
imports: import io
  import os
  import copy
  import math
  import json
  import numpy
  import pandas
  import matplotlib.pyplot
  import torch
  import torch.nn
  import torch.optim
  import torch.nn.functional
  from tqdm import tqdm
  from qlib.constant import EPS
  from qlib.log import get_module_logger
  from qlib.model.base import Model
  from qlib.contrib.data.dataset import MTSDatasetH
label_desc: Use image node_id 3 for calling the TRAModel obj's underlying member method code with example usage: obj.train_epoch(epoch, data_set, is_pretrain) and returns: total_loss