并行视频解码:多进程与多线程#

本教程探索将单个长视频的大量帧进行并行解码的多种方式,并比较三种策略:

  1. FFmpeg 内部线程并行

  2. 使用 Joblib 的多进程并行

  3. 使用 Joblib 的多线程并行

我们使用 joblib 进行并行,其 API 便于在进程或线程间分发任务;当然也可替换为其他线程/进程池管理工具。

工具函数与数据准备#

下载一个短视频并通过重复拼接生成更长的视频,以模拟需要高效处理的长视频场景。

from typing import List
import torch
import httpx
import tempfile
from pathlib import Path
import subprocess
from time import perf_counter_ns
from datetime import timedelta

from joblib import Parallel, delayed, cpu_count
from torchcodec.decoders import VideoDecoder

def bench(f, *args, num_exp=3, warmup=1, **kwargs):
    for _ in range(warmup):
        f(*args, **kwargs)
    times = []
    for _ in range(num_exp):
        start = perf_counter_ns()
        result = f(*args, **kwargs)
        end = perf_counter_ns()
        times.append(end - start)
    return torch.tensor(times).float(), result

def report_stats(times, unit="s"):
    mul = {
        "ns": 1,
        "µs": 1e-3,
        "ms": 1e-6,
        "s": 1e-9,
    }[unit]
    times = times * mul
    std = times.std().item()
    med = times.median().item()
    print(f"median = {med:.2f}{unit} ± {std:.2f}")
    return med

def split_indices(indices: List[int], num_chunks: int) -> List[List[int]]:
    chunk_size = len(indices) // num_chunks
    chunks = []
    for i in range(num_chunks - 1):
        chunks.append(indices[i * chunk_size:(i + 1) * chunk_size])
    chunks.append(indices[(num_chunks - 1) * chunk_size:])
    return chunks

def generate_long_video(temp_dir: str):
    url = "https://videos.pexels.com/video-files/854132/854132-sd_640_360_25fps.mp4"
    response = httpx.get(url, headers={"User-Agent": ""})
    if response.status_code != 200:
        raise RuntimeError(f"Failed to download video. {response.status_code = }.")
    short_video_path = Path(temp_dir) / "short_video.mp4"
    with open(short_video_path, 'wb') as f:
        for chunk in response.iter_content():
            f.write(chunk)
    long_video_path = Path(temp_dir) / "long_video.mp4"
    ffmpeg_command = [
        "ffmpeg", "-y",
        "-stream_loop", "49",
        "-i", str(short_video_path),
        "-c", "copy",
        str(long_video_path)
    ]
    subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return short_video_path, long_video_path

temp_dir = tempfile.mkdtemp()
short_video_path, long_video_path = generate_long_video(temp_dir)

decoder = VideoDecoder(long_video_path, seek_mode="approximate")
metadata = decoder.metadata

short_duration = timedelta(seconds=VideoDecoder(short_video_path).metadata.duration_seconds)
long_duration = timedelta(seconds=metadata.duration_seconds)
print(f"Original video duration: {int(short_duration.total_seconds() // 60)}m{int(short_duration.total_seconds() % 60):02d}s")
print(f"Long video duration: {int(long_duration.total_seconds() // 60)}m{int(long_duration.total_seconds() % 60):02d}s")
print(f"Video resolution: {metadata.width}x{metadata.height}")
print(f"Average FPS: {metadata.average_fps:.1f}")
print(f"Total frames: {metadata.num_frames}")

采样策略#

每 2 秒采样 1 帧,模拟仅处理子集帧的常见场景。

TARGET_FPS = 2
step = max(1, round(metadata.average_fps / TARGET_FPS))
all_indices = list(range(0, metadata.num_frames, step))

print(f"Sampling 1 frame every {TARGET_FPS} seconds")
print(f"We'll skip every {step} frames")
print(f"Total frames to decode: {len(all_indices)}")

方法一:顺序解码(基线)#

单实例顺序解码。

def decode_sequentially(indices: List[int], video_path=long_video_path):
    decoder = VideoDecoder(video_path, seek_mode="approximate")
    return decoder.get_frames_at(indices)

times, result_sequential = bench(decode_sequentially, all_indices)
sequential_time = report_stats(times, unit="s")

方法二:FFmpeg 线程并行#

通过解码器参数 num_ffmpeg_threads 使用 FFmpeg 的内部多线程能力。

def decode_with_ffmpeg_parallelism(
    indices: List[int],
    num_threads: int,
    video_path=long_video_path
):
    decoder = VideoDecoder(video_path, num_ffmpeg_threads=num_threads, seek_mode="approximate")
    return decoder.get_frames_at(indices)

NUM_CPUS = cpu_count()
times, result_ffmpeg = bench(decode_with_ffmpeg_parallelism, all_indices, num_threads=NUM_CPUS)
ffmpeg_time = report_stats(times, unit="s")
speedup = sequential_time / ffmpeg_time
print(f"Speedup compared to sequential: {speedup:.2f}x with {NUM_CPUS} FFmpeg threads.")

方法三:多进程并行#

使用 Joblib 的 loky 后端在多进程中分发任务。

def decode_with_multiprocessing(
    indices: List[int],
    num_processes: int,
    video_path=long_video_path
):
    chunks = split_indices(indices, num_chunks=num_processes)
    results = Parallel(n_jobs=num_processes, backend="loky", verbose=0)(
        delayed(decode_sequentially)(chunk, video_path) for chunk in chunks
    )
    return torch.cat([frame_batch.data for frame_batch in results], dim=0)

times, result_multiprocessing = bench(decode_with_multiprocessing, all_indices, num_processes=NUM_CPUS)
multiprocessing_time = report_stats(times, unit="s")
speedup = sequential_time / multiprocessing_time
print(f"Speedup compared to sequential: {speedup:.2f}x with {NUM_CPUS} processes.")

方法四:多线程并行#

TorchCodec 在解码时释放 GIL,多线程通常有效。

def decode_with_multithreading(
    indices: List[int],
    num_threads: int,
    video_path=long_video_path
):
    chunks = split_indices(indices, num_chunks=num_threads)
    results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)(
        delayed(decode_sequentially)(chunk, video_path) for chunk in chunks
    )
    return torch.cat([frame_batch.data for frame_batch in results], dim=0)

times, result_multithreading = bench(decode_with_multithreading, all_indices, num_threads=NUM_CPUS)
multithreading_time = report_stats(times, unit="s")
speedup = sequential_time / multithreading_time
print(f"Speedup compared to sequential: {speedup:.2f}x with {NUM_CPUS} threads.")

结果校验#

验证四种方式的结果一致性。

torch.testing.assert_close(result_sequential.data, result_ffmpeg.data, atol=0, rtol=0)
torch.testing.assert_close(result_sequential.data, result_multiprocessing, atol=0, rtol=0)
torch.testing.assert_close(result_sequential.data, result_multithreading, atol=0, rtol=0)
print("All good!")

清理临时文件#

import shutil
shutil.rmtree(temp_dir)