class_docstr: A training loop.
class_docstr_tok: ["A","training","loop","."]
docstr: Train the KGE model.
  
  .. note ::
  
      Gradient clipping is a technique to avoid the exploding gradient problem. Clip by norm and clip by value
      are two alternative implementations.
  
  :param triples_factory:
      The training triples.
  :param num_epochs:
      The number of epochs to train the model.
  :param batch_size:
      If set the batch size to use for mini-batch training. Otherwise find the largest possible batch_size
      automatically.
  :param slice_size: >0
      The divisor for the scoring function when using slicing. This is only possible for LCWA training loops in
      general and only for models that have the slicing capability implemented.
  :param label_smoothing: (0 <= label_smoothing < 1)
      If larger than zero, use label smoothing.
  :param sampler: (None or 'schlichtkrull')
      The type of sampler to use. At the moment sLCWA in R-GCN is the only user of schlichtkrull sampling.
  :param continue_training:
      If set to False, (re-)initialize the model's weights. Otherwise continue training.
  :param only_size_probing:
      The evaluation is only performed for two batches to test the memory footprint, especially on GPUs.
  :param use_tqdm: Should a progress bar be shown for epochs?
  :param use_tqdm_batch: Should a progress bar be shown for batching (inside the epoch progress bar)?
  :param tqdm_kwargs:
      Keyword arguments passed to :mod:`tqdm` managing the progress bar.
  :param stopper:
      An instance of :class:`pykeen.stopper.EarlyStopper` with settings for checking
      if training should stop early
  :param sub_batch_size:
      If provided split each batch into sub-batches to avoid memory issues for large models / small GPUs.
  :param num_workers:
      The number of child CPU workers used for loading data. If None, data are loaded in the main process.
  :param clear_optimizer:
      Whether to delete the optimizer instance after training (as the optimizer might have additional memory
      consumption due to e.g. moments in Adam).
  :param checkpoint_directory:
      An optional directory to store the checkpoint files. If None, a subdirectory named ``checkpoints`` in the
      directory defined by :data:`pykeen.constants.PYKEEN_HOME` is used. Unless the environment variable
      ``PYKEEN_HOME`` is overridden, this will be ``~/.pykeen/checkpoints``.
  :param checkpoint_name:
      The filename for saving checkpoints. If the given filename exists already, that file will be loaded and used
      to continue training.
  :param checkpoint_frequency:
      The frequency of saving checkpoints in minutes. Setting it to 0 will save a checkpoint after every epoch.
  :param checkpoint_on_failure:
      Whether to save a checkpoint in cases of a RuntimeError or MemoryError. This option differs from ordinary
      checkpoints, since ordinary checkpoints are only saved after a successful epoch. When saving checkpoints
      due to failure of the training loop there is no guarantee that all random states can be recovered correctly,
      which might cause problems with regards to the reproducibility of that specific training loop. Therefore,
      these checkpoints are saved with a distinct checkpoint name, which will be
      ``PyKEEN_just_saved_my_day_{datetime}.pt`` in the given checkpoint_root.
  :param drop_last:
      Whether to drop the last batch in each epoch to prevent smaller batches. Defaults to False, except if the
      model contains batch normalization layers. Can be provided explicitly to override.
  :param callbacks:
      An optional :class:`pykeen.training.TrainingCallback` or collection of callback instances that define
      one of several functionalities. Their interface was inspired by Keras.
  :param callbacks_kwargs:
      additional keyword-based parameter to instantiate the training callback.
  :param gradient_clipping_max_norm:
      The maximum gradient norm for use with gradient clipping. If None, no gradient norm clipping is used.
  :param gradient_clipping_norm_type:
      The gradient norm type to use for maximum gradient norm, cf. :func:`torch.nn.utils.clip_grad_norm_`
  :param gradient_clipping_max_abs_value:
      The maximum absolute value in gradients, cf. :func:`torch.nn.utils.clip_grad_value_`. If None, no
      gradient clipping will be used.
  :param pin_memory:
      whether to use memory pinning in the data loader, cf.
      https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning
  
  :return:
      The losses per epoch.
docstr_tok: ["Train","the","KGE","model",".","..","note",":",":","Gradient","clipping","is","a","technique","to","avoid","the","exploding","gradient","problem",".","Clip","by","norm","and","clip","by","value","are","two","alternative","implementations",".",":","param","triples_factory",":","The","training","triples",".",":","param","num_epochs",":","The","number","of","epochs","to","train","the","model",".",":","param","batch_size",":","If","set","the","batch","size","to","use","for","mini-batch","training",".","Otherwise","find","the","largest","possible","batch_size","automatically",".",":","param","slice_size",":",">","0","The","divisor","for","the","scoring","function","when","using","slicing",".","This","is","only","possible","for","LCWA","training","loops","in","general","and","only","for","models","that","have","the","slicing","capability","implemented",".",":","param","label_smoothing",":","(","0","<","=","label_smoothing","<","1",")","If","larger","than","zero",",","use","label","smoothing",".",":","param","sampler",":","(","None","or","'schlichtkrull","'",")","The","type","of","sampler","to","use",".","At","the","moment","sLCWA","in","R-GCN","is","the","only","user","of","schlichtkrull","sampling",".",":","param","continue_training",":","If","set","to","False",",","(","re-",")","initialize","the","model","'s","weights",".","Otherwise","continue","training",".",":","param","only_size_probing",":","The","evaluation","is","only","performed","for","two","batches","to","test","the","memory","footprint",",","especially","on","GPUs",".",":","param","use_tqdm",":","Should","a","progress","bar","be","shown","for","epochs","?",":","param","use_tqdm_batch",":","Should","a","progress","bar","be","shown","for","batching","(","inside","the","epoch","progress","bar",")","?",":","param","tqdm_kwargs",":","Keyword","arguments","passed","to",":","mod",":","`","tqdm","`","managing","the","progress","bar",".",":","param","stopper",":","An","instance","of",":","class",":","`","pykeen.stopper.EarlyStopper","`","with","settings","for","checking","if","training","should","stop","early",":","param","sub_batch_size",":","If","provided","split","each","batch","into","sub-batches","to","avoid","memory","issues","for","large","models","\/","small","GPUs",".",":","param","num_workers",":","The","number","of","child","CPU","workers","used","for","loading","data",".","If","None",",","data","are","loaded","in","the","main","process",".",":","param","clear_optimizer",":","Whether","to","delete","the","optimizer","instance","after","training","(","as","the","optimizer","might","have","additional","memory","consumption","due","to","e.g",".","moments","in","Adam",")",".",":","param","checkpoint_directory",":","An","optional","directory","to","store","the","checkpoint","files",".","If","None",",","a","subdirectory","named","``","checkpoints","``","in","the","directory","defined","by",":","data",":","`","pykeen.constants.PYKEEN_HOME","`","is","used",".","Unless","the","environment","variable","``","PYKEEN_HOME","``","is","overridden",",","this","will","be","``","~\/.pykeen\/checkpoints","``",".",":","param","checkpoint_name",":","The","filename","for","saving","checkpoints",".","If","the","given","filename","exists","already",",","that","file","will","be","loaded","and","used","to","continue","training",".",":","param","checkpoint_frequency",":","The","frequency","of","saving","checkpoints","in","minutes",".","Setting","it","to","0","will","save","a","checkpoint","after","every","epoch",".",":","param","checkpoint_on_failure",":","Whether","to","save","a","checkpoint","in","cases","of","a","RuntimeError","or","MemoryError",".","This","option","differs","from","ordinary","checkpoints",",","since","ordinary","checkpoints","are","only","saved","after","a","successful","epoch",".","When","saving","checkpoints","due","to","failure","of","the","training","loop","there","is","no","guarantee","that","all","random","states","can","be","recovered","correctly",",","which","might","cause","problems","with","regards","to","the","reproducibility","of","that","specific","training","loop",".","Therefore",",","these","checkpoints","are","saved","with","a","distinct","checkpoint","name",",","which","will","be","``","PyKEEN_just_saved_my_day_","{","datetime","}",".pt","``","in","the","given","checkpoint_root",".",":","param","drop_last",":","Whether","to","drop","the","last","batch","in","each","epoch","to","prevent","smaller","batches",".","Defaults","to","False",",","except","if","the","model","contains","batch","normalization","layers",".","Can","be","provided","explicitly","to","override",".",":","param","callbacks",":","An","optional",":","class",":","`","pykeen.training.TrainingCallback","`","or","collection","of","callback","instances","that","define","one","of","several","functionalities",".","Their","interface","was","inspired","by","Keras",".",":","param","callbacks_kwargs",":","additional","keyword-based","parameter","to","instantiate","the","training","callback",".",":","param","gradient_clipping_max_norm",":","The","maximum","gradient","norm","for","use","with","gradient","clipping",".","If","None",",","no","gradient","norm","clipping","is","used",".",":","param","gradient_clipping_norm_type",":","The","gradient","norm","type","to","use","for","maximum","gradient","norm",",","cf",".",":","func",":","`","torch.nn.utils.clip_grad_norm_","`",":","param","gradient_clipping_max_abs_value",":","The","maximum","absolute","value","in","gradients",",","cf",".",":","func",":","`","torch.nn.utils.clip_grad_value_","`",".","If","None",",","no","gradient","clipping","will","be","used",".",":","param","pin_memory",":","whether","to","use","memory","pinning","in","the","data","loader",",","cf",".","https",":","\/\/pytorch.org\/docs\/stable\/notes\/cuda.html","#","cuda-memory-pinning",":","return",":","The","losses","per","epoch","."]
code: def train(
      self,
      triples_factory: CoreTriplesFactory,
      num_epochs: int = 1,
      batch_size: Optional[int] = None,
      slice_size: Optional[int] = None,
      label_smoothing: float = 0.0,
      sampler: Optional[str] = None,
      continue_training: bool = False,
      only_size_probing: bool = False,
      use_tqdm: bool = True,
      use_tqdm_batch: bool = True,
      tqdm_kwargs: Optional[Mapping[str, Any]] = None,
      stopper: Optional[Stopper] = None,
      sub_batch_size: Optional[int] = None,
      num_workers: Optional[int] = None,
      clear_optimizer: bool = False,
      checkpoint_directory: Union[None, str, pathlib.Path] = None,
      checkpoint_name: Optional[str] = None,
      checkpoint_frequency: Optional[int] = None,
      checkpoint_on_failure: bool = False,
      drop_last: Optional[bool] = None,
      callbacks: TrainingCallbackHint = None,
      callbacks_kwargs: TrainingCallbackKwargsHint = None,
      gradient_clipping_max_norm: Optional[float] = None,
      gradient_clipping_norm_type: Optional[float] = None,
      gradient_clipping_max_abs_value: Optional[float] = None,
      pin_memory: bool = True,
  ) -> Optional[List[float]]:
      """Train the KGE model.
  
      .. note ::
  
          Gradient clipping is a technique to avoid the exploding gradient problem. Clip by norm and clip by value
          are two alternative implementations.
  
      :param triples_factory:
          The training triples.
      :param num_epochs:
          The number of epochs to train the model.
      :param batch_size:
          If set the batch size to use for mini-batch training. Otherwise find the largest possible batch_size
          automatically.
      :param slice_size: >0
          The divisor for the scoring function when using slicing. This is only possible for LCWA training loops in
          general and only for models that have the slicing capability implemented.
      :param label_smoothing: (0 <= label_smoothing < 1)
          If larger than zero, use label smoothing.
      :param sampler: (None or 'schlichtkrull')
          The type of sampler to use. At the moment sLCWA in R-GCN is the only user of schlichtkrull sampling.
      :param continue_training:
          If set to False, (re-)initialize the model's weights. Otherwise continue training.
      :param only_size_probing:
          The evaluation is only performed for two batches to test the memory footprint, especially on GPUs.
      :param use_tqdm: Should a progress bar be shown for epochs?
      :param use_tqdm_batch: Should a progress bar be shown for batching (inside the epoch progress bar)?
      :param tqdm_kwargs:
          Keyword arguments passed to :mod:`tqdm` managing the progress bar.
      :param stopper:
          An instance of :class:`pykeen.stopper.EarlyStopper` with settings for checking
          if training should stop early
      :param sub_batch_size:
          If provided split each batch into sub-batches to avoid memory issues for large models / small GPUs.
      :param num_workers:
          The number of child CPU workers used for loading data. If None, data are loaded in the main process.
      :param clear_optimizer:
          Whether to delete the optimizer instance after training (as the optimizer might have additional memory
          consumption due to e.g. moments in Adam).
      :param checkpoint_directory:
          An optional directory to store the checkpoint files. If None, a subdirectory named ``checkpoints`` in the
          directory defined by :data:`pykeen.constants.PYKEEN_HOME` is used. Unless the environment variable
          ``PYKEEN_HOME`` is overridden, this will be ``~/.pykeen/checkpoints``.
      :param checkpoint_name:
          The filename for saving checkpoints. If the given filename exists already, that file will be loaded and used
          to continue training.
      :param checkpoint_frequency:
          The frequency of saving checkpoints in minutes. Setting it to 0 will save a checkpoint after every epoch.
      :param checkpoint_on_failure:
          Whether to save a checkpoint in cases of a RuntimeError or MemoryError. This option differs from ordinary
          checkpoints, since ordinary checkpoints are only saved after a successful epoch. When saving checkpoints
          due to failure of the training loop there is no guarantee that all random states can be recovered correctly,
          which might cause problems with regards to the reproducibility of that specific training loop. Therefore,
          these checkpoints are saved with a distinct checkpoint name, which will be
          ``PyKEEN_just_saved_my_day_{datetime}.pt`` in the given checkpoint_root.
      :param drop_last:
          Whether to drop the last batch in each epoch to prevent smaller batches. Defaults to False, except if the
          model contains batch normalization layers. Can be provided explicitly to override.
      :param callbacks:
          An optional :class:`pykeen.training.TrainingCallback` or collection of callback instances that define
          one of several functionalities. Their interface was inspired by Keras.
      :param callbacks_kwargs:
          additional keyword-based parameter to instantiate the training callback.
      :param gradient_clipping_max_norm:
          The maximum gradient norm for use with gradient clipping. If None, no gradient norm clipping is used.
      :param gradient_clipping_norm_type:
          The gradient norm type to use for maximum gradient norm, cf. :func:`torch.nn.utils.clip_grad_norm_`
      :param gradient_clipping_max_abs_value:
          The maximum absolute value in gradients, cf. :func:`torch.nn.utils.clip_grad_value_`. If None, no
          gradient clipping will be used.
      :param pin_memory:
          whether to use memory pinning in the data loader, cf.
          https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning
  
      :return:
          The losses per epoch.
      """
      self._should_stop = False
  
      # In some cases, e.g. using Optuna for HPO, the cuda cache from a previous run is not cleared
      torch.cuda.empty_cache()
  
      # A checkpoint root is always created to ensure a fallback checkpoint can be saved
      if checkpoint_directory is None:
          checkpoint_directory = PYKEEN_CHECKPOINTS
      checkpoint_directory = pathlib.Path(checkpoint_directory)
      checkpoint_directory.mkdir(parents=True, exist_ok=True)
      logger.debug("using checkpoint_root at %s", checkpoint_directory)
  
      # If a checkpoint file is given, it must be loaded if it exists already
      save_checkpoints = False
      checkpoint_path = None
      best_epoch_model_file_path = None
      last_best_epoch = None
      if checkpoint_name:
          checkpoint_path = checkpoint_directory.joinpath(
              checkpoint_name
          )
          if checkpoint_path.is_file():
              (
                  best_epoch_model_file_path,
                  last_best_epoch,
              ) = self._load_state(
                  path=checkpoint_path,
                  triples_factory=triples_factory,
              )
              if stopper is not None:
                  stopper_dict = stopper.load_summary_dict_from_training_loop_checkpoint(
                      path=checkpoint_path
                  )
                  # If the stopper dict has any keys, those are written back to the stopper
                  if stopper_dict:
                      stopper._write_from_summary_dict(**stopper_dict)
                  else:
                      logger.warning(
                          "the training loop was configured with a stopper but no stopper configuration was "
                          "saved in the checkpoint",
                      )
              continue_training = True
          else:
              logger.info(
                  f"=> no checkpoint found at '{checkpoint_path}'. Creating a new file."
              )
          # The checkpoint frequency needs to be set to save checkpoints
          if checkpoint_frequency is None:
              checkpoint_frequency = 30
          save_checkpoints = True
      elif checkpoint_frequency is not None:
          logger.warning(
              "A checkpoint frequency was set, but no checkpoint file was given. No checkpoints will be created",
          )
  
      checkpoint_on_failure_file_path = None
      if checkpoint_on_failure:
          # In case a checkpoint frequency was set, we warn that no checkpoints will be saved
          date_string = datetime.now().strftime("%Y%m%d_%H_%M_%S")
          # If no checkpoints were requested, a fallback checkpoint is set in case the training loop crashes
          checkpoint_on_failure_file_path = (
              checkpoint_directory.joinpath(
                  PYKEEN_DEFAULT_CHECKPOINT.replace(
                      ".", f"_{date_string}."
                  ),
              )
          )
  
      # If the stopper loaded from the training loop checkpoint stopped the training, we return those results
      if getattr(stopper, "stopped", False):
          result: Optional[List[float]] = self.losses_per_epochs
      else:
          # send model to device before going into the internal training loop
          self.model = self.model.to(
              get_preferred_device(self.model, allow_ambiguity=True)
          )
  
          # the exit stack ensure that we clean up temporary files when an error occurs
          with ExitStack() as exit_stack:
              # When using early stopping models have to be saved separately at the best epoch, since the training
              # loop will due to the patience continue to train after the best epoch and thus alter the model
              if (
                  stopper is not None
                  and not only_size_probing
                  and last_best_epoch is None
                  and best_epoch_model_file_path is None
              ):
                  # note: NamedTemporaryFile does not seem to work
                  # Create a path
                  temporary_directory = exit_stack.enter_context(
                      TemporaryDirectory()
                  )
                  best_epoch_model_file_path = pathlib.Path(
                      temporary_directory
                  ).joinpath("best_model.pt")
  
              result = self._train(
                  num_epochs=num_epochs,
                  batch_size=batch_size,
                  slice_size=slice_size,
                  label_smoothing=label_smoothing,
                  sampler=sampler,
                  continue_training=continue_training,
                  only_size_probing=only_size_probing,
                  use_tqdm=use_tqdm,
                  use_tqdm_batch=use_tqdm_batch,
                  tqdm_kwargs=tqdm_kwargs,
                  stopper=stopper,
                  sub_batch_size=sub_batch_size,
                  num_workers=num_workers,
                  save_checkpoints=save_checkpoints,
                  checkpoint_path=checkpoint_path,
                  checkpoint_frequency=checkpoint_frequency,
                  checkpoint_on_failure_file_path=checkpoint_on_failure_file_path,
                  best_epoch_model_file_path=best_epoch_model_file_path,
                  last_best_epoch=last_best_epoch,
                  drop_last=drop_last,
                  callbacks=callbacks,
                  callbacks_kwargs=callbacks_kwargs,
                  gradient_clipping_max_norm=gradient_clipping_max_norm,
                  gradient_clipping_norm_type=gradient_clipping_norm_type,
                  gradient_clipping_max_abs_value=gradient_clipping_max_abs_value,
                  triples_factory=triples_factory,
                  pin_memory=pin_memory,
              )
  
      # Ensure the release of memory
      torch.cuda.empty_cache()
  
      # Clear optimizer
      if clear_optimizer:
          self.optimizer = None
          self.lr_scheduler = None
  
      return result
  
imports: import gc
  import logging
  import os
  import pathlib
  import pickle
  import random
  import time
  from abc import ABC, abstractmethod
  from contextlib import ExitStack
  from datetime import datetime
  from hashlib import md5
  from tempfile import NamedTemporaryFile, TemporaryDirectory
  from typing import IO, Any, ClassVar, Generic, List, Mapping, Optional, Tuple, TypeVar, Union
  import numpy
  import torch
  from class_resolver import HintOrType, OptionalKwargs
  from class_resolver.contrib.torch import lr_scheduler_resolver, optimizer_resolver
  from torch.optim.optimizer import Optimizer
  from torch.utils.data import DataLoader
  from tqdm.autonotebook import tqdm, trange
  from .callbacks import GradientAbsClippingTrainingCallback, GradientNormClippingTrainingCallback, LearningRateSchedulerTrainingCallback, MultiTrainingCallback, OptimizerTrainingCallback, StopperTrainingCallback, TrackerTrainingCallback, TrainingCallback, TrainingCallbackHint, TrainingCallbackKwargsHint
  from ..constants import PYKEEN_CHECKPOINTS, PYKEEN_DEFAULT_CHECKPOINT
  from ..lr_schedulers import LRScheduler
  from ..models import RGCN, Model
  from ..stoppers import Stopper
  from ..trackers import ResultTracker, tracker_resolver
  from ..triples import CoreTriplesFactory, TriplesFactory
  from ..typing import InductiveMode
  from ..utils import format_relative_comparison, get_batchnorm_modules, get_preferred_device, is_cuda_oom_error, is_cudnn_error, normalize_string
label_desc: Use image node_id 6 for calling the TrainingLoop obj's underlying member method code with example usage: obj.train(triples_factory, num_epochs, batch_size, slice_size, label_smoothing, sampler, continue_training, only_size_probing, use_tqdm, use_tqdm_batch, tqdm_kwargs, stopper, sub_batch_size, num_workers, clear_optimizer, checkpoint_directory, checkpoint_name, checkpoint_frequency, checkpoint_on_failure, drop_last, callbacks, callbacks_kwargs, gradient_clipping_max_norm, gradient_clipping_norm_type, gradient_clipping_max_abs_value, pin_memory) and returns: result