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 _init_model(self):
      self.logger.info("init TRAModel...")
  
      self.model = eval(self.model_type)(**self.model_config).to(device)
      print(self.model)
  
      self.tra = TRA(self.model.output_size, **self.tra_config).to(
          device
      )
      print(self.tra)
  
      if self.init_state:
          self.logger.warning(f"load state dict from `init_state`")
          state_dict = torch.load(self.init_state, map_location="cpu")
          self.model.load_state_dict(state_dict["model"])
          res = load_state_dict_unsafe(self.tra, state_dict["tra"])
          self.logger.warning(str(res))
  
      if self.reset_router:
          self.logger.warning(f"reset TRA.router parameters")
          self.tra.fc.reset_parameters()
          self.tra.router.reset_parameters()
  
      if self.freeze_model:
          self.logger.warning(f"freeze model parameters")
          for param in self.model.parameters():
              param.requires_grad_(False)
  
      if self.freeze_predictors:
          self.logger.warning(f"freeze TRA.predictors parameters")
          for param in self.tra.predictors.parameters():
              param.requires_grad_(False)
  
      self.logger.info(
          "# model params: %d"
          % sum(
              p.numel()
              for p in self.model.parameters()
              if p.requires_grad
          )
      )
      self.logger.info(
          "# tra params: %d"
          % sum(
              p.numel()
              for p in self.tra.parameters()
              if p.requires_grad
          )
      )
  
      self.optimizer = optim.Adam(
          list(self.model.parameters()) + list(self.tra.parameters()),
          lr=self.lr,
      )
  
      self.fitted = False
      self.global_step = -1
  
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 2 for calling the TRAModel obj's underlying member method code with example usage: obj._init_model() without return types