👟 Trainer

一个基于PyTorch的通用模型训练器，具有简洁的代码基础和独特的见解。

安装

从GitHub安装：

git clone https://github.com/coqui-ai/Trainer
cd Trainer
make install

从PyPI安装：

pip install trainer

建议从GitHub安装，因为它更稳定。

实现模型

继承并重写TrainerModel()中的函数。

使用自动优化训练模型

参见MNIST示例。

使用高级优化训练模型

通过👟，您可以根据需要定义整个优化周期，就像下面的GAN示例一样。它为更高级的训练循环提供了更多的底层控制和灵活性。

您只需使用scaled_backward()函数来处理混合精度训练。

...

def optimize(self, batch, trainer):
    imgs, _ = batch

    # 采样噪声
    z = torch.randn(imgs.shape[0], 100)
    z = z.type_as(imgs)

    # 训练判别器
    imgs_gen = self.generator(z)
    logits = self.discriminator(imgs_gen.detach())
    fake = torch.zeros(imgs.size(0), 1)
    fake = fake.type_as(imgs)
    loss_fake = trainer.criterion(logits, fake)

    valid = torch.ones(imgs.size(0), 1)
    valid = valid.type_as(imgs)
    logits = self.discriminator(imgs)
    loss_real = trainer.criterion(logits, valid)
    loss_disc = (loss_real + loss_fake) / 2

    # 更新判别器
    _, _ = self.scaled_backward(loss_disc, None, trainer, trainer.optimizer[0])

    if trainer.total_steps_done % trainer.grad_accum_steps == 0:
        trainer.optimizer[0].step()
        trainer.optimizer[0].zero_grad()

    # 训练生成器
    imgs_gen = self.generator(z)

    valid = torch.ones(imgs.size(0), 1)
    valid = valid.type_as(imgs)

    logits = self.discriminator(imgs_gen)
    loss_gen = trainer.criterion(logits, valid)

    # 更新生成器
    _, _ = self.scaled_backward(loss_gen, None, trainer, trainer.optimizer[1])
    if trainer.total_steps_done % trainer.grad_accum_steps == 0:
        trainer.optimizer[1].step()
        trainer.optimizer[1].zero_grad()
    return {"model_outputs": logits}, {"loss_gen": loss_gen, "loss_disc": loss_disc}

...

参见GAN训练示例，其中使用了梯度累积。

使用批量大小查找器进行训练

有关使用批量大小查找器进行训练的测试脚本，请参见此处。

批量大小查找器从默认的BS（默认为2048，也可以由用户定义）开始，搜索可以适应您的硬件的最大批量大小。您应该预期它会运行多次训练，直到找到合适的批量大小。要使用它，您需要调用trainer.fit_with_largest_batch_size(starting_batch_size=2048)，而不是调用trainer.fit()，其中starting_batch_size是您想要开始搜索的批量大小。如果您想尽可能多地使用GPU内存，这非常有用。

使用DDP进行训练

$ python -m trainer.distribute --script path/to/your/train.py --gpus "0,1"

我们不使用.spawn()来启动多GPU训练，因为它会导致某些限制。

所有内容都必须是可序列化的。
.spawn()在子进程中训练模型，主进程中的模型不会更新。
当N很大时，具有N个进程的DataLoader会变得非常慢。

使用Accelerate进行训练

在TrainingArgs中将use_accelerate设置为True将启用使用Accelerate进行训练。

您也可以将其用于多GPU或分布式训练。

CUDA_VISIBLE_DEVICES="0,1,2" accelerate launch --multi_gpu --num_processes 3 train_recipe_autoregressive_prompt.py

参见Accelerate文档。

添加回调

👟支持回调以自定义您的运行。您可以在模型实现中设置回调，也可以显式地将它们提供给Trainer。

请查看trainer.utils.callbacks以查看可用的回调。

以下是如何为权重重新初始化向👟Trainer对象提供显式回调的示例。

def my_callback(trainer):
    print(" > 我的回调被调用了。")

trainer = Trainer(..., callbacks={"on_init_end": my_callback})
trainer.fit()

性能分析示例

根据需要创建torch性能分析器，并将其传递给trainer。

import torch
profiler = torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./profiler/"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
)
prof = trainer.profile_fit(profiler, epochs=1, small_run=64)
然后运行Tensorboard