并行视频解码:多进程与多线程#
本教程探索将单个长视频的大量帧进行并行解码的多种方式,并比较三种策略:
FFmpeg 内部线程并行
使用 Joblib 的多进程并行
使用 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)