class_docstr: Trainer for the AlphaTensor model. The trainer does not require an
  explicit loss since the loss is computed by the model itself. The trainer
  is responsible for both the training step and the acting one, storing
  acting performance in a buffer.
class_docstr_tok: ["Trainer","for","the","AlphaTensor","model",".","The","trainer","does","not","require","an","explicit","loss","since","the","loss","is","computed","by","the","model","itself",".","The","trainer","is","responsible","for","both","the","training","step","and","the","acting","one",",","storing","acting","performance","in","a","buffer","."]
docstr: Initializes the trainer.
  
  Args:
      model (AlphaTensorModel): The model to train.
      tensor_size (int): Flattened size of the matrices to be multiplied.
      n_steps (int): Number of steps used to get a single action out of
      a triplet.
      batch_size (int): Batch size.
      optimizer (torch.optim.Optimizer): The optimizer used to train the
      model.
      device (str): The name of the torch device used for training.
      len_data (int): Number of training samples used (both actor
      generated and synthetic).
      pct_synth (float): Initial percentage of synthetic samples used
      for training.
      n_synth_data (int): Number of synthetic training samples.
      limit_rank (int): Maximum rank for synthetically-generated
      matrices.
      n_cob (int): Number of change of basis (cob) used for a single
      training sample.
      cob_prob (float): Probability of applying a change of basis.
      data_augmentation (bool): Whether to randomly swap the last
      operation of an episode with another operation.
      loss_params (Tuple[float, float]): Alpha and Beta parameters used
      in the loss function.
      random_seed (int): Randomizing seed.
      checkpoint_dir (str): Directory used to store model checkpoints.
      checkpoint_data_dir (str): Directory used to store games as JSON
      files.
      extra_devices (List[str]): Extra devices names used for multi-GPU
      training.
docstr_tok: ["Initializes","the","trainer",".","Args",":","model","(","AlphaTensorModel",")",":","The","model","to","train",".","tensor_size","(","int",")",":","Flattened","size","of","the","matrices","to","be","multiplied",".","n_steps","(","int",")",":","Number","of","steps","used","to","get","a","single","action","out","of","a","triplet",".","batch_size","(","int",")",":","Batch","size",".","optimizer","(","torch.optim.Optimizer",")",":","The","optimizer","used","to","train","the","model",".","device","(","str",")",":","The","name","of","the","torch","device","used","for","training",".","len_data","(","int",")",":","Number","of","training","samples","used","(","both","actor","generated","and","synthetic",")",".","pct_synth","(","float",")",":","Initial","percentage","of","synthetic","samples","used","for","training",".","n_synth_data","(","int",")",":","Number","of","synthetic","training","samples",".","limit_rank","(","int",")",":","Maximum","rank","for","synthetically-generated","matrices",".","n_cob","(","int",")",":","Number","of","change","of","basis","(","cob",")","used","for","a","single","training","sample",".","cob_prob","(","float",")",":","Probability","of","applying","a","change","of","basis",".","data_augmentation","(","bool",")",":","Whether","to","randomly","swap","the","last","operation","of","an","episode","with","another","operation",".","loss_params","(","Tuple","[","float",",","float","]",")",":","Alpha","and","Beta","parameters","used","in","the","loss","function",".","random_seed","(","int",")",":","Randomizing","seed",".","checkpoint_dir","(","str",")",":","Directory","used","to","store","model","checkpoints",".","checkpoint_data_dir","(","str",")",":","Directory","used","to","store","games","as","JSON","files",".","extra_devices","(","List","[","str","]",")",":","Extra","devices","names","used","for","multi-GPU","training","."]
code: def __init__(
      self,
      model: AlphaTensorModel,
      tensor_size: int,
      n_steps: int,
      batch_size: int,
      optimizer: torch.optim.Optimizer,
      device: str,
      len_data: int,
      pct_synth: float,
      n_synth_data: int,
      limit_rank: int,
      n_cob: int,
      cob_prob: float,
      data_augmentation: bool,
      loss_params: Tuple[float, float] = None,
      random_seed: int = None,
      checkpoint_dir: str = None,
      checkpoint_data_dir: Path = None,
      extra_devices: List[str] = None,
  ):
      """Initializes the trainer.
  
      Args:
          model (AlphaTensorModel): The model to train.
          tensor_size (int): Flattened size of the matrices to be multiplied.
          n_steps (int): Number of steps used to get a single action out of
          a triplet.
          batch_size (int): Batch size.
          optimizer (torch.optim.Optimizer): The optimizer used to train the
          model.
          device (str): The name of the torch device used for training.
          len_data (int): Number of training samples used (both actor
          generated and synthetic).
          pct_synth (float): Initial percentage of synthetic samples used
          for training.
          n_synth_data (int): Number of synthetic training samples.
          limit_rank (int): Maximum rank for synthetically-generated
          matrices.
          n_cob (int): Number of change of basis (cob) used for a single
          training sample.
          cob_prob (float): Probability of applying a change of basis.
          data_augmentation (bool): Whether to randomly swap the last
          operation of an episode with another operation.
          loss_params (Tuple[float, float]): Alpha and Beta parameters used
          in the loss function.
          random_seed (int): Randomizing seed.
          checkpoint_dir (str): Directory used to store model checkpoints.
          checkpoint_data_dir (str): Directory used to store games as JSON
          files.
          extra_devices (List[str]): Extra devices names used for multi-GPU
          training.
      """
      self.model = model
      self.optimizer = optimizer
      self.device = device
      self.dataset = TensorGameDataset(
          len_data,
          pct_synth,
          tensor_size,
          n_synth_data,
          limit_rank,
          f_prob_distribution,
          device=device,
          n_steps=n_steps,
          action_memory_len=(model.tensor_length - 1),
          random_seed=random_seed,
      )
      self.batch_size = batch_size
      self.max_rank = limit_rank
      if loss_params is None:
          self.alpha = 1
          self.beta = 1
      else:
          self.alpha, self.beta = loss_params
      self.checkpoint_dir = Path(
          checkpoint_dir if checkpoint_dir else BASE_CHECKPOINT_DIR
      )
      self.checkpoint_dir.mkdir(exist_ok=True, parents=True)
      self.checkpoint_data_dir = (
          checkpoint_data_dir
          if checkpoint_data_dir
          else Path(BASE_CHECKPOINT_DATA_DIR)
      )
      self.checkpoint_data_dir.mkdir(exist_ok=True, parents=True)
      self.change_of_basis = ChangeOfBasis(
          tensor_size, n_cob, cob_prob, device, random_seed
      )
      self.data_augmentation = data_augmentation
      self.extra_devices = extra_devices
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/training.py
imports: from pathlib import Path
  from typing import Tuple, List
  import torch.optim
  import tqdm
  from torch.utils.data import DataLoader
  from open_alpha_tensor.config import BASE_CHECKPOINT_DATA_DIR, BASE_CHECKPOINT_DIR
  from open_alpha_tensor.core.actors.stage import actor_prediction
  from open_alpha_tensor.core.data.basis_change import ChangeOfBasis
  from open_alpha_tensor.core.data.dataset import TensorGameDataset
  from open_alpha_tensor.core.data.generation import f_prob_distribution
  from open_alpha_tensor.core.data.utils import map_action_to_triplet
  from open_alpha_tensor.core.modules.alpha_tensor import AlphaTensorModel
label_desc: Use image node_id 1 to create a new Trainer object  with example: obj = Trainer(model, tensor_size, n_steps, batch_size, optimizer, device, len_data, pct_synth, n_synth_data, limit_rank, n_cob, cob_prob, data_augmentation, loss_params, random_seed, checkpoint_dir, checkpoint_data_dir, extra_devices)