快捷方式

通过类文件对象流式传输数据

在此示例中,我们将演示如何解码流式传输的数据。也就是说,当文件不位于本地时,我们将演示如何仅下载解码所需帧的数据段。我们通过 Python 的 类文件对象 实现此功能。我们的示例使用的是视频文件,因此我们使用 VideoDecoder 类来解码它。但此处的所有经验也同样适用于音频文件和 AudioDecoder 类。

首先,一些样板代码。我们定义了两个函数:一个用于从给定 URL 下载内容,另一个用于计时给定函数的执行。

import torch
import requests
from time import perf_counter_ns


def get_url_content(url):
    response = requests.get(url, headers={"User-Agent": ""})
    if response.status_code != 200:
        raise RuntimeError(f"Failed to download video. {response.status_code = }.")
    return response.content


def bench(f, average_over=10, warmup=2):
    for _ in range(warmup):
        f()

    times = []
    for _ in range(average_over):
        start = perf_counter_ns()
        f()
        end = perf_counter_ns()
        times.append(end - start)

    times = torch.tensor(times) * 1e-6  # ns to ms
    std = times.std().item()
    med = times.median().item()
    print(f"{med = :.2f}ms +- {std:.2f}")

性能:先下载 vs. 流式传输

我们将研究在解码任何帧之前下载整个视频的成本,与在解码时能够流式传输视频数据相比。为了演示一个极端情况,我们将始终只解码视频的第一帧,同时改变我们获取视频数据的方式。

本教程中使用的视频在互联网上公开可用。我们对其进行了初步下载,以便了解其大小和内容。

from torchcodec.decoders import VideoDecoder

nasa_url = "https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4"

pre_downloaded_raw_video_bytes = get_url_content(nasa_url)
decoder = VideoDecoder(pre_downloaded_raw_video_bytes)

print(f"Video size in MB: {len(pre_downloaded_raw_video_bytes) // 1024 // 1024}")
print(decoder.metadata)
Video size in MB: 253
VideoStreamMetadata:
  duration_seconds_from_header: 206.039167
  begin_stream_seconds_from_header: 0.0
  bit_rate: 9958354.0
  codec: h264
  stream_index: 0
  begin_stream_seconds_from_content: 0.0
  end_stream_seconds_from_content: 206.039167
  width: 1920
  height: 1080
  num_frames_from_header: 6175
  num_frames_from_content: 6175
  average_fps_from_header: 29.97003
  pixel_aspect_ratio: 1
  duration_seconds: 206.039167
  begin_stream_seconds: 0.0
  end_stream_seconds: 206.039167
  num_frames: 6175
  average_fps: 29.970029921543997

我们可以看到,视频大约为 253 MB,分辨率为 1920x1080,每秒约 30 帧,时长接近 3 分半钟。由于我们只想解码第一帧,因此不下载整个视频显然会使我们受益!

让我们先测试三种场景

  1. 从我们刚刚下载的*现有*视频解码。这是我们的基准性能,因为我们将下载成本降至 0。

  2. 在解码之前下载整个视频。这是我们想避免的最坏情况。

  3. 将 URL 直接提供给 VideoDecoder 类,它会将 URL 传递给 FFmpeg。然后 FFmpeg 将决定在解码之前下载多少视频。

请注意,在我们的场景中,我们始终将 VideoDecoder 类的 seek_mode 参数设置为 "approximate"。我们这样做是为了避免在初始化期间扫描整个视频,即使我们只想解码第一帧,这也会导致下载整个视频。有关更多信息,请参阅 精确与近似搜索模式:性能和精度比较

def decode_from_existing_download():
    decoder = VideoDecoder(
        source=pre_downloaded_raw_video_bytes,
        seek_mode="approximate",
    )
    return decoder[0]


def download_before_decode():
    raw_video_bytes = get_url_content(nasa_url)
    decoder = VideoDecoder(
        source=raw_video_bytes,
        seek_mode="approximate",
    )
    return decoder[0]


def direct_url_to_ffmpeg():
    decoder = VideoDecoder(
        source=nasa_url,
        seek_mode="approximate",
    )
    return decoder[0]


print("Decode from existing download:")
bench(decode_from_existing_download)
print()

print("Download before decode:")
bench(download_before_decode)
print()

print("Direct url to FFmpeg:")
bench(direct_url_to_ffmpeg)
Decode from existing download:
med = 238.30ms +- 1.08

Download before decode:
med = 1589.95ms +- 151.48

Direct url to FFmpeg:
med = 285.81ms +- 6.19

解码已下载的视频显然是最快的。每次我们只想解码第一帧就必须下载整个视频,这比解码现有视频要慢得多。提供直接 URL 要好得多,但我们仍然可能下载了比我们需要更多的东西。

我们可以做得更好,方法是使用一个实现自己的读取和搜索方法的类文件对象,该对象仅按需从 URL 下载数据。我们不需要自己实现,而是可以使用来自 fsspec 模块的此类对象,该模块提供了 Python 的文件系统接口。请注意,使用 fsspec 库的这些功能还需要 aiohttp 模块。你可以使用 pip install fsspec aiohttp 安装两者。

import fsspec


def stream_while_decode():
    # The `client_kwargs` are passed down to the aiohttp module's client
    # session; we need to indicate that we need to trust the environment
    # settings for proxy configuration. Depending on your environment, you may
    # not need this setting.
    with fsspec.open(nasa_url, client_kwargs={'trust_env': True}) as file_like:
        decoder = VideoDecoder(file_like, seek_mode="approximate")
        return decoder[0]


print("Stream while decode: ")
bench(stream_while_decode)
Stream while decode:
med = 257.39ms +- 0.87

通过类文件对象流式传输数据比先下载视频快得多。而且它不仅比提供直接 URL 快,而且更通用。 VideoDecoder 支持直接 URL,因为底层的 FFmpeg 函数支持它们。但是支持的协议类型取决于该 FFmpeg 版本支持的内容。类文件对象可以适应任何类型的资源,包括特定于你自己的基础设施且 FFmpeg 未知的资源。

工作原理

在 Python 中,类文件对象是任何暴露用于读取、写入和搜索的特殊方法的对象。虽然这些方法显然是面向文件的,但类文件对象不必由实际文件支持。就 Python 而言,如果一个对象表现得像文件,那它就是文件。这是一个强大的概念,因为它使得读取或写入数据的库可以假定类文件接口。呈现新颖资源的其他库可以通过为它们的资源提供类文件包装器来轻松使用。

在我们的例子中,我们只需要用于解码的读取和搜索方法。所需的确切方法签名在下面的示例中。我们通过包装一个实际文件并计算每次调用每个方法的次数来演示此功能,而不是包装新颖资源。

from pathlib import Path
import tempfile

# Create a local file to interact with.
temp_dir = tempfile.mkdtemp()
nasa_video_path = Path(temp_dir) / "nasa_video.mp4"
with open(nasa_video_path, "wb") as f:
    f.write(pre_downloaded_raw_video_bytes)


# A file-like class that is backed by an actual file, but it intercepts reads
# and seeks to maintain counts.
class FileOpCounter:
    def __init__(self, file):
        self._file = file
        self.num_reads = 0
        self.num_seeks = 0

    def read(self, size: int) -> bytes:
        self.num_reads += 1
        return self._file.read(size)

    def seek(self, offset: int, whence: int) -> bytes:
        self.num_seeks += 1
        return self._file.seek(offset, whence)


# Let's now get a file-like object from our class defined above, providing it a
# reference to the file we created. We pass our file-like object to the decoder
# rather than the file itself.
file_op_counter = FileOpCounter(open(nasa_video_path, "rb"))
counter_decoder = VideoDecoder(file_op_counter, seek_mode="approximate")

print("Decoder initialization required "
      f"{file_op_counter.num_reads} reads and "
      f"{file_op_counter.num_seeks} seeks.")

init_reads = file_op_counter.num_reads
init_seeks = file_op_counter.num_seeks

first_frame = counter_decoder[0]

print("Decoding the first frame required "
      f"{file_op_counter.num_reads - init_reads} additional reads and "
      f"{file_op_counter.num_seeks - init_seeks} additional seeks.")
Decoder initialization required 9 reads and 11 seeks.
Decoding the first frame required 2 additional reads and 1 additional seeks.

虽然我们定义了一个简单的类,主要是为了演示,但它实际上对于诊断不同解码操作所需的读取和搜索量很有用。我们还引入了一个我们应该回答的神秘之处:为什么*初始化*解码器比解码第一帧需要更多的读取和搜索?答案是,在我们解码器的实现中,我们实际上调用了一个特殊的 FFmpeg 函数,该函数解码前几帧以返回更丰富的元数据。

还值得注意的是,Python 的类文件接口只是故事的一半。FFmpeg 也有自己的机制,用于在解码期间将读取和搜索定向到用户定义的函数。 VideoDecoder 对象负责将你定义的 Python 方法连接到 FFmpeg。你所要做的就是用 Python 定义你的方法,其余的交给我们。

性能:本地文件路径 vs. 本地类文件对象

由于我们定义了一个本地文件,让我们进行一次额外的性能测试。我们现在有两种方式将本地文件提供给 VideoDecoder

  1. 通过一个*路径*,其中 VideoDecoder 对象将负责打开该路径上的本地文件。

  2. 通过一个*类文件对象*,你自行打开文件并向 VideoDecoder 提供类文件对象。

一个显而易见的问题是:哪个更快?下面的代码测试了这个问题。

def decode_from_existing_file_path():
    decoder = VideoDecoder(nasa_video_path, seek_mode="approximate")
    return decoder[0]


def decode_from_existing_open_file_object():
    with open(nasa_video_path, "rb") as file:
        decoder = VideoDecoder(file, seek_mode="approximate")
        return decoder[0]


print("Decode from existing file path:")
bench(decode_from_existing_file_path)
print()

print("Decode from existing open file object:")
bench(decode_from_existing_open_file_object)
Decode from existing file path:
med = 238.10ms +- 0.61

Decode from existing open file object:
med = 238.30ms +- 1.91

值得庆幸的是,答案是这两种方式从本地文件解码所需的时间大致相同。这个结果意味着在您自己的代码中,您可以使用任何一种更方便的方式。这个结果表明,实际读取和复制数据的成本主导了解码过程中调用 Python 方法的成本。

最后,让我们清理一下我们创建的本地资源。

import shutil
shutil.rmtree(temp_dir)

脚本总运行时间: (0 分钟 40.094 秒)

由 Sphinx-Gallery 生成的画廊

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源