2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license trainingarguments is the subset of the arguments we use in our example scripts which relate to the training loop itself using hfargumentparser we can turn this class into argparsehttps docs python org3libraryargparsemoduleargparse arguments that can be specified on the command line parameters outputdir str the output directory where the model predictions and checkpoints will be written overwriteoutputdir bool optional defaults to false if true overwrite the content of the output directory use this to continue training if outputdir points to a checkpoint directory dotrain bool optional defaults to false whether to run training or not this argument is not directly used by trainer it s intended to be used by your trainingevaluation scripts instead see the example scriptshttps github comhuggingfacetransformerstreemainexamples for more details doeval bool optional whether to run evaluation on the validation set or not will be set to true if evaluationstrategy is different from no this argument is not directly used by trainer it s intended to be used by your trainingevaluation scripts instead see the example scriptshttps github comhuggingfacetransformerstreemainexamples for more details dopredict bool optional defaults to false whether to run predictions on the test set or not this argument is not directly used by trainer it s intended to be used by your trainingevaluation scripts instead see the example scriptshttps github comhuggingfacetransformerstreemainexamples for more details evaluationstrategy str or trainerutils intervalstrategy optional defaults to no the evaluation strategy to adopt during training possible values are no no evaluation is done during training steps evaluation is done and logged every evalsteps epoch evaluation is done at the end of each epoch perdevicetrainbatchsize int optional defaults to 8 the batch size per gputpu corecpu for training perdeviceevalbatchsize int optional defaults to 8 the batch size per gputpu corecpu for evaluation gradientaccumulationsteps int optional defaults to 1 number of updates steps to accumulate the gradients for before performing a backwardupdate pass tip warningtrue when using gradient accumulation one step is counted as one step with backward pass therefore logging evaluation save will be conducted every gradientaccumulationsteps xxxstep training examples tip learningrate float optional defaults to 5e5 the initial learning rate for adam weightdecay float optional defaults to 0 the weight decay to apply if not zero adambeta1 float optional defaults to 0 9 the beta1 hyperparameter for the adam optimizer adambeta2 float optional defaults to 0 999 the beta2 hyperparameter for the adam optimizer adamepsilon float optional defaults to 1e8 the epsilon hyperparameter for the adam optimizer maxgradnorm float optional defaults to 1 0 maximum gradient norm for gradient clipping numtrainepochsfloat optional defaults to 3 0 total number of training epochs to perform maxsteps int optional defaults to 1 if set to a positive number the total number of training steps to perform overrides numtrainepochs for a finite dataset training is reiterated through the dataset if all data is exhausted until maxsteps is reached warmupratio float optional defaults to 0 0 ratio of total training steps used for a linear warmup from 0 to learningrate warmupsteps int optional defaults to 0 number of steps used for a linear warmup from 0 to learningrate overrides any effect of warmupratio loggingdir str optional tensorboardhttps www tensorflow orgtensorboard log directory will default to runscurrentdatetimehostname loggingstrategy str or trainerutils intervalstrategy optional defaults to steps the logging strategy to adopt during training possible values are no no logging is done during training epoch logging is done at the end of each epoch steps logging is done every loggingsteps loggingfirststep bool optional defaults to false whether to log and evaluate the first globalstep or not loggingsteps int optional defaults to 500 number of update steps between two logs if loggingstrategysteps savestrategy str or trainerutils intervalstrategy optional defaults to steps the checkpoint save strategy to adopt during training possible values are no no save is done during training epoch save is done at the end of each epoch steps save is done every savesteps savesteps int optional defaults to 500 number of updates steps before two checkpoint saves if savestrategysteps savetotallimit int optional if a value is passed will limit the total amount of checkpoints deletes the older checkpoints in outputdir nocuda bool optional defaults to false whether to not use cuda even when it is available or not seed int optional defaults to 42 random seed that will be set at the beginning of training fp16 bool optional defaults to false whether to use 16bit mixed precision training through nvidia apex instead of 32bit training fp16optlevel str optional defaults to o1 for fp16 training apex amp optimization level selected in o0 o1 o2 and o3 see details on the apex documentationhttps nvidia github ioapexamp localrank int optional defaults to 1 during distributed training the rank of the process tpunumcores int optional when training on tpu the number of tpu cores automatically passed by launcher script debug bool optional defaults to false whether to activate the trace to record computation graphs and profiling information or not dataloaderdroplast bool optional defaults to false whether to drop the last incomplete batch if the length of the dataset is not divisible by the batch size or not evalsteps int optional defaults to 1000 number of update steps before two evaluations pastindex int optional defaults to 1 some models like transformerxl modeldoctransformerxl or docxlnet modeldocxlnet can make use of the past hidden states for their predictions if this argument is set to a positive int the trainer will use the corresponding output usually index 2 as the past state and feed it to the model at the next training step under the keyword argument mems tpuname str optional the name of the tpu the process is running on tpuzone str optional the zone of the tpu the process is running on if not specified we will attempt to automatically detect from metadata gcpproject str optional google cloud project name for the cloud tpuenabled project if not specified we will attempt to automatically detect from metadata runname str optional a descriptor for the run notably used for wandb logging xla bool optional whether to activate the xla compilation or not set to float16 at first set to bfloat16 in case of tpu if you only want to use a specific subset of gpus use cudavisibledevices0 the strategy used for distributed training the number of replicas cpus gpus or tpu cores used in this training whether or not the current process should produce log the actual batch size for training may differ from pergputrainbatchsize in distributed training the actual batch size for evaluation may differ from pergpuevalbatchsize in distributed training the number of replicas cpus gpus or tpu cores used in this training 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license trainingarguments is the subset of the arguments we use in our example scripts which relate to the training loop itself using hfargumentparser we can turn this class into argparse https docs python org 3 library argparse module argparse arguments that can be specified on the command line parameters output_dir str the output directory where the model predictions and checkpoints will be written overwrite_output_dir bool optional defaults to false if true overwrite the content of the output directory use this to continue training if output_dir points to a checkpoint directory do_train bool optional defaults to false whether to run training or not this argument is not directly used by trainer it s intended to be used by your training evaluation scripts instead see the example scripts https github com huggingface transformers tree main examples for more details do_eval bool optional whether to run evaluation on the validation set or not will be set to true if evaluation_strategy is different from no this argument is not directly used by trainer it s intended to be used by your training evaluation scripts instead see the example scripts https github com huggingface transformers tree main examples for more details do_predict bool optional defaults to false whether to run predictions on the test set or not this argument is not directly used by trainer it s intended to be used by your training evaluation scripts instead see the example scripts https github com huggingface transformers tree main examples for more details evaluation_strategy str or trainer_utils intervalstrategy optional defaults to no the evaluation strategy to adopt during training possible values are no no evaluation is done during training steps evaluation is done and logged every eval_steps epoch evaluation is done at the end of each epoch per_device_train_batch_size int optional defaults to 8 the batch size per gpu tpu core cpu for training per_device_eval_batch_size int optional defaults to 8 the batch size per gpu tpu core cpu for evaluation gradient_accumulation_steps int optional defaults to 1 number of updates steps to accumulate the gradients for before performing a backward update pass tip warning true when using gradient accumulation one step is counted as one step with backward pass therefore logging evaluation save will be conducted every gradient_accumulation_steps xxx_step training examples tip learning_rate float optional defaults to 5e 5 the initial learning rate for adam weight_decay float optional defaults to 0 the weight decay to apply if not zero adam_beta1 float optional defaults to 0 9 the beta1 hyperparameter for the adam optimizer adam_beta2 float optional defaults to 0 999 the beta2 hyperparameter for the adam optimizer adam_epsilon float optional defaults to 1e 8 the epsilon hyperparameter for the adam optimizer max_grad_norm float optional defaults to 1 0 maximum gradient norm for gradient clipping num_train_epochs float optional defaults to 3 0 total number of training epochs to perform max_steps int optional defaults to 1 if set to a positive number the total number of training steps to perform overrides num_train_epochs for a finite dataset training is reiterated through the dataset if all data is exhausted until max_steps is reached warmup_ratio float optional defaults to 0 0 ratio of total training steps used for a linear warmup from 0 to learning_rate warmup_steps int optional defaults to 0 number of steps used for a linear warmup from 0 to learning_rate overrides any effect of warmup_ratio logging_dir str optional tensorboard https www tensorflow org tensorboard log directory will default to runs current_datetime_hostname logging_strategy str or trainer_utils intervalstrategy optional defaults to steps the logging strategy to adopt during training possible values are no no logging is done during training epoch logging is done at the end of each epoch steps logging is done every logging_steps logging_first_step bool optional defaults to false whether to log and evaluate the first global_step or not logging_steps int optional defaults to 500 number of update steps between two logs if logging_strategy steps save_strategy str or trainer_utils intervalstrategy optional defaults to steps the checkpoint save strategy to adopt during training possible values are no no save is done during training epoch save is done at the end of each epoch steps save is done every save_steps save_steps int optional defaults to 500 number of updates steps before two checkpoint saves if save_strategy steps save_total_limit int optional if a value is passed will limit the total amount of checkpoints deletes the older checkpoints in output_dir no_cuda bool optional defaults to false whether to not use cuda even when it is available or not seed int optional defaults to 42 random seed that will be set at the beginning of training fp16 bool optional defaults to false whether to use 16 bit mixed precision training through nvidia apex instead of 32 bit training fp16_opt_level str optional defaults to o1 for fp16 training apex amp optimization level selected in o0 o1 o2 and o3 see details on the apex documentation https nvidia github io apex amp local_rank int optional defaults to 1 during distributed training the rank of the process tpu_num_cores int optional when training on tpu the number of tpu cores automatically passed by launcher script debug bool optional defaults to false whether to activate the trace to record computation graphs and profiling information or not dataloader_drop_last bool optional defaults to false whether to drop the last incomplete batch if the length of the dataset is not divisible by the batch size or not eval_steps int optional defaults to 1000 number of update steps before two evaluations past_index int optional defaults to 1 some models like transformerxl model_doc transformerxl or doc xlnet model_doc xlnet can make use of the past hidden states for their predictions if this argument is set to a positive int the trainer will use the corresponding output usually index 2 as the past state and feed it to the model at the next training step under the keyword argument mems tpu_name str optional the name of the tpu the process is running on tpu_zone str optional the zone of the tpu the process is running on if not specified we will attempt to automatically detect from metadata gcp_project str optional google cloud project name for the cloud tpu enabled project if not specified we will attempt to automatically detect from metadata run_name str optional a descriptor for the run notably used for wandb logging xla bool optional whether to activate the xla compilation or not set to float16 at first set to bfloat16 in case of tpu if you only want to use a specific subset of gpus use cuda_visible_devices 0 the strategy used for distributed training the number of replicas cpus gpus or tpu cores used in this training whether or not the current process should produce log tf logging is handled by keras not the trainer the actual batch size for training may differ from per_gpu_train_batch_size in distributed training the actual batch size for evaluation may differ from per_gpu_eval_batch_size in distributed training the number of replicas cpus gpus or tpu cores used in this training
import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


logger = logging.get_logger(__name__)

if is_tf_available():
    import tensorflow as tf


@dataclass
class TFTrainingArguments(TrainingArguments):
    

    framework = "tf"
    tpu_name: Optional[str] = field(
        default=None,
        metadata={"help": "Name of TPU"},
    )

    tpu_zone: Optional[str] = field(
        default=None,
        metadata={"help": "Zone of TPU"},
    )

    gcp_project: Optional[str] = field(
        default=None,
        metadata={"help": "Name of Cloud TPU-enabled project"},
    )

    poly_power: float = field(
        default=1.0,
        metadata={"help": "Power for the Polynomial decay LR scheduler."},
    )

    xla: bool = field(default=False, metadata={"help": "Whether to activate the XLA compilation or not"})

    @cached_property
    def _setup_strategy(self) -> Tuple["tf.distribute.Strategy", int]:
        requires_backends(self, ["tf"])
        logger.info("Tensorflow: setting up strategy")

        gpus = tf.config.list_physical_devices("GPU")

        
        if self.fp16:
            tf.keras.mixed_precision.set_global_policy("mixed_float16")

        if self.no_cuda:
            strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
        else:
            try:
                if self.tpu_name:
                    tpu = tf.distribute.cluster_resolver.TPUClusterResolver(
                        self.tpu_name, zone=self.tpu_zone, project=self.gcp_project
                    )
                else:
                    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
            except ValueError:
                if self.tpu_name:
                    raise RuntimeError(f"Couldn't connect to TPU {self.tpu_name}!")
                else:
                    tpu = None

            if tpu:
                
                if self.fp16:
                    tf.keras.mixed_precision.set_global_policy("mixed_bfloat16")

                tf.config.experimental_connect_to_cluster(tpu)
                tf.tpu.experimental.initialize_tpu_system(tpu)

                strategy = tf.distribute.TPUStrategy(tpu)

            elif len(gpus) == 0:
                strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
            elif len(gpus) == 1:
                strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
            elif len(gpus) > 1:
                
                strategy = tf.distribute.MirroredStrategy()
            else:
                raise ValueError("Cannot find the proper strategy, please check your environment properties.")

        return strategy

    @property
    def strategy(self) -> "tf.distribute.Strategy":
        
        requires_backends(self, ["tf"])
        return self._setup_strategy

    @property
    def n_replicas(self) -> int:
        
        requires_backends(self, ["tf"])
        return self._setup_strategy.num_replicas_in_sync

    @property
    def should_log(self):
        
        return False  

    @property
    def train_batch_size(self) -> int:
        
        if self.per_gpu_train_batch_size:
            logger.warning(
                "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
                "version. Using `--per_device_train_batch_size` is preferred."
            )
        per_device_batch_size = self.per_gpu_train_batch_size or self.per_device_train_batch_size
        return per_device_batch_size * self.n_replicas

    @property
    def eval_batch_size(self) -> int:
        
        if self.per_gpu_eval_batch_size:
            logger.warning(
                "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
                "version. Using `--per_device_eval_batch_size` is preferred."
            )
        per_device_batch_size = self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
        return per_device_batch_size * self.n_replicas

    @property
    def n_gpu(self) -> int:
        
        requires_backends(self, ["tf"])
        warnings.warn(
            "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead.",
            FutureWarning,
        )
        return self._setup_strategy.num_replicas_in_sync