PyTorchのDataLoaderで動画を並列化して読み込むためのハック
Posted On 2022-07-30
動画の前処理はフレーム単位の画像処理をするためとても重いですが、特に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()
流れは次の通りです。
- joblib+lz4でフレームをすべて書き出す
- lz4を静止画として読み込む
- クォータービューの動画を上下左右半分ずつ、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の中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー