Python

PytorchのDataloaderに関するTips

EIGHT

Posted on 8th 3月 2022 by EIGHT

今回はPytorchで入力するデータセットを作成するためのDataloaderについて解説する。

Datasetについて

Datasetとは入力するデータをすべて保存しておくためのclassである。 Datasetは非常に便利なので、是非使い方をマスターしてほしい。

Dataloaderの最小構成はこうである。


class DataSet_A(Dataset):
    def __init__(self):
        super().__init__()
                pass

    def __len__(self):
        return self.len

    def __getitem__(self, index):
                return x,t

Datasetで前処理をすべて行うべきである。例えば画像分類であれば、前処理として

画像ファイルの読み込み
画像ファイルの加工サイズをあわせるグレースケールにする等
numpyやtensorへの変換
正解ラベルの作成

などの処理が必要になるだろう。これらをDatasetに記述してゆく。

initに書くべきこと

initに書くべきことは確定している情報だ。例えば画像ファイルのresizeする規格の情報や、画像フォルダのパスなどを書くのがよいだろう。

lenに書くべきこと

lenに書くべきことは全体のデータセットのうち、どこまでを使うかだ。基本はすべてのデータを使いたいのでself.lenを返すようにしているが、例えばデータのうち最初の1000個だけを使いたい場合は1000を返せばよい。

getitemに書くべきこと

getitemに書くべきことは実際の前処理だ。 getitemの引数index番目の入力データと正解データをgetitemによって出力している。ここで、学習用のデータセットを作る場合と予測用のデータセットを作る場合で処理を共通化するために、Datasetには前処理用の関数も用意しておくべきだろう。

それらを踏まえると画像認識用のclassは次のようになる。



class DataSet_A(Dataset):
    def __init__(self):
        super().__init__()
        self.image_size = (200,200)
                # フォルダに入ってる画像パスをすべて取得する
        self.image_paths = glob.glob("/Dir/to/path/")

    def __len__(self):
        return self.len

    def __getitem__(self, index):
                # index番目の画像パスを取得
                image_path = self.image_paths[index]
                # 画像をOpenCVで読み込み
                image = cv2.imread(image_path,1)
                #画像の前処理
                image = self.reshape_image(image)

                #正解ラベル
                label = 0
                return image,label

        def reshape_image(self, image)
                #画像を加工する
                image = cv2.resize(image,dsize = self.image_size)
                ...

                return image

これで前処理とデータセットの作成は完成である。特筆すべきは、reshape_image関数だ。 getitemに前処理を全部書いてもいいのだが、このように分けておくことで予測するときに画像の前処理をDataset_A.reshape_image(image)とすることで完了できる利点がある。

Dataloaderについて

Dataloaderとは、作成したDatasetを複数個にまとめてバッチ化するclassのことだ。コードを書く時は次のようにするだけでよい。

train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)

ここまできたら後は学習するだけである。

株式会社ファントムへのお問い合わせ

群馬県でPythonを使ったAIやソフトウェアを開発している株式会社ファントムが運営しています。

Related Articles