huggingface如何加载本地数据集进行大模型训练

背景：

一般情况下，我们使用huggingface都是从其网站上直接加载数据集进行训练，代码如下：

from datasets import load_dataset

food = load_dataset("food101")

写了上面的代码，那么load_dataset函数就会自动从huggingface上面下载数据集到本地，然后缓存起来。

可对于我们自己的数据集该如何加载呢？尤其是图片数据集，对于模型来说，他需要的数据格式如下：

{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=512x512 at 0x7F52AFC8AC50>,
 'label': 79}

从上述代码可以看到，image是一个PIL格式的图片，而我们本地是一个图片文件，如何将本地的图片转换成PIL的图片呢？

下面我们就来讲一讲如何加载本地数据集。

数据集目录结构

dataset
	--image
		--1.jsp
		--2.jsp
		--……
	--image.json
	--label.json

目录结构说明

1、image文件夹

存储所有的图片

2、image.json

存储map结构的json字符串，key是图片的名称，value是图片对应的标签分类

{
	"1.jpg": 0,
	"2.jpg": 0,
	"3.jpg": 1,
	"4.jpg": 1,
	"5.jpg": 4,
	"6.jpg": 4,
	"7.jpg": 2,
	"8.jpg": 2,
	"9.jpg": 3,
	"10.jpg": 3
}

3、label.json

存储map结构的json字符串，key是标签的名称，value是标签的分类

{
	"apple": 0,
	"pear": 1,
	"strawberry": 2,
	"peach": 3,
	"chestnut": 4
}

代码

如何将上述的目录结构存储的数据转换成模型需要的格式呢？

话不多说，直接上代码

import json
import os
from PIL import Image
from datasets import Dataset

path = 'D:/项目/AI/数据集/image/vit_dataset'


def gen(path):
    image_json = os.path.join(path, "image.json")
    with open(image_json, 'r') as f:
        # 读取JSON数据
        data = json.load(f)
    for key, value in data.items():
        imagePath = os.path.join(path, "image")
        imagePath = os.path.join(imagePath, key)
        image = Image.open(imagePath)
        yield {'image': image, 'label': value}


ds = Dataset.from_generator(gen, gen_kwargs={"path": path})
ds = ds.train_test_split(test_size=0.2)

def get_label(path):
    label_json = os.path.join(path, "label.json")
    with open(label_json, 'r') as f:
        # 读取JSON数据
        data = json.load(f)
    label2id, id2label = dict(), dict()
    for key, value in data.items():
        label2id[key] = str(value)
        id2label[str(value)] = key
    return label2id, id2label


print(ds)
print(ds['train'][0])
label2id, id2label = get_label(path)
print(label2id)
print(id2label)

运行结果：

DatasetDict({
    train: Dataset({
        features: ['image', 'label'],
        num_rows: 8
    })
    test: Dataset({
        features: ['image', 'label'],
        num_rows: 2
    })
})
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=332x332 at 0x1ED6A84C690>, 'label': 0}
{'apple': '0', 'pear': '1', 'strawberry': '2', 'peach': '3', 'chestnut': '4'}
{'0': 'apple', '1': 'pear', '2': 'strawberry', '3': 'peach', '4': 'chestnut'}

从输出结果，我们发现我们已经获得了模型需要的数据结构。主要关注运行结果的倒数第三行。

总结：

利用Dataset.from_generator()函数，通过定义一个生成器，就能根据将我们本地自定义的数据转换成大模型需要的任何的格式类型。