使用 Trainer API 训练自定义 AI 模型并将其适配到 🤗

Hafedh Hichri

不是-lain

关注

介绍

训练 LLM 需要大量的计算能力和多个 GPU，而且共享/版本控制这些 AI 模型往往很棘手。

但是，现在可以比以往任何时候都更容易地直接从免费的开源可靠的机器学习平台推送到和加载您的权重，该平台利用了 ModelHubMixin 类。

至于为什么您需要使用 Trainer API 来适配您的模型，那是因为它允许您的模型几乎不需要任何训练脚本即可进行训练。它甚至默认提供了 分布式训练 兼容性。

设置和数据处理

首先，让我们下载依赖项

pip install -q datasets evaluate accelerate "huggingface_hub>=0.22"

现在，让我们用 HF 写入令牌登录

from huggingface_hub import notebook_login
notebook_login()

现在开始编码

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from datasets import load_dataset, Image
from transformers import DefaultDataCollator, TrainingArguments, Trainer
from huggingface_hub import PyTorchModelHubMixin
import evaluate
import numpy as np

原始数据集可在 这里 找到

dataset = load_dataset("mnist")

# convert the "image" column to a pillow image
dataset = dataset.cast_column("image", Image())

重要

您的目标列需要命名为 labels

否则，您需要在训练参数中传递您的 label_names

dataset = dataset.rename_column("label","labels")

# convert to Pytorch tensors
transform = transforms.Compose([transforms.ToTensor()])
def to_pt(batch):
    batch["image"] = [transform(image.convert("RGB")) for image in batch["image"]]
    return batch
train = dataset["train"].with_transform(to_pt)
test = dataset["test"].with_transform(to_pt)

与 dataset.map 不同，dataset.with_transform 仅在访问数据时应用转换，您可以在文档中了解更多信息。

不幸的是，它与 streaming 数据集不兼容，如果您对此感兴趣，可以在其 github 仓库中贡献。

模型

将您的模型适配到 huggingface

要使您的模型与 huggingface 兼容，您只需继承相应的类，就是这样。如果您想用更多可选参数来修饰您的 README，请在此处传递它们

您只需继承其中一个 Mixin 类，其他一切都是可选的

class BasicNet(nn.Module,PyTorchModelHubMixin,tags=["image-classification"]):
  (...)

mixin 类的工作方式是，它们获取 __init__ 方法中使用的参数并将其存储在 config.json 文件中。调用 push_to_hub 方法时，我们将模型权重和 config.json 推送到 hub。使用 from_pretrained 时，我们下载模型权重和 config 文件，使用存储在 config 文件中的参数实例化模型，并注入我们的权重。

总结：

• 它们为您的模型添加了 3 个方法，类似于 transformers 库的工作方式，即
  • save_pretrained（在本地保存权重）
  • from_pretrained（从 hub 或本地加载并初始化模型）
  • push_to_hub（将权重和配置推送到 hub）
• 您可以传入其他可选的元数据，以使您的 README 文件脱颖而出

Trainer API 兼容性

既然我们已经强调了这一点，让我们进入使我们的模型与 Trainer API 兼容的部分。

这严重依赖于 forward 方法

forward 方法需要一个名为 labels 的参数，否则我们需要在 TrainingArguments 中使用 label_names

class BasicNet(nn.Module,PyTorchModelHubMixin,tags=["image-classification"]):
    def __init__(self,channels):
        super().__init__()
        self.criterion = nn.CrossEntropyLoss()
        self.conv1 = nn.Conv2d(channels, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
        self.softmax = nn.Softmax(dim=-1)
    def forward(self, image,labels=None):
        # the labels parameter allows us to finetune our model
        # with the Trainer API easily
        x = F.relu(F.max_pool2d(self.conv1(image), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        logits = F.log_softmax(x,dim=-1)
        if labels != None :
          # this will make your AI compatible with the trainer API
          loss = self.criterion(logits, labels)
          return {"loss": loss, "logits": logits}
        return logits

forward 方法需要至少包含两个内容：

• image: 或 X，表示将通过模型的数据
• labels: 或 Y，如果存在，我们将计算并以字典类型或 ModelOutput 类型返回损失

根据我们的数据集列，我们可以调整 forward 方法以匹配它，请注意数据是以 kwargs 的形式传递的。

训练模型

首先，让我们初始化模型

# The number 3 will be stored in the config.json file
model = BasicNet(channels=3) #RGB

让我们在这里添加一个可选的计算指标部分

accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

collate 函数用于以适当的方式对批次中的数据进行整理

在将 collate 函数传递给 Trainer API 之前，您可以使用以下代码对其进行测试：collate_fn(train.select(range(8)))，这将选择数据集的前 8 个样本并将其传递给 collate_fn，然后您可以验证所有数据形状是否符合您的需求。

def collate_fn(examples):
    images = []
    labels = []
    for example in examples:
        images.append((example["image"]))
        labels.append(example["labels"])

    images = torch.stack(images)
    labels = torch.tensor(labels)
    return {"image": images, "labels": labels}

现在我们定义训练我们AI模型所需的参数，建议您阅读有关训练参数（TrainingArguments）的文档，以找出符合您要求的参数。

training_args = TrainingArguments(
    output_dir="my_mnist_model",
    # remove_unused_columns=False,
    evaluation_strategy="steps",
    save_strategy="epoch",
    learning_rate=5e-4,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    num_train_epochs=2,
    logging_steps=100,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train,
    eval_dataset=test,
    compute_metrics=compute_metrics,
)

trainer.train()

Trainer API 只推送权重而不包含配置（初始化参数）

要解决此问题，您应该在训练完成后手动推送您的模型

model.push_to_hub("my_mnist_model")

您就大功告成了 🥳🥳🥳

从中心加载模型

要从 hub 加载您的模型，无需再次重新初始化它，也无需手动下载权重，您可以直接使用该类为您完成此操作。

new_model = BasicNet.from_pretrained("not-lain/my_mnist_model")

from safetensors.torch import load_file
from huggingface_hub import snapshot_download
snapshot_download(repo_id="not-lain/my_mnist_model",local_dir="temp_folder")

# or use the parameters in the config.json
new_model = BasicNet(3)
weights = load_file("temp_folder/model.safetensors")
new_model.load_state_dict(weights)

结尾

希望本文能为您提供有用的信息，并让您学到很多。如果您喜欢这篇博文，请考虑点赞 🤗。如果您有任何问题或疑问，请随时发起讨论并标记我，或直接私信我。

有趣的是，我对本文中使用的 ModelHubMixin 类做了一点贡献，所以希望您会喜欢 🤗

感谢阅读这篇博文，爱你们所有人 ❤️