ハギングフェイスが大規模モデルトレーニング用にローカルデータセットをロードする方法

背景：

通常、huggingface を使用する場合、トレーニングのために Web サイトからデータセットを直接ロードします。コードは次のとおりです。

from datasets import load_dataset

food = load_dataset("food101")

上記のコードを記述した後、load_dataset 関数は自動的にデータセットをハグフェイスからローカルにダウンロードし、キャッシュします。

しかし、独自のデータセットをロードするにはどうすればよいでしょうか?特に画像データセットの場合、モデルに必要なデータ形式は次のとおりです。

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

上記のコードからわかるように、イメージは PIL 形式の画像であり、ローカルイメージはイメージファイルです。ローカルイメージを PIL イメージに変換するにはどうすればよいですか?

ローカルデータセットをロードする方法について説明します。

データセットのディレクトリ構造

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

ディレクトリ構造の説明

1.画像フォルダ

すべての写真を保存する

2、画像.json

マップ構造を格納する json 文字列。キーは画像の名前、値は画像に対応するラベル分類です。

{
	"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、ラベル.json

マップ構造を格納する json 文字列。キーはラベルの名前、値はラベルの分類です。

{
	"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'}

出力から、モデルに必要なデータ構造を取得したことがわかります。主に実行結果の最後から 3 行目に注目してください。

要約:

Dataset.from_generator() 関数を使用してジェネレーターを定義すると、ローカルでカスタマイズされたデータを大規模モデルで必要な任意の形式タイプに変換できます。