Pytorchのデータ前処理【Dataset・Transform・DataLoader】を実装

前書き

今までTensorflowを活用していたのですが、toPytorchを勉強しています。
今日は基礎をざっと紹介していきます。

melheaven.hatenadiary.jp

Pytorchの機能

  • torch → Tensorの作成や操作
  • torch.Tensor → torch.Tensorクラスのパッケージ化
  • torch.autograd → 自動微分機能
  • torch.nn → ニューラルネットを構成するコンポーネントの提供
  • torch.nn.functional → 関数をメソッドして提供
  • torch.nn.init → 初期値の設定
  • torch.optim → SGDやAdam、RMSpropなど最適化オプティマイザやスケジューラなど
  • torchvision.datasets → MNISTなどのデータセットの提供
  • torchvision.models → AlexNetやRasNetなどの既存モデルネットワークの提供
  • torchvision.transforms → 画像Augumationやresize、正規化などの前処理機能の提供

Tensor

Pytorchでは入力データをTensorで扱います。
Tensorとは任意の次元の配列で、pytorch.Tensorクラスで用意されています。

Datasetの作成

Datasetとはデータと対応するラベルが1組になっているデータです。

ちなみに現状はこんな感じ。
inputは入力画像、annotationはアノテーションしたメタデータです。
どちらも画像形式をNumpy配列で表現している状態です。

print('input   :: train: %d , test: %d'%(len(input_train), len(input_val)))
print('annotation  :: train: %d , test: %d'%(len(annotation_train), len(annotation_val)))
----------------------------
input   :: train: 227 , test: 57
annotation  :: train: 227 , test: 57

訓練データと検証データ別に前処理を行なって、
前処理を終えたものをデータセットとして返しています。

train_dataset = VOCDataset(input_train, annotation_train, phase="train", transform=DataTransform(
    input_size=475, color_mean=color_mean, color_std=color_std))

val_dataset = VOCDataset(input_val, annotation_val, phase="val", transform=DataTransform(
    input_size=475, color_mean=color_mean, color_std=color_std))

VOCDatasetメソッド

VOCDatasetでは入力画像とアノテーション画像を取得して、
PIL(pillow)で画像処理をしています。
DataTransformメソッドではresizeや正規化に繋がります。

class VOCDataset(Dataset):
    def __init__(self, img_list, anno_list, phase, transform):
        self.img_list = img_list
        self.anno_list = anno_list
        self.phase = phase
        self.transform = transform

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, index):
        img, anno_class_img = self.pull_item(index)
        return img, anno_class_img

    def pull_item(self, index):
        image_file_path = self.img_list[index]
        img = Image.open(image_file_path)   # [高さ][幅][色RGB]
        anno_file_path = self.anno_list[index]
        anno_class_img = Image.open(anno_file_path)   # [高さ][幅]
        img, anno_class_img = self.transform(self.phase, img, anno_class_img)
        return img, anno_class_img

前処理Transformの定義

Transformは前処理担当。
ここで具体的に画像の正規化や拡張を定義しています。
訓練データの場合はダイナミックにAugumationをしていますが、
検証データでは再現性を意識したいので、Augumationはしません。

class DataTransform():
    def __init__(self, input_size, color_mean, color_std):
        self.data_transform = {
            'train': Compose([
                Scale(scale=[0.5, 1.5]),  # 画像の拡大
                RandomRotation(angle=[-10, 10]),  # 回転
                RandomMirror(),  # ランダムミラー
                Resize(input_size),  # リサイズ(input_size)
                Normalize_Tensor(color_mean, color_std)  # 色情報の標準化とテンソル化
            ]),
            'val': Compose([
                Resize(input_size),  # リサイズ(input_size)
                Normalize_Tensor(color_mean, color_std)  # 色情報の標準化とテンソル化
            ])
        }

    def __call__(self, phase, img, anno_class_img):
        return self.data_transform[phase](img, anno_class_img)

DataLoaderの作成

DataLoaderでDatasetからデータをバッチサイズに分けて返しています。
ここで返しているdataloaderはNumpy配列→Tensor型になっています。
学習時にtrain or valで訓練 or 検証を切り替えられるように、
dataloaders_dictを作成しておきました。
(batchsize, channel, height, width)みたいなshapeになっています。

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

val_dataloader = DataLoader(
    val_dataset, batch_size=batch_size, shuffle=False)

dataloaders_dict = {"train": train_dataloader, "val": val_dataloader}