Optimum 文档

训练器

您正在查看 主分支 版本,需要从源代码安装。如果您希望使用常规的 pip 安装,请查看最新的稳定版本(v1.23.1)。
Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

训练器

ORTTrainer

optimum.onnxruntime.ORTTrainer

< >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module] = None args: ORTTrainingArguments = None data_collator: typing.Optional[DataCollator] = None train_dataset: typing.Optional[torch.utils.data.dataset.Dataset] = None eval_dataset: typing.Union[torch.utils.data.dataset.Dataset, typing.Dict[str, torch.utils.data.dataset.Dataset], NoneType] = None tokenizer: typing.Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase] = None model_init: typing.Union[typing.Callable[[], transformers.modeling_utils.PreTrainedModel], NoneType] = None compute_metrics: typing.Union[typing.Callable[[transformers.trainer_utils.EvalPrediction], typing.Dict], NoneType] = None callbacks: typing.Optional[typing.List[transformers.trainer_callback.TrainerCallback]] = None optimizers: typing.Tuple[torch.optim.optimizer.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None) preprocess_logits_for_metrics: typing.Union[typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor], NoneType] = None )

参数

  • model (PreTrainedModeltorch.nn.Module可选) — 要训练、评估或用于预测的模型。如果未提供,则必须传递 model_init

    ORTTrainer 经过优化,可与 transformers 库提供的 PreTrainedModel 配合使用。您仍然可以使用定义为 torch.nn.Module 的自定义模型进行 ONNX 运行时后端训练和 PyTorch 后端推理,只要它们的工作方式与 🤗 Transformers 模型相同。

  • args (ORTTrainingArguments可选) — 用于调整训练的参数。如果未提供,则默认为 ORTTrainingArguments 的基本实例,其中 output_dir 设置为当前目录中名为 tmp_trainer 的目录。
  • data_collator (DataCollator可选) — 用于从train_dataseteval_dataset的元素列表中形成批次的函数。如果未提供tokenizer,则默认为default_data_collator,否则为DataCollatorWithPadding的实例。
  • train_dataset (torch.utils.data.Datasettorch.utils.data.IterableDataset可选) — 用于训练的数据集。如果它是一个Dataset,则model.forward()方法不接受的列会自动删除。请注意,如果它是一个带有某些随机化的torch.utils.data.IterableDataset,并且您正在分布式环境中进行训练,则您的可迭代数据集应使用一个内部属性generator,该属性是一个torch.Generator,用于随机化,并且必须在所有进程中相同(并且ORTTrainer将在每个 epoch 手动设置此generator的种子),或者具有一个set_epoch()方法,该方法在内部设置所用 RNG 的种子。
  • eval_dataset (Union[torch.utils.data.Dataset, Dict[str, torch.utils.data.Dataset]],可选) — 用于评估的数据集。如果它是一个Dataset,则model.forward()方法不接受的列会自动删除。如果它是一个字典,它将在每个数据集上进行评估,并将字典键添加到指标名称之前。
  • tokenizer (PreTrainedTokenizerBase可选) — 用于预处理数据的标记器。如果提供,将用于在批处理输入时自动将输入填充到最大长度,并且它将与模型一起保存,以便更容易地重新运行中断的训练或重用微调的模型。
  • model_init (Callable[[], PreTrainedModel]可选) — 用于实例化要使用的模型的函数。如果提供,则每次调用ORTTrainer.train都将从该函数给出的模型的新实例开始。该函数可以没有参数,或者有一个包含 optuna/Ray Tune/SigOpt 试验对象的单个参数,以便能够根据超参数(例如层数、内部层的尺寸、dropout 概率等)选择不同的架构。
  • compute_metrics (Callable[[EvalPrediction], Dict]可选) — 用于计算评估指标的函数。必须接收一个EvalPrediction并返回一个字符串到指标值的字典。
  • callbacks (TrainerCallback列表,可选) — 用于自定义训练循环的回调列表。会将这些回调添加到此处详细介绍的默认回调列表中。如果您想删除其中一个使用的默认回调,请使用ORTTrainer.remove_callback方法。
  • optimizers (Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]可选) — 包含要使用的优化器和调度器的元组。将默认为模型上的AdamW实例和由get_linear_schedule_with_warmup控制的调度器(由args控制)。
  • preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]可选) — 在每个评估步骤缓存 logits 之前对其进行预处理的函数。必须接收两个张量,logits 和标签,并返回处理后的 logits。此函数所做的修改将反映在compute_metrics接收到的预测中。请注意,如果数据集没有标签,则标签(第二个参数)将为None

ORTTrainer 是一个简单但功能齐全的 ONNX Runtime 训练和评估循环,针对 🤗 Transformers 进行了优化。

重要属性

  • model — 始终指向核心模型。如果使用 transformers 模型,它将是 PreTrainedModel 的子类。
  • model_wrapped — 在一个或多个其他模块包装原始模型的情况下,始终指向最外部的模型。这是应该用于前向传递的模型。例如,在 DeepSpeed 下,内部模型首先被包装在 ORTModule 中,然后在 DeepSpeed 中,然后再次在 torch.nn.DistributedDataParallel 中。如果内部模型没有被包装,则 self.model_wrappedself.model 相同。
  • is_model_parallel — 模型是否已切换到模型并行模式(与数据并行不同,这意味着某些模型层分布在不同的 GPU 上)。
  • place_model_on_device — 是否自动将模型放置到设备上 - 如果使用模型并行或 deepspeed,或者如果默认的 ORTTrainingArguments.place_model_on_device 被覆盖为返回 False,则将其设置为 False
  • is_in_train — 模型当前是否正在运行 train(例如,当在 train 中调用 evaluate 时)。

create_optimizer

< >

( )

设置优化器。

我们提供了一个合理且效果良好的默认设置。如果您想使用其他优化器,可以在 ORTTrainer 的初始化过程中通过 optimizers 传递一个元组,或者在子类中重写此方法。

get_ort_optimizer_cls_and_kwargs

< >

( args: ORTTrainingArguments )

参数

  • args (ORTTrainingArguments) — 训练会话的训练参数。

根据 ORTTrainingArguments 返回在 ONNX Runtime 中实现的优化器类和优化器参数。

train

< >

( resume_from_checkpoint: typing.Union[bool, str, NoneType] = None trial: typing.Union[ForwardRef('optuna.Trial'), typing.Dict[str, typing.Any]] = None ignore_keys_for_eval: typing.Optional[typing.List[str]] = None **kwargs )

参数

  • resume_from_checkpoint (strbool可选) — 如果是 str,则为先前 ORTTrainer 实例保存的已保存检查点的本地路径。如果是 bool 且等于 True,则加载先前 ORTTrainer 实例保存在 args.output_dir 中的最后一个检查点。如果存在,训练将从此处加载的模型/优化器/调度程序状态恢复。
  • trial (optuna.TrialDict[str, Any]可选) — 超参数搜索的试验运行或超参数字典。
  • ignore_keys_for_eval (List[str]可选) — 模型输出中(如果它是字典)应在训练期间收集评估预测时忽略的一系列键。
  • kwargs (Dict[str, Any]可选) — 用于隐藏已弃用参数的其他关键字参数

使用 ONNX Runtime 加速器进行训练的主要入口点。

ORTSeq2SeqTrainer

optimum.onnxruntime.ORTSeq2SeqTrainer

< >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module] = None args: ORTTrainingArguments = None data_collator: typing.Optional[DataCollator] = None train_dataset: typing.Optional[torch.utils.data.dataset.Dataset] = None eval_dataset: typing.Union[torch.utils.data.dataset.Dataset, typing.Dict[str, torch.utils.data.dataset.Dataset], NoneType] = None tokenizer: typing.Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase] = None model_init: typing.Union[typing.Callable[[], transformers.modeling_utils.PreTrainedModel], NoneType] = None compute_metrics: typing.Union[typing.Callable[[transformers.trainer_utils.EvalPrediction], typing.Dict], NoneType] = None callbacks: typing.Optional[typing.List[transformers.trainer_callback.TrainerCallback]] = None optimizers: typing.Tuple[torch.optim.optimizer.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None) preprocess_logits_for_metrics: typing.Union[typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor], NoneType] = None )

评估

< >

( eval_dataset: typing.Optional[torch.utils.data.dataset.Dataset] = None ignore_keys: typing.Optional[typing.List[str]] = None metric_key_prefix: str = 'eval' **gen_kwargs )

参数

  • eval_dataset (Dataset, 可选) — 如果您希望覆盖 self.eval_dataset,请传递一个数据集。如果它是一个 Dataset,则模型的 model.forward() 方法无法接受的列会自动删除。它必须实现 __len__ 方法。
  • ignore_keys (List[str], 可选) — 模型输出(如果它是字典)中在收集预测时应忽略的一系列键。
  • metric_key_prefix (str, 可选,默认为 "eval") — 用作指标键前缀的可选前缀。例如,如果前缀为 "eval"(默认值),则指标“bleu”将命名为“eval_bleu”
  • max_length (int, 可选) — 使用生成方法预测时使用的最大目标长度。
  • num_beams (int, 可选) — 用于生成方法预测时的波束搜索的波束数。1 表示不进行波束搜索。gen_kwargs — 其他 generate 特定的关键字参数。

运行评估并返回指标。

调用脚本将负责提供计算指标的方法,因为它们是依赖于任务的(将其传递给初始化的 compute_metrics 参数)。

您还可以子类化并覆盖此方法以注入自定义行为。

预测

< >

( test_dataset: 数据集 ignore_keys: typing.Optional[typing.List[str]] = None metric_key_prefix: str = 'test' **gen_kwargs )

参数

  • test_dataset (Dataset) — 用于运行预测的数据集。如果它是一个 Dataset,则model.forward()方法不接受的列会自动删除。必须实现__len__方法
  • ignore_keys (List[str], 可选) — 模型输出(如果是字典)中应在收集预测时忽略的键列表。
  • metric_key_prefix (str, 可选,默认为 "eval") — 用于作为指标键前缀的可选前缀。例如,如果前缀为"eval"(默认值),则指标“bleu”将命名为“eval_bleu”
  • max_length (int, 可选) — 使用生成方法预测时使用的最大目标长度。
  • num_beams (int, 可选) — 使用生成方法预测时将使用的波束搜索的波束数。1 表示不进行波束搜索。gen_kwargs — 其他 generate 特定关键字参数。

运行预测并返回预测和潜在的指标。

根据数据集和用例,测试数据集可能包含标签。在这种情况下,此方法还会返回指标,就像在evaluate()中一样。

如果预测或标签具有不同的序列长度(例如,因为您在令牌分类任务中执行动态填充),则预测将被填充(在右侧)以允许连接到一个数组中。填充索引为 -100。

返回:NamedTuple 具有以下键的命名元组

  • predictions (np.ndarray): test_dataset上的预测。
  • label_ids (np.ndarray, 可选): 标签(如果数据集包含一些)。
  • metrics (Dict[str, float], 可选): 潜在的指标字典(如果数据集包含标签)。

ORTTrainingArguments

optimum.onnxruntime.ORTTrainingArguments

< >

( output_dir: str overwrite_output_dir: bool = False do_train: bool = False do_eval: bool = False do_predict: bool = False evaluation_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'no' prediction_loss_only: bool = False per_device_train_batch_size: int = 8 per_device_eval_batch_size: int = 8 per_gpu_train_batch_size: typing.Optional[int] = None per_gpu_eval_batch_size: typing.Optional[int] = None gradient_accumulation_steps: int = 1 eval_accumulation_steps: typing.Optional[int] = None eval_delay: typing.Optional[float] = 0 learning_rate: float = 5e-05 weight_decay: float = 0.0 adam_beta1: float = 0.9 adam_beta2: float = 0.999 adam_epsilon: float = 1e-08 max_grad_norm: float = 1.0 num_train_epochs: float = 3.0 max_steps: int = -1 lr_scheduler_type: typing.Union[transformers.trainer_utils.SchedulerType, str] = 'linear' lr_scheduler_kwargs: typing.Optional[typing.Dict] = <factory> warmup_ratio: float = 0.0 warmup_steps: int = 0 log_level: typing.Optional[str] = 'passive' log_level_replica: typing.Optional[str] = 'warning' log_on_each_node: bool = True logging_dir: typing.Optional[str] = None logging_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'steps' logging_first_step: bool = False logging_steps: float = 500 logging_nan_inf_filter: bool = True save_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'steps' save_steps: float = 500 save_total_limit: typing.Optional[int] = None save_safetensors: typing.Optional[bool] = True save_on_each_node: bool = False save_only_model: bool = False no_cuda: bool = False use_cpu: bool = False use_mps_device: bool = False seed: int = 42 data_seed: typing.Optional[int] = None jit_mode_eval: bool = False use_ipex: bool = False bf16: bool = False fp16: bool = False fp16_opt_level: str = 'O1' half_precision_backend: str = 'auto' bf16_full_eval: bool = False fp16_full_eval: bool = False tf32: typing.Optional[bool] = None local_rank: int = -1 ddp_backend: typing.Optional[str] = None tpu_num_cores: typing.Optional[int] = None tpu_metrics_debug: bool = False debug: typing.Union[str, typing.List[transformers.debug_utils.DebugOption]] = '' dataloader_drop_last: bool = False eval_steps: typing.Optional[float] = None dataloader_num_workers: int = 0 past_index: int = -1 run_name: typing.Optional[str] = None disable_tqdm: typing.Optional[bool] = None remove_unused_columns: typing.Optional[bool] = True label_names: typing.Optional[typing.List[str]] = None load_best_model_at_end: typing.Optional[bool] = False metric_for_best_model: typing.Optional[str] = None greater_is_better: typing.Optional[bool] = None ignore_data_skip: bool = False fsdp: typing.Union[typing.List[transformers.trainer_utils.FSDPOption], str, NoneType] = '' fsdp_min_num_params: int = 0 fsdp_config: typing.Optional[str] = None fsdp_transformer_layer_cls_to_wrap: typing.Optional[str] = None deepspeed: typing.Optional[str] = None label_smoothing_factor: float = 0.0 optim: typing.Optional[str] = 'adamw_hf' optim_args: typing.Optional[str] = None adafactor: bool = False group_by_length: bool = False length_column_name: typing.Optional[str] = 'length' report_to: typing.Optional[typing.List[str]] = None ddp_find_unused_parameters: typing.Optional[bool] = None ddp_bucket_cap_mb: typing.Optional[int] = None ddp_broadcast_buffers: typing.Optional[bool] = None dataloader_pin_memory: bool = True dataloader_persistent_workers: bool = False skip_memory_metrics: bool = True use_legacy_prediction_loop: bool = False push_to_hub: bool = False resume_from_checkpoint: typing.Optional[str] = None hub_model_id: typing.Optional[str] = None hub_strategy: typing.Union[transformers.trainer_utils.HubStrategy, str] = 'every_save' hub_token: typing.Optional[str] = None hub_private_repo: bool = False hub_always_push: bool = False gradient_checkpointing: bool = False gradient_checkpointing_kwargs: typing.Optional[dict] = None include_inputs_for_metrics: bool = False fp16_backend: str = 'auto' push_to_hub_model_id: typing.Optional[str] = None push_to_hub_organization: typing.Optional[str] = None push_to_hub_token: typing.Optional[str] = None mp_parameters: str = '' auto_find_batch_size: bool = False full_determinism: bool = False torchdynamo: typing.Optional[str] = None ray_scope: typing.Optional[str] = 'last' ddp_timeout: typing.Optional[int] = 1800 torch_compile: bool = False torch_compile_backend: typing.Optional[str] = None torch_compile_mode: typing.Optional[str] = None dispatch_batches: typing.Optional[bool] = None split_batches: typing.Optional[bool] = False include_tokens_per_second: typing.Optional[bool] = False include_num_input_tokens_seen: typing.Optional[bool] = False neftune_noise_alpha: float = None use_module_with_loss: typing.Optional[bool] = False save_onnx: typing.Optional[bool] = False onnx_prefix: typing.Optional[str] = None onnx_log_level: typing.Optional[str] = 'WARNING' )

参数

  • optim (strtraining_args.ORTOptimizerNamestransformers.training_args.OptimizerNames, 可选,默认为 "adamw_hf") — 要使用的优化器,包括 Transformers 中的优化器:adamw_hf、adamw_torch、adamw_apex_fused 或 adafactor。以及 ONNX Runtime 实现的优化器:adamw_ort_fused。

ORTSeq2SeqTrainingArguments

optimum.onnxruntime.ORTSeq2SeqTrainingArguments

  • optim (strtraining_args.ORTOptimizerNamestransformers.training_args.OptimizerNames可选,默认为 "adamw_hf") — 使用的优化器,包括 Transformers 中的优化器:adamw_hf、adamw_torch、adamw_apex_fused 或 adafactor。以及 ONNX Runtime 实现的优化器:adamw_ort_fused。
< > 在 GitHub 上更新