训练 Transformers HuggingFace 模型时,每隔 step/epoch 打印输入/输出/梯度/损失

Print input / output / grad / loss at every step/epoch when training Transformers HuggingFace model

我正在研究 HuggingFace 变形金刚并使用此处的玩具示例: https://huggingface.co/transformers/custom_datasets.html#fine-tuning-with-trainer

我真正需要的是:能够在每一步打印输入、输出、梯度和损失。 使用 Pytorch 训练循环很简单,但使用 HuggingFace Trainer 并不明显。 目前我有下一个想法:像这样创建一个 CustomCallback

class MyCallback(TrainerCallback):
    "A callback that prints a grad at every step"

    def on_step_begin(self, args, state, control, **kwargs):
        print("next step")
        print(kwargs['model'].classifier.out_proj.weight.grad.norm())

args = TrainingArguments(
    output_dir='test_dir',
    overwrite_output_dir=True,
    num_train_epochs=1,
    logging_steps=100,
    report_to="none",
    fp16=True,
    disable_tqdm=True,
)


trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    callbacks=[MyCallback],
)

trainer.train()

这样我就可以为任何模型层打印梯度和权重。 但是我还是想不通怎么打印input/output(比如我想在nan上查)和loss?

P.S。我还阅读了一些有关 forward_hook 的内容,但仍然找不到合适的代码示例。

虽然使用挂钩和自定义回调是解决问题的正确方法,但我找到了更好的解决方案 - 使用内置实用程序查找损失/权重/输入/输出中的 nan/Infhttps://huggingface.co/transformers/internal/trainer_utils.html#transformers.debug_utils.DebugUnderflowOverflow 因为 4.6.0 变形金刚有这样的选项。

您可以在 forward 函数中手动使用它,或者像这样为 TrainingArguments 使用附加选项:

args = TrainingArguments(
    output_dir='test_dir',
    overwrite_output_dir=True,
    num_train_epochs=1,
    logging_steps=100,
    report_to="none",
    fp16=True,
    disable_tqdm=True,
    debug="debug underflow_overflow"
)