こしあん
2022-07-30

PyTorchのDataLoaderで動画を並列化して読み込むためのハック


4.4k{icon} {views}

動画の前処理はフレーム単位の画像処理をするためとても重いですが、特にOpenCVで動画を読み込む場合、OpenCVの特性上並列化が難しいという面倒くさい状況に遭遇します。この記事では、全フレームを書き出して、DataLoaderで並列的に扱えるようにし、joblib+ファイル圧縮がJPEGやPNGより高速に回せることを示します。

はじめに

静止画をPyTorchのDataLoaderで並列化して読み込むのは簡単です。しかし、これが動画になると急に難しくなります。どういうシチュエーションかというと、

  • mp4のような動画の各フレームに対して推論をかけたい
  • フレーム単位の読み込みは、OpenCVのcv2.VideoCaptureで行う
  • しかし、cv2.VideoCaptureが並列化に対応しておらず、ループ内で前処理を同一ループで行わないといけず、GPU使用率がなかなか上がらない

という状況です。今回紹介するのは、「フレームをすべて書き出して静止画として読み込む」方法です。これ以外にもやり方あるので参考程度にみていただければと思います。

サンプル動画を作る

まずは次のようなサンプル動画を作ります。

import cv2
import numpy as np

def make_video():
    video_writer = cv2.VideoWriter(
        "sample.webm",
        cv2.VideoWriter_fourcc(*'vp80'),
        2, (1280, 720))

    for i in range(20):
        frame = np.zeros((720, 1280, 3), np.uint8)
        for j in range(4):
            x = np.zeros((360, 640, 3), np.uint8)
            x = cv2.putText(
                x, f"{i+1:02}-{j+1}", 
                (200, 100), cv2.FONT_HERSHEY_SCRIPT_SIMPLEX, 3.0,
                (255, 255, 255), thickness=2)
            frame[360*(j//2):360*(j//2+1), 640*(j%2):640*(j%2+1), :] = x

        video_writer.write(frame)

    video_writer.release()

if __name__ == "__main__":
    make_video()

フレームを読み込んだあと、分割してバッチ化する想定なので、このようなクォータービューな動画になっています。

VideoCaptureを何も考えずにDataLoaderに読み込ませる

num_workers=0は成功する

まず、並列化しないで(num_workers=0)読み込むと、普通に成功します。

import cv2
import torch
import torch.utils.data

class VideoCaptureDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.video_reader = cv2.VideoCapture("sample.webm")
        self.n_frames = int(self.video_reader.get(cv2.CAP_PROP_FRAME_COUNT))

    def __len__(self):
        return self.n_frames

    def __getitem__(self, idx):
        success, image_frame = self.video_reader.read()
        return image_frame, idx

def main():
    dataset = VideoCaptureDataset()
    dataloader = torch.utils.data.DataLoader(
        dataset=dataset, batch_size=1, shuffle=False, num_workers=0)

    for frame, idx in dataloader:
        print(idx, frame.shape)

if __name__ == "__main__":
    main()

コンソールの出力は次のようになります。事実上のシーケンシャルリードなので、順序も何も問題もありません。

tensor([0]) torch.Size([1, 720, 1280, 3])
tensor([1]) torch.Size([1, 720, 1280, 3])
:   :  :
tensor([18]) torch.Size([1, 720, 1280, 3])
tensor([19]) torch.Size([1, 720, 1280, 3])

num_workers=1以上にすると失敗する

同一コードで、「num_workers=1以上」にしてみます。

    dataloader = torch.utils.data.DataLoader(
        dataset=dataset, batch_size=1, shuffle=False, num_workers=1)

すると次のようなエラーが出ます。

  File "C:\Program Files\Python37\lib\multiprocessing\process.py", line 112, in start
    self._popen = self._Popen(self)
  File "C:\Program Files\Python37\lib\multiprocessing\context.py", line 223, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "C:\Program Files\Python37\lib\multiprocessing\context.py", line 322, in _Popen
    return Popen(process_obj)
  File "C:\Program Files\Python37\lib\multiprocessing\popen_spawn_win32.py", line 89, in __init__
    reduction.dump(process_obj, to_child)
  File "C:\Program Files\Python37\lib\multiprocessing\reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
TypeError: can't pickle cv2.VideoCapture objects

これは何を言いたいのかというと、「VideoCaptureがマルチプロセスに対応していないため、並列化できない」ということです。VideoCaptureのやっていることが、1フレーム目から順にフレームを読んでいくという、シーケンシャルリードなので、並列化は難しいだろうなというのが直感的にもわかります。

個人的に興味深かったのは、「num_workers=1でも失敗する」ということです。num_workers=1というのは、前処理のプロセスを外だしして、ニューラルネットワークのプロセスと前処理のプロセスを切り分けて効率化しようというものですが、これすら認めてくれないのは意外でした。

VideoCaptureを並列化できない対策

この対策はいくつか考えられます。

動画を分割する

例えば、0~1000フレーム、1000~2000フレームのように動画自体を分割してしまう方法です。

わかりやすい方法ですが、分割した動画を読み込むときに同じ問題に直面します。

並列化対応可能な動画読み込みライブラリを探す

直球な方法です。一例としては、torchvisionのVideoReaderを使う方法です。

ただし、このVideoReaderが2022年7月のv0.13.1時点では、pip installで落としてきたtorchvisionには付属しておらず、VideoReaderを使いたい場合はtorchvisionをソースコードからビルドする必要があります

特にCUDA対応したい場合、かなり大変なハマり方をしたので(エラーが出るときと出ないときがある)、今回は割愛します。

一旦フレーム単位で書き出して、画像として読み込む

おそらく一番わかりやすく、最もよく使われている方法ではないかと思います。画像として読み込めば並列化は可能です。今回紹介するのはこのやり方です。

ただし、このやり方は3つのデメリットがあります。

  • フレームの書き出し、読み込みのI/Oのボトルネックが発生する
  • フレームを圧縮かけないで書き出せばI/Oボトルネックは軽くなるが、ストレージが大量に必要になる
  • JPEGとして書き出せば速いし軽いが、画質が悪化する

最近はクラウド等でもSSDのサーバーが多くなってきて、I/Oボトルネックをそこまで気にしなくてよくなりましたが、圧縮かけると極端にスループットが落ちるため注意が必要です。

つまり、I/Oボトルネックやシリアライザーの性能について気にしながら、読み書きしないといけなくなります。

OpenCV単体のスループット

まずOpenCV単体のスループットを検証します。検証環境は次のとおりです。

  • ColabのCPUインスタンス(CPUは2個)
  • HD(1920×1080)の動画を読み込みしてFPSを計測

YouTubeからHAMBURGER BOYSの動画をお借りしました。

YouTubeからのダウンロード

yt-dlpをインストールします

!pip install yt-dlp

動画の画質リストを取得します

!yt-dlp -F https://www.youtube.com/watch?v=ntZRtGpoEQo
!yt-dlp -F https://www.youtube.com/watch?v=ntZRtGpoEQo
[youtube] ntZRtGpoEQo: Downloading webpage
[youtube] ntZRtGpoEQo: Downloading android player API JSON
[info] Available formats for ntZRtGpoEQo:
ID  EXT   RESOLUTION FPS │   FILESIZE    TBR PROTO │ VCODEC         VBR ACODEC      ABR ASR MORE INFO
─────────────────────────────────────────────────────────────────────────────────────────────────────────────
sb2 mhtml 48x27        0 │                   mhtml │ images                                 storyboard
sb1 mhtml 80x45        1 │                   mhtml │ images                                 storyboard
sb0 mhtml 160x90       1 │                   mhtml │ images                                 storyboard
139 m4a   audio only     │    1.46MiB    49k https │ audio only         mp4a.40.5   49k 22k low, m4a_dash
249 webm  audio only     │    1.47MiB    49k https │ audio only         opus        49k 48k low, webm_dash
250 webm  audio only     │    1.93MiB    65k https │ audio only         opus        65k 48k low, webm_dash
140 m4a   audio only     │    3.88MiB   129k https │ audio only         mp4a.40.2  129k 44k medium, m4a_dash
251 webm  audio only     │    3.82MiB   128k https │ audio only         opus       128k 48k medium, webm_dash
17  3gp   176x144     12 │    2.36MiB    79k https │ mp4v.20.3      79k mp4a.40.2    0k 22k 144p
160 mp4   256x144     24 │    2.21MiB    74k https │ avc1.4d400c    74k video only          144p, mp4_dash
278 webm  256x144     24 │    2.52MiB    84k https │ vp9            84k video only          144p, webm_dash
133 mp4   426x240     24 │    4.14MiB   138k https │ avc1.4d4015   138k video only          240p, mp4_dash
242 webm  426x240     24 │    5.42MiB   181k https │ vp9           181k video only          240p, webm_dash
134 mp4   640x360     24 │   10.14MiB   338k https │ avc1.4d401e   338k video only          360p, mp4_dash
18  mp4   640x360     24 │   16.51MiB   551k https │ avc1.42001E   551k mp4a.40.2    0k 44k 360p
243 webm  640x360     24 │    9.94MiB   332k https │ vp9           332k video only          360p, webm_dash
135 mp4   854x480     24 │   20.67MiB   690k https │ avc1.4d401e   690k video only          480p, mp4_dash
244 webm  854x480     24 │   18.13MiB   605k https │ vp9           605k video only          480p, webm_dash
22  mp4   1280x720    24 │ ~ 42.95MiB  1402k https │ avc1.64001F  1402k mp4a.40.2    0k 44k 720p
136 mp4   1280x720    24 │   38.11MiB  1273k https │ avc1.4d401f  1273k video only          720p, mp4_dash
247 webm  1280x720    24 │   37.45MiB  1250k https │ vp9          1250k video only          720p, webm_dash
137 mp4   1920x1080   24 │   67.49MiB  2253k https │ avc1.640028  2253k video only          1080p, mp4_dash
248 webm  1920x1080   24 │   66.96MiB  2236k https │ vp9          2236k video only          1080p, webm_dash
271 webm  2560x1440   24 │  171.33MiB  5720k https │ vp9          5720k video only          1440p, webm_dash
313 webm  3840x2160   24 │  422.49MiB 14106k https │ vp9         14106k video only          2160p, webm_dash

HD+mp4は「137」でした。4Kとかの高解像度だとWebM形式で圧縮されていました。

ダウンロード

!yt-dlp -f 137 https://www.youtube.com/watch?v=ntZRtGpoEQo
[youtube] ntZRtGpoEQo: Downloading webpage
[youtube] ntZRtGpoEQo: Downloading android player API JSON
[info] ntZRtGpoEQo: Downloading 1 format(s): 137
[download] Destination: HAMBURGER BOYS _ 厚岸オイスタースタイル -厚岸町-[Official Video] [ntZRtGpoEQo].mp4
[download] 100% of 67.49MiB in 00:18

パスを変数に登録しておきましょう。

video_path = "HAMBURGER BOYS _ 厚岸オイスタースタイル -厚岸町-[Official Video] [ntZRtGpoEQo].mp4"

スループット計測

import cv2
import time

video_cap = cv2.VideoCapture(video_path)
print("Num frames :", video_cap.get(cv2.CAP_PROP_FRAME_COUNT), "\n")

start_time = time.time()
cnt = 0
while video_cap.isOpened():
    success, image_frame = video_cap.read()
    if not success: break
    if cnt == 0:
        print(image_frame.shape)
    elif cnt % 1000 == 0:
        print(cnt, cnt/(time.time()-start_time), "FPS")
    cnt += 1

print(cnt, cnt/(time.time()-start_time), "FPS")

このように適当に読んでいきます。正直スループットは環境によりけりですが、ColabのCPUインスタンス(メモリは通常)の場合、

Num frames : 6030.0 

(1080, 1920, 3)
1000 124.10039284070153 FPS
2000 122.18650210846748 FPS
3000 123.89120631731222 FPS
4000 122.02155541900483 FPS
5000 122.1418569193976 FPS
6000 122.83553509573085 FPS
6030 123.09827719719543 FPS

123.09FPSでした。shapeが(1080, 1920, 3)なので、HD解像度であるのが確認できます。

画像として読み書きする場合のスループット

わかりやすいのが、JPEGかPNGにフレーム単位で出力してしまう方法です。この読み書きのスループットとサイズを計測してみましょう。画像の読み書きはOpenCVを使っています

※かなり適当に書いたコードです

import os
import glob

def write_image_throughput(format="jpg", quality=90):
    print("Format :", format, " / Quality :", quality)

    os.system("rm -rf _FRAMES")
    os.makedirs("_FRAMES")

    video_cap = cv2.VideoCapture(video_path)

    # Write Frames
    start_time = time.time()
    cnt = 0
    while video_cap.isOpened():
        success, image_frame = video_cap.read()
        if not success: break
        if format=="jpg":
            cv2.imwrite(f"_FRAMES/{cnt:05}."+format, image_frame, 
                        [int(cv2.IMWRITE_JPEG_QUALITY), quality])
        elif format=="png":
            cv2.imwrite(f"_FRAMES/{cnt:05}."+format, image_frame)
        else:
            raise NotImplementedError()
        cnt += 1

    print("Write throughput : ", cnt/(time.time()-start_time), "FPS")

    # Read Frames
    files = sorted(glob.glob("_FRAMES/*"))
    start_time = time.time()
    for cnt, f in enumerate(files):
        image_frame = cv2.imread(f)

    print("Read throughput : ", cnt/(time.time()-start_time), "FPS")

    # Filesize
    print("Directory size (MB):")
    !du -sm _FRAMES/
    print("")

    # Copy sample
    os.system(f"cp _FRAMES/00741.{format} 0741_{format}_{quality}.{format}")

formats = ["jpg", "jpg", "jpg", "jpg", "jpg", "png"]
qualities = [70, 80, 90, 95, 100, -1]
for f, q in zip(formats, qualities):
    write_image_throughput(f, q)

画像として出力する場合のスループット/容量一覧

ファイル形式 品質 容量(MB) Write FPS Read FPS
オンメモリ 123.09
JPEG 70 774 25.66 41.55
JPEG 80 934 25.37 41.07
JPEG 90 1307 24.77 39.01
JPEG 95 1834 24.01 36.82
JPEG 100 3477 20.22 29.55
PNG 7918 12.12 24.3

FPSは容量とほぼ連動する形になりました。PNGになると可逆圧縮ですが、流石にスループットが犠牲になりますね。

画質ごとのフレームサンプル

各画質ごとの特定フレームを切り出した結果を示します

joblibで圧縮をかける

もう1個のやり方が、joblibのようなシリアライザーを使う方法です。これもそのままだと容量が膨大になってしまうので、何らかの圧縮をかけます(ファイル圧縮なのでこれらはすべて可逆圧縮です)。PNGよりも速い方法があればこちらのほうが効率良いことになります。

joblibの公式ドキュメントに圧縮を使ったI/Oの改善方法がありました。これによると、GZipなどよりも、LZ4のほうが効率が良いということが書かれています。過去の私の経験的にこれはそうなので、その通りの結果が出てくると思います。検証してみましょう。

import joblib

def write_joblib_throughput(compress="raw", ratio=None):
    print("Compress :", compress, " / Ratio :", ratio)

    os.system("rm -rf _FRAMES")
    os.makedirs("_FRAMES")

    video_cap = cv2.VideoCapture(video_path)

    # Write Frames
    start_time = time.time()
    cnt = 0
    while video_cap.isOpened():
        success, image_frame = video_cap.read()
        if not success: break
        if compress=="raw":
            joblib.dump(image_frame, f"_FRAMES/{cnt:05}.pkl")
        elif compress in ["gz", "lz4"]:
            joblib.dump(image_frame, f"_FRAMES/{cnt:05}.pkl."+compress, compress=ratio)
        else:
            raise NotImplementedError()
        cnt += 1

    print("Write throughput : ", cnt/(time.time()-start_time), "FPS")

    # Read Frames
    files = sorted(glob.glob("_FRAMES/*"))
    start_time = time.time()
    for cnt, f in enumerate(files):
        image_frame = joblib.load(f)

    print("Read throughput : ", cnt/(time.time()-start_time), "FPS")

    # Filesize
    print("Directory size (MB):")
    !du -sm _FRAMES/
    print("")

compresses = ["raw", "gz", "gz", "lz4", "lz4", "lz4", "lz4"]
ratios = [None, ("gzip", 1), ("gzip", 3), ("lz4", 1), ("lz4", 2), ("lz4", 3), ("lz4", 5)]
for c, r in zip(compresses, ratios):
    write_joblib_throughput(c, r)

joblib圧縮の結果

ファイル形式 圧縮率 容量(MB) Write FPS Read FPS
圧縮なし 35780 37.52 33.97
gzip 1 8698 8.98 24.82
gzip 3 8292 8.25 26.16
lz4 1 11714 37.23 53.88
lz4 2 11714 39.66 48.51
lz4 3 9635 11.58 58.11
lz4 5 9459 9.44 55.41
PNG(可逆) 7918 12.12 24.3
JPEG(非可逆) 70 774 25.66 41.55

gzipはpngより重い、遅いで特に良いところはありませんが、lz4は特に良いところがあります。

PNGよりも1.47倍重くなるのを許容すれば、PNGどころか、非可逆のJPEGよりも速い結果となりました。これは活用できそうです。

動画のDataLoaderの並列化へ

先にコードを示します。冒頭の例を並列可能な形に書き換えたものがこちらです。

import cv2
import torch
import torch.utils.data
import os
import glob
import joblib
import numpy as np

def write_frames():
    video_cap = cv2.VideoCapture("sample.webm")
    os.makedirs("frames", exist_ok=True)

    frame_cnt = 0
    while video_cap.isOpened():
        success, image_frame = video_cap.read()
        if not success: break
        joblib.dump(image_frame, f"frames/{frame_cnt:03}.lz4", compress="lz4")
        frame_cnt += 1

class VideoFrameDataset(torch.utils.data.Dataset):
    def __init__(self):
        self.files = sorted(glob.glob("frames/*.lz4"))

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

    def __getitem__(self, idx):
        data = joblib.load(f"frames/{idx:03}.lz4")
        data = data.reshape(2, 360, 2, 640, 3).transpose(
            [0, 2, 1, 3, 4]).reshape(4, 360, 640, 3)
        indices = 10*idx + np.arange(4)
        return data, indices

def collate_fn(batches):
    images, indices = [], []
    for b in batches:
        images.append(b[0])
        indices.append(b[1])
    images = np.concatenate(images, axis=0)
    indices = np.concatenate(indices, axis=0)
    return images, indices   

def main():
    write_frames()
    dataset = VideoFrameDataset()
    dataloader = torch.utils.data.DataLoader(
        dataset=dataset, batch_size=2, shuffle=False, 
        num_workers=4, collate_fn=collate_fn)

    video_writer = cv2.VideoWriter(
        "dataloader_worker4.webm",
        cv2.VideoWriter_fourcc(*'vp80'),
        6, (640, 360)
    )

    for batch in dataloader:
        frames, indices = batch
        for i in range(frames.shape[0]):
            video_writer.write(frames[i])
    video_writer.release()

if __name__ == "__main__":
    main()

流れは次の通りです。

  1. joblib+lz4でフレームをすべて書き出す
  2. lz4を静止画として読み込む
  3. クォータービューの動画を上下左右半分ずつ、4個に分割
  4. バッチ化

というものです。生成動画は次の通りです。

num_workers=4で並列化されているのに、動画を取り扱うことができました

並列化しているのに順序が保証されている

ここで面白いのが、num_workersを1以上にしてもバッチ内、バッチ間の順番が保証されている点です。DataLoaderの値を変えて、インデックスをみてみましょう。

// いずれも同じ出力
// batch_size=1, num_workers=2
// batch_size=1, num_workers=3
// batch_size=1, num_workers=4
[0 1 2 3]
[10 11 12 13]
[20 21 22 23]
[30 31 32 33]
[40 41 42 43]
[50 51 52 53]
[60 61 62 63]
[70 71 72 73]
[80 81 82 83]
[90 91 92 93]
[100 101 102 103]
[110 111 112 113]
[120 121 122 123]
[130 131 132 133]
[140 141 142 143]
[150 151 152 153]
[160 161 162 163]
[170 171 172 173]
[180 181 182 183]
[190 191 192 193]

// すべて同じ出力
// batch_size=2, num_workers=2
// batch_size=2, num_workers=3
// batch_size=2, num_workers=4
[ 0  1  2  3 10 11 12 13]
[20 21 22 23 30 31 32 33]
[40 41 42 43 50 51 52 53]
[60 61 62 63 70 71 72 73]
[80 81 82 83 90 91 92 93]
[100 101 102 103 110 111 112 113]
[120 121 122 123 130 131 132 133]
[140 141 142 143 150 151 152 153]
[160 161 162 163 170 171 172 173]
[180 181 182 183 190 191 192 193]

通常並列処理というと、順番が保証されないのですが、PyTorchのDataLoaderで並列化すると、PyTorch側でよしなにやってくれて、順番が保証されるという特徴があるそうです。バッチサイズやnum_workersに関わりなく、実質シーケンシャルリードとして扱えるのは、フレーム間の順序が大事な動画では便利そうですね。

まとめ

この記事では、動画のフレームをいったん全部ファイルとして書き出して、DataLoaderに静止画として読ませることで並列化する手法を検証しました。動画の前処理を並列化する方法はこれだけではありませんが、これが必要なレベルで前処理が重かったので何らかの役に立てば幸いです。

現実的には、

  • JPEGで吐き出す(画質落ちてもいいから、容量重視)
  • LZ4で吐き出す(容量いっぱいとってもいいから、速度と画質重視)

この2つではないかと思います。JPEGで良ければそれでOKですが、見ての通りJPEGの品質とスループットは逆相関の関係にあるので、JPEGもJPEGでは難しい点はあると思います。

最後にすべてのパターンのスループットの結果をまとめて終わりにします。

ファイル形式 可逆か? 品質 容量(MB) Write FPS Read FPS
オンメモリ 可逆 123.09
JPEG 非可逆 70 774 25.66 41.55
JPEG 非可逆 80 934 25.37 41.07
JPEG 非可逆 90 1307 24.77 39.01
JPEG 非可逆 95 1834 24.01 36.82
JPEG 非可逆 100 3477 20.22 29.55
PNG 可逆 7918 12.12 24.3
圧縮なし 可逆 35780 37.52 33.97
gzip 可逆 1 8698 8.98 24.82
gzip 可逆 3 8292 8.25 26.16
lz4 可逆 1 11714 37.23 53.88
lz4 可逆 2 11714 39.66 48.51
lz4 可逆 3 9635 11.58 58.11
lz4 可逆 5 9459 9.44 55.41


Shikoan's ML Blogの中の人が運営しているサークル「じゅ~しぃ~すくりぷと」の本のご案内

技術書コーナー

北海道の駅巡りコーナー


Add a Comment

メールアドレスが公開されることはありません。 が付いている欄は必須項目です