2.4.2Trainer训练自定义的模型

comments: true
title: Trainer训练自定义模型​

前言

理论上来讲，基于torch或tensorflow实现的模型，并且正确重载对应方法（forward、from_pretrained等方法），那么就可以借助Trainer这个强大的工具避开繁琐的代码编写，提高工作效率。

下面将借用动手学深度学习内的线性回归的简洁实现来封装基于nn.Module实现的模型，以进一步使用Trainer。将重点讲解封装步骤，d2l部分不再展开。

代码

引入函数库

import torch
from torch import nn
from d2l import torch as d2l
from transformers import Trainer, TrainingArguments

true_w = torch.tensor([2, -3.4])
true_b = 4.2
​

在这里人为设定线性回归的真实权重为2和-3.4，偏置为4.2，接下来的目标是让神经网络无限逼近这个权重。

定义数据集

class CustDatasetForRegression(torch.utils.data.Dataset):
    def __init__(self, true_w, true_b, num_samples):
        self.true_w = true_w
        self.true_b = true_b
        self.num_samples = num_samples

        self.features, self.labels = d2l.synthetic_data(true_w, true_b, num_samples)

    def __getitem__(self, idx):
        item = {"inputs": self.features[idx], "labels": self.labels[idx]}
        return item

    def __len__(self):
        return len(self.features)
​

定义数据集时，确保重载后的__getitem__方法返回的字典的key对应于模型重载后的forward方法里面的形参名称。

data = CustDatasetForRegression(true_w, true_b, 1000)
​

定义模型

class CustomModelForRegression(nn.Module):
    def __init__(self):
        super(CustomModelForRegression, self).__init__()
        self.net = nn.Sequential(nn.Linear(2, 1))

    def forward(self, inputs, labels=None):
        logits = self.net(inputs)

        if labels is not None:
            loss_fn = nn.MSELoss()
            loss = loss_fn(logits, labels)
            return {"logits": logits, "loss": loss}
        else:
            return {"logits": logits}

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        model = cls(*model_args, **kwargs)
        state_dict = torch.load(f"{pretrained_model_name_or_path}/model.bin")
        model.load_state_dict(state_dict)
        return model

    def save_pretrained(self, save_directory):
        self.to(torch.device("cpu"))
        torch.save(self.state_dict(), f"{save_directory}/model.bin")

    def predict(self, inputs, device):
        device = device or self.device
        with torch.no_grad():
            inputs = inputs.to(device)
            out = self(inputs, None)
            return out["logits"].flatten()
​

自定义模型的步骤：

1. 重载前向传播逻辑。需要注意的是，应当有是否传入labels的判断，无论返回的是ModelOutput类，还是Python的字典，都应包含logits，当有标签传入时，还应有loss字段。
2. 重载加载预训练模型的方法。
3. 重载保存模型的方法。
4. 重载用于预测的方法。

创建模型

model = CustomModelForRegression()
​

CustomModelForRegression(
  (net): Sequential(
    (0): Linear(in_features=2, out_features=1, bias=True)
  )
)
​

创建训练器

optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
​

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=20,
    logging_strategy="epoch",
    per_device_train_batch_size=512,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=data,
    optimizers=(optimizer, None),
)

trainer.train()
​

保存与加载

model.save_pretrained("./model/")

model.from_pretrained("./model/")
​

推理

Parameter containing:
tensor([[ 1.9996, -3.4005]], requires_grad=True)
 Parameter containing:
tensor([4.2004], requires_grad=True)
​

model.predict(torch.tensor([[2.0, 3.0], [6.0, 7.0]]), torch.device("cpu"))
​

tensor([-2.0019, -7.6055])
​

可以看到基本和答案吻合了。

参考资料