概要

PyTorchで動画データを読み込むにあたって，
動画をテンソル型に変換し，バッチ化するために前処理が必要となります．
動画の前処理は，

動画中のどのフレームを使用するかを決める時間方向の前処理
空間方向の前処理 (リサイズやクロップ，フリップなど)

の二つに分けられます．
今まではいずれも自分で実装していたのですが，
2番目の空間方向の前処理はnn.Sequentialを用いることで，
自前実装なしで行えるということを知ったので，こちらにメモを残しておきます．

今回使用したコードはこちらにまとめてあります．

github.com

概要

動画の前処理

前述の通り動画の前処理は二つに分けられます．

動画中のどのフレームを使用するかを決める時間方向の前処理
空間方向の前処理 (リサイズやクロップ，フリップなど)

この際に空間方向の前処理は，フレーム間の連続性を保証するために，
フレーム間で同じパラメータで前処理を行う必要があります．
例えば，動画中の奇数フレーム目だけ左右のフリップがされていて，
偶数フレームだけフリップがされていない場合，動画としてはおかしくなってしまいます．
他にも，クロップ位置がフレーム間で同じでないと，フレーム間でつながりのない動画になってしまいます．

そこで，フレーム間の前処理のパラメータが同じになるように工夫する必要があります．

PIL Imageを用いた動画の前処理方法

PIL Imageを使用した前処理では，空間方向の前処理のパラメータが同じになるように，
前処理を行う動画の１フレーム目だけパラメータをランダムで決定し，
以降のフレームでは決定したフレームを使用するという処理を行う必要があります．

空間方向の前処理を行うクラスの実装例 (水平方向のフリップ)

import random

from PIL import Image
from torchvision.transforms import functional as F
from torchvision.transforms import transforms


class RandomHorizontalFlip(transforms.RandomHorizontalFlip):
    def __init__(self, p: float = 0.5) -> None:
        super().__init__(p)
        self.randomize_parameters()

    def __call__(self, img: Image.Image) -> Image.Image:
        """
        Args:
            img: Image to be flipped.
        Returns:
            Randomly flipped image.
        """
        if self.random_p < self.p:
            return F.hflip(img)
        return img

    def randomize_parameters(self) -> None:
        self.random_p = random.random()

動画のデータセットクラスの__getitem__での前処理の実行例

import os
from typing import Any, Dict, List, Optional

import h5py
import pandas as pd
import torch
from PIL import Image
from torch.utils.data import Dataset


class VideoDataset(Dataset):
    def __init__(
        self,
        csv_file: str,
        spatial_transform: Optional[SpatialCompose] = None,
        temporal_transform: Optional[TemporalCompose] = None,
    ) -> None:
        super().__init__()
        # 動画のパスなどを保持するDataFrame
        self.df = pd.read_csv(csv_file)

        # 動画の読み込みを行うインスタンス
        self.loader = VideoLoader(temporal_transform)

        # 空間方向の前処理を行うオブジェクト
        self.spatial_transform = spatial_transform

    def __len__(self) -> int:
        return len(self.df)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        video_path = self.df.iloc[idx]["video_path"]
        name = os.path.splitext(os.path.basename(video_path))[0]

        # 動画の読み込み
        clip = self.loader(video_path)

        if self.spatial_transform is not None:
            # パラメータを一度だけランダムで指定
            self.spatial_transform.randomize_parameters()
            # 固定されたパラメータで，フレームごとに動画の前処理
            clip = [self.spatial_transform(img) for img in clip]

        clip_tensor = torch.stack(clip, dim=0).permute(1, 0, 2, 3)

        sample = {
            "clip": clip_tensor,
            "name": name,
        }

        return sample

このように動画内で前処理のパラメータが保持されるように，
前処理のクラスやデータセットクラスを実装する必要があるのですが，
torchvision.transformsで実装されている前処理のクラスを，
わざわざ継承して実装し直す必要があり，かなり面倒でした．

`nn.Sequential`を用いた動画の前処理

torchvisionのv0.8.0からtorchvision.transformsの前処理のクラスが，
PIL Imageだけでなく，torch.Tensorにも対応しました．
バッチの次元を持つtorch.Tensorにも処理を行うことが可能で，
その際バッチ内では同一のパラメータで前処理が行われるとのことです．
しかもGPUを使用することができるため，前処理が高速化も行うことができます．

import torch
import torch.nn as nn
import torchvision.transforms as T


spatial_transforms = nn.Sequential(
    T.RandomCrop(224),
    T.RandomHorizontalFlip(p=0.3),
    T.ConvertImageDtype(torch.float),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
)

batched_image = torch.randint(0, 256, size=(4, 3, 256, 256), dtype=torch.uint8)
out_image_batched = spatial_transforms(batched_image)

# GPUを使用する場合
out_image_batched_cuda = spatial_transforms(tensor_image.cuda())

本当に空間方向の前処理がバッチ内で同じかを確かめてみます．

import torch
import torch.nn as nn
from torchvision.transforms import Compose, RandomResizedCrop, ConvertImageDtype


# 適当な画像
frame = torch.randint(0, 256, (3, 512, 512))

# shape -> (16, 3, H, W)
tensor = torch.stack([frame.clone() for _ in range(16)], dim=0)

transform = nn.Sequential(
    RandomResizedCrop(224),
    ConvertImageDtype(torch.float32),
)
res = transform(tensor)

# クロップ位置が一緒か検証
for i in range(16):
    for j in range(i + 1, 16):
        # no assertion error
        assert torch.all(res[i] == res[j])

エラーは確認されず，バッチ内で前処理のパラメータが同じであることが確認できました．

速度の比較

二つの前処理の速度を比較してみました．
今回用いたコードは，こちらです．
なお動画の形式として，動画の各フレームをhdf5ファイルとしてまとめた場合と，
動画の各フレームをjpgに分割した場合の二通りで比較してみました．
データ数が128動画，入力フレームが64，バッチサイズが16で，
5エポック分データを読み込んで，1エポック平均のデータロードにかかる時間を計測しました．

github.com

CPU上での自前実装による前処理と，nn.Sequentialによる前処理の比較

いずれの場合も自前実装の方が高速になりました．
CPU上での処理に関しては自前実装で行う方が良さそうです．

---------- Experimental settings ----------
Device: cpu num_workers: 2.
n_data: 128 batch_size: 16  input_n_frames: 64.

Measuring loading time with HDF5.
---------- Start loading data ----------
PIL-based dataloader: Ave.  14.09 sec.
tensor-based dataloader: Ave.  33.91 sec.

Measuring loading time with JPG images.
---------- Start loading data ----------
PIL-based dataloader: Ave.  18.37 sec.
tensor-based dataloader: Ave.  33.58 sec.

GPU上での自前実装による前処理と，nn.Sequentialによる前処理の比較

hdf5の読み込みでは大きな差はありませんが，
画像を用いた場合はnn.Sequentialによる前処理の方が高速でした．
GPUを使用できる場合は，テンソルの計算を並列化できるため，処理が高速になっています．

---------- Experimental settings ----------
Device: cuda    num_workers: 2.
n_data: 128     batch_size: 16  input_n_frames: 64.

Measuring loading time with HDF5.
---------- Start loading data ----------
PIL-based dataloader: Ave.  14.29 sec.
tensor-based dataloader: Ave.  14.43 sec.

Measuring loading time with JPG images.
---------- Start loading data ----------
PIL-based dataloader: Ave.  18.64 sec.
tensor-based dataloader: Ave.  11.11 sec.

まとめ

nn.Sequentialを用いて動画の前処理を行う方法について紹介しました．
計算リソースが十分に使える場合は，nn.Sequentialを用いた前処理によって，実装の手間を省き，かつ学習の高速化が図れそうです．
しかしながら動画認識の深層学習モデルはメモリを多く占有してしまうため，前処理のためにメモリを割くことが難しいこともあるかもしれません．
そのような場合には自前の実装を使用するのもありかもしれません．

今回の結果は単一の実験設定でのもので，前処理を増やしたり，バッチサイズやnum_workersなどを変えることで，
結果も変わると思います．このあたりも引き続き調査していきたいです．
また，今回はjpgあるいはhdf5からの動画の読み込みを対象としていましたが，
mp4からの直接の読み込みでも比較を行いたいと思います．

Reference

リンク

yiskw note

機械学習やプログラミングについて気まぐれで書きます

【PyTorch】nn.Sequentialで動画の前処理を行う

概要

動画の前処理

PIL Imageを用いた動画の前処理方法

`nn.Sequential`を用いた動画の前処理

速度の比較

まとめ

Reference

概要

動画の前処理

PIL Imageを用いた動画の前処理方法

nn.Sequentialを用いた動画の前処理

速度の比較

まとめ

Reference

`nn.Sequential`を用いた動画の前処理