HuggingFace (transformers) 自定义图像数据集、使用 DeiT 模型、Trainer 进行训练回归任务

资料

Hugging Face 官方文档https://huggingface.co/
Hugging Face 代码链接https://github.com/huggingface/transformers

1. 环境准备

1. 创建 conda 环境
2. 激活 conda 环境
3. 下载 transformers 依赖
4. 下载 transformers 中需要处理数据集的依赖
5. 下载 pytorch 依赖因为这里使用的 transformers 是基于 PyTorch 实现的所以需要导入 pytorch 依赖
6. 下载 tensorboard 依赖。训练过程中使用 TensorBoard 可视化

conda create -n hugging python=3.7 
conda activate hugging
conda install -c huggingface transformers
conda install datasets
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
conda install tensorboard

7. 打开 PyCharm配置 Interpreter
   依次点击File -> Settings
   
   然后选择刚才创建的 conda 环境
   

2 任务及数据集描述

需求说明有一个视线估计任务输入为人脸图像输出为该人脸图像在手机屏幕上的注视点坐标 (x, y)。

数据集的目录结构如下

\GazeCapture_new
	-- Image
		-- 00002
			-- face
				-- 00000.jpg
				-- 00001.jpg
				-- .....
			-- grid
				-- .....
			-- left
				-- ....
			-- right
				-- .....
		-- 00003
			-- face
				-- .....
			-- grid
				-- .....
			-- left
				-- ....
			-- right
				-- .....
		-- ......
	-- Label
		-- train
			-- 00002.label
			-- .....
		-- test
			-- 03024.label
			-- .....
		-- val
			-- ......

每一个标签文件中的内容如 00002.label 存储的内容

Face Left Right Grid Xcam, Ycam Xdot, Ydot Device
00002\face\00000.jpg 00002\left\00000.jpg 00002\right\00000.jpg 00002\grid\00000.jpg 1.064,-6.0055 160,284 iPhone6
00002\face\00001.jpg 00002\left\00001.jpg 00002\right\00001.jpg 00002\grid\00001.jpg 1.064,-6.0055 160,284 iPhone6
00002\face\00002.jpg 00002\left\00002.jpg 00002\right\00002.jpg 00002\grid\00002.jpg 1.064,-6.0055 160,284 iPhone6
00002\face\00003.jpg 00002\left\00003.jpg 00002\right\00003.jpg 00002\grid\00003.jpg 1.064,-6.0055 160,284 iPhone6
.......

• Face 表示脸部图片的存储路径。
• Left 表示左眼图片的存储路径。
• Right 表示右眼图片的存储路径。
• Grid 表示网格图片的存储路径。
• Xcam, Ycam 是标签表示人脸图片对应的视线位置的 (x, y) 坐标单位为厘米。 后续的训练过程使用这两个值作为标签。
• Xdot, Ydot 表示人脸图片对应的视线位置的 (x, y) 坐标单位为像素。
• Device 表示采集设备型号。

如果想要使用我的数据集先把代码跑通这里提供我使用的部分数据集作为参考但由于不是完整的数据集所以训练效果不是很好仅供跑通代码作为参考。
https://drive.google.com/file/d/1gM-wzkaEcnw0GEKQ2eedpYlvjuqhp3gA/view?usp=sharing

3. DataSet

注意Dataset 一定不要完全粘贴我的代码一定要按照自己的数据集编写对应代码。只有以下几点需要和我一模一样

1. 自定义类继承 Dataset自定义的类名可以自行命名。
2. 重写 __init__、__len__、__getitem__这三个方法方法内的具体逻辑根据自己的数据集修改。
3. __getitem__ 方法的返回值形式一定要是 {"labels": xxx, "pixel_values": xxx}。

import os.path

from torch.utils.data import Dataset
from transform import transform
import numpy as np

# 读取数据如果是训练数据随即打乱数据顺序
def get_label_list(label_path):
    # 存储所有标签文件中的所有内容
    full_lines = []
    # 获取所有标签文件的名称如 00002.label, 00003.label, ......
    label_names = os.listdir(label_path)
    # 遍历每一个标签文件并读取其中内容
    for label_name in label_names:
        # 标签文件全路径如 D:\datasets\GazeCapture_new\Label\train\00002.label
        label_abs_path = os.path.join(label_path, label_name)
        # 读取每一个标签文件中的内容
        with open(label_abs_path) as flist:
            # 存储该标签文件中的所有内容
            full_line = []
            for line in flist:
                full_line.append(line.strip())
            # 移除首行表头 'Face Left Right Grid Xcam, Ycam Xdot, Ydot Device'
            full_line.pop(0)
            full_lines.extend(full_line)
    return full_lines


class GazeCaptureDataset(Dataset):
    def __init__(self, root_path, data_type):
        self.data_dir = root_path
        # 标签文件的根路径如 D:\datasets\GazeCapture_new\Label\train
        label_root_path = os.path.join(root_path + '/Label', data_type)
        # 获取所有标签文件中的所有内容
        self.full_lines = get_label_list(label_root_path)
        # 每一行内容的分隔符
        self.delimiter = ' '
        # 数据集长度也就是一共有多少个图片
        self.num_samples = len(self.full_lines)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # 标签文件的一行对应一个训练实例
        line = self.full_lines[idx]
        # 将标签文件中的一行内容按照分隔符进行分割
        Face, Left, Right, Grid, XYcam, XYdot, Device = line.split(self.delimiter)
        # 获取网络的输入人脸图片
        face_path = os.path.join(self.data_dir + '/Image/', Face)
        # 读取人脸图像
        with open(face_path, 'rb') as f:
            img = f.read()
        # 将人脸图像进行格式转化缩放、裁剪、标准化
        pixel_values = transform(img)
        # 获取标签值
        labels = np.array(XYcam.split(","), np.float32)
        # 注意返回值的形式一定要是 {"labels": xxx, "pixel_values": xxx}
        result = {"labels": labels}
        result["pixel_values"] = pixel_values
        return result

transform.py 工具类的代码如下

# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import cv2
from PIL import Image


# 定义decode_image函数将图片转为Numpy格式r
def decode_image(img, to_rgb=True):
    data = np.frombuffer(img, dtype='uint8')
    img = cv2.imdecode(data, 1)
    if to_rgb:
        assert img.shape[2] == 3, 'invalid shape of image[%s]' % (
            img.shape)
        img = img[:, :, ::-1]

    return img


# 定义resize_image函数对图片大小进行调整
def resize_image(img, size=None, resize_short=None, interpolation=-1):
    interpolation = interpolation if interpolation >= 0 else None
    if resize_short is not None and resize_short > 0:
        resize_short = resize_short
        w = None
        h = None
    elif size is not None:
        resize_short = None
        w = size if type(size) is int else size[0]
        h = size if type(size) is int else size[1]
    else:
        raise ValueError("invalid params for ReisizeImage for '\
            'both 'size' and 'resize_short' are None")

    img_h, img_w = img.shape[:2]
    if resize_short is not None:
        percent = float(resize_short) / min(img_w, img_h)
        w = int(round(img_w * percent))
        h = int(round(img_h * percent))
    else:
        w = w
        h = h
    if interpolation is None:
        return cv2.resize(img, (w, h))
    else:
        return cv2.resize(img, (w, h), interpolation=interpolation)


# 定义crop_image函数对图片进行裁剪
def crop_image(img, size):
    if type(size) is int:
        size = (size, size)
    else:
        size = size  # (h, w)

    w, h = size
    img_h, img_w = img.shape[:2]
    w_start = (img_w - w) // 2
    h_start = (img_h - h) // 2

    w_end = w_start + w
    h_end = h_start + h
    return img[h_start:h_end, w_start:w_end, :]


# 定义normalize_image函数对图片进行归一化
def normalize_image(img, scale=None, mean=None, std=None, order= ''):
    if isinstance(scale, str):
        scale = eval(scale)
    scale = np.float32(scale if scale is not None else 1.0 / 255.0)
    mean = mean if mean is not None else [0.485, 0.456, 0.406]
    std = std if std is not None else [0.229, 0.224, 0.225]

    shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
    mean = np.array(mean).reshape(shape).astype('float32')
    std = np.array(std).reshape(shape).astype('float32')

    if isinstance(img, Image.Image):
        img = np.array(img)
    assert isinstance(img, np.ndarray), "invalid input 'img' in NormalizeImage"
    # 对图片进行归一化
    return (img.astype('float32') * scale - mean) / std


# 定义to_CHW_image函数对图片进行通道变换将原通道为‘hwc’的图像转为‘chw‘
def to_CHW_image(img):
    if isinstance(img, Image.Image):
        img = np.array(img)
    # 对图片进行通道变换
    return img.transpose((2, 0, 1))


# 图像预处理方法汇总
def transform(data, mode='train'):

    # 图像解码
    data = decode_image(data)
    # 图像缩放
    data = resize_image(data, resize_short=224)
    # 图像裁剪
    data = crop_image(data, size=224)
    # 标准化
    data = normalize_image(data, scale=1./255., mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    # 通道变换
    data = to_CHW_image(data)
    return data

4. 训练

from transformers import TrainingArguments
from transformers import DeiTForImageClassification
from torch import nn
from transformers import Trainer
from transformers import DeiTConfig
from dataset import GazeCaptureDataset

# 数据集根路径
root_path = r"D:\datasets\GazeCapture_new"
# 1.定义 Dataset
train_dataset = GazeCaptureDataset(root_path, data_type='train')
val_dataset = GazeCaptureDataset(root_path, data_type='val')

# 2.定义 DeiT 图像模型
'''
num_labels 表示图像的输出值为 2即 (x, y) 两个坐标值
problem_type="regression" 表示任务是回归任务
'''
configuration = DeiTConfig(num_labels=2, problem_type="regression")
model = DeiTForImageClassification(configuration)


# 3.训练
## 3.1 训练参数
'''
output_dir模型预测和 checkpoint 的输出目录。
evaluation_strategy 训练过程中采用的验证策略。可能的取值有
    "no": 训练过程中不验证
    "steps": 在每个 eval_steps 中执行(并记录)验证。
    "epoch": 在每个 epoch 结束时进行验证。
eval_steps=100每 100 次训练执行一次验证。
per_device_train_batch_size/per_device_eval_batch_size用于训练/验证的 batch size。
logging_dirTensorBoard 日志目录。默认为 *output_dir/runs/CURRENT_DATETIME_HOSTNAME*。
logging_steps=50每隔 50 步写入 TensorBoard
save_strategy 训练期间采用的 checkpoint 保存策略。可能取值为:
    "no": 训练期间不保存 checkpoint
    "epoch": 每个 epoch 结束后保存 checkpoint
    "steps": 每个 save_steps 结束后保存 checkpoint
save_steps=100每 100 次训练保存一次 checkpoint
'''
training_args = TrainingArguments(output_dir="gaze_trainer",
                                  evaluation_strategy="steps",
                                  eval_steps=100,
                                  per_device_train_batch_size=2,
                                  per_device_eval_batch_size=2,
                                  logging_dir='./logs',
                                  logging_steps=50,
                                  save_strategy="steps",
                                  save_steps=100)
## 3.2 自定义 Trainer
class RegressionTrainer(Trainer):
    # 重写计算 loss 的函数
    def compute_loss(self, model, inputs, return_outputs=False):
        # 获取标签值
        labels = inputs.get("labels")
        # 获取输入值
        x = inputs.get("pixel_values")
        # 模型输出值
        outputs = model(x)
        logits = outputs.get('logits')
        # 定义损失函数为平滑 L1 损失
        loss_fct = nn.SmoothL1Loss()
        # 计算输出值和标签的损失
        loss = loss_fct(logits, labels)
        return (loss, outputs) if return_outputs else loss

## 3.3 定义Trainer对象
trainer = RegressionTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)

## 3.4 开始训练
trainer.train()

更多 Trainer 参数参考https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments

5. 查看 Tensorboard

在当前工程目录下打开命令行执行

(hugging) PS D:\PycharmProjects\hugging> tensorboard --logdir ./logs

然后打开浏览器访问 http://localhost:6006/ 即可看到训练过程的 TensorBoard 可视化结果