简介:HuggingFace (transformers) 是一个广泛使用的自然语言处理库,它提供了许多预训练的模型和工具,用于快速构建和训练深度学习模型。其中,DeiT 模型是一种基于 Transformer 架构的图像变换器模型,专为视觉任务设计,包括图像分类、目标检测和分割等。而 Trainer 是 HuggingFace 提供的用于训练模型的工具之一,它可以自动化许多训练步骤,包括数据加载、模型定义、训练和评估等。
HuggingFace (transformers) 是一个广泛使用的自然语言处理库,它提供了许多预训练的模型和工具,用于快速构建和训练深度学习模型。其中,DeiT 模型是一种基于 Transformer 架构的图像变换器模型,专为视觉任务设计,包括图像分类、目标检测和分割等。而 Trainer 是 HuggingFace 提供的用于训练模型的工具之一,它可以自动化许多训练步骤,包括数据加载、模型定义、训练和评估等。
在这个教程中,我们将介绍如何使用 HuggingFace (transformers) 库中的 DeiT 模型和 Trainer 进行自定义图像数据集的训练回归任务。
datasets 模块中的 Dataset 类来创建一个自定义的图像数据集。使用上述代码中的
import datasetsfrom PIL import Imageimport requestsfrom io import BytesIOimport pandas as pdimport os# 下载并加载自定义图像数据集def load_custom_dataset(data_folder):images = []labels = []for filename in os.listdir(data_folder):if filename.endswith('.jpg'):image_path = os.path.join(data_folder, filename)with open(image_path, 'rb') as f:img = Image.open(f).convert('RGB')img = img.resize((224, 224))img = img.rotate(0)images.append(img)labels.append(int(filename.split('.')[0]))return images, labels# 定义自定义图像数据集类class CustomImageDataset(datasets.ImageFolder):def __init__(self, data_folder, split='train'):self.data_folder = data_folderself.split = splitself.images, self.labels = load_custom_dataset(data_folder)self.num_classes = len(set(self.labels))def __len__(self):return len(self.images)def __getitem__(self, idx):return self.images[idx], self.labels[idx]
CustomImageDataset 类和数据集所在文件夹路径,我们可以创建一个自定义的图像数据集实例。接下来,我们定义一个训练集和验证集的组合器类,用于在训练和验证时调用数据集的不同部分。在这里,我们使用 train 和 val 作为文件夹名来存储训练集和验证集的图像数据。