{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Accelerated Video Decoding with NVDEC\n",
        "\n",
        "This tutorial shows how to use Nvidia's hardware video decoding (NVDEC)† with TorchAudio.\n",
        "\n",
        "**NOTE**\n",
        "\n",
        "> This tutorial is authored in Google Colab, and is tailored to Google Colab's specifications.\n",
        ">\n",
        "> Please check out this tutorial in [Google Colab](https://colab.research.google.com/drive/1DDah_IaGULEO66CfQWltRqaVheBkiXdN#sandboxMode=true).\n",
        ">\n",
        "> If you install FFmpeg following this tutorial, please adjust the build configuration accordingly.\n",
        "\n",
        "To use NVDEC with TorchAudio, the following items are required.\n",
        "\n",
        "1. Nvidia GPU with hardware video encoder.\n",
        "2. FFmpeg libraries compiled with NVDEC support.\n",
        "3. PyTorch / TorchAudio with CUDA support.\n",
        "\n",
        "TorchAudio's binary distributions are compiled against FFmpeg 4 libraries, and they contain the logic required for hardware-based decoding.\n",
        "\n",
        "In the following sections, we build FFmpeg 4 libraries with NVDEC support and enable hardware acceleration through TorchAudio's `StreamReader` API. We then compare the time it takes to decode the same MP4 video with CPU and NVDEC.\n",
        "\n",
        "† For details on NVDEC and FFmpeg, please refer to the following articles.\n",
        "\n",
        "* https://docs.nvidia.com/video-technologies/video-codec-sdk/nvdec-video-decoder-api-prog-guide/\n",
        "* https://docs.nvidia.com/video-technologies/video-codec-sdk/ffmpeg-with-nvidia-gpu/#compiling-ffmpeg\n",
        "* https://developer.nvidia.com/blog/nvidia-ffmpeg-transcoding-guide/\n"
      ],
      "metadata": {
        "id": "3BxT9HiUjMyv"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08XIDshiiFp-"
      },
      "source": [
        "## Check the available GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gCegZlbgutNM",
        "outputId": "d4552883-dc7d-4989-a895-a6d6bb958e72"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Wed May 18 03:17:57 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   36C    P8    28W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S8sX7UqrSPck"
      },
      "source": [
        "## Update PyTorch and TorchAudio with nightly builds\n",
        "\n",
        "Until TorchAudio 0.12 is released, we need to use the nightly builds of PyTorch and TorchAudio."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TdQUL-fMSTsn",
        "outputId": "fc8c57be-5c95-4b3e-a436-84ad7f7152a3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found existing installation: torchaudio 0.11.0+cu113\n",
            "Uninstalling torchaudio-0.11.0+cu113:\n",
            "  Successfully uninstalled torchaudio-0.11.0+cu113\n",
            "Found existing installation: torch 1.11.0+cu113\n",
            "Uninstalling torch-1.11.0+cu113:\n",
            "  Successfully uninstalled torch-1.11.0+cu113\n",
            "Looking in indexes: https://pypi.org/simple, https://download.pytorch.org/whl/nightly/cu113\n",
            "Collecting torch\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu113/torch-1.12.0.dev20220517%2Bcu113-cp37-cp37m-linux_x86_64.whl (2098.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 2098.0 MB 7.8 kB/s \n",
            "\u001b[?25hCollecting torchaudio\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu113/torchaudio-0.12.0.dev20220517%2Bcu113-cp37-cp37m-linux_x86_64.whl (3.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.7 MB 27.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch) (4.2.0)\n",
            "Installing collected packages: torch, torchaudio\n",
            "Successfully installed torch-1.12.0.dev20220517+cu113 torchaudio-0.12.0.dev20220517+cu113\n"
          ]
        }
      ],
      "source": [
        "!pip3 uninstall -y torchaudio torch\n",
        "!pip3 install --pre torch torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu113"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RAziwx_MqJzr"
      },
      "source": [
        "## Build FFmpeg libraries with Nvidia NVDEC support\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VcKsO-zurtWM"
      },
      "source": [
        "### Install NVIDIA Video Codec Headers\n",
        "\n",
        "To build FFmpeg with NVDEC, we first install the headers that FFmpeg uses to interact with Video Codec SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vgbusLNCq2mV",
        "outputId": "9b7ef073-0ceb-468e-9ac6-438895989113"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'nv-codec-headers'...\n",
            "remote: Enumerating objects: 808, done.\u001b[K\n",
            "remote: Counting objects: 100% (808/808), done.\u001b[K\n",
            "remote: Compressing objects: 100% (688/688), done.\u001b[K\n",
            "remote: Total 808 (delta 437), reused 0 (delta 0)\u001b[K\n",
            "Receiving objects: 100% (808/808), 154.95 KiB | 755.00 KiB/s, done.\n",
            "Resolving deltas: 100% (437/437), done.\n",
            "sed 's#@@PREFIX@@#/usr/local#' ffnvcodec.pc.in > ffnvcodec.pc\n",
            "install -m 0755 -d '/usr/local/include/ffnvcodec'\n",
            "install -m 0644 include/ffnvcodec/*.h '/usr/local/include/ffnvcodec'\n",
            "install -m 0755 -d '/usr/local/lib/pkgconfig'\n",
            "install -m 0644 ffnvcodec.pc '/usr/local/lib/pkgconfig'\n"
          ]
        }
      ],
      "source": [
        "!git clone https://git.videolan.org/git/ffmpeg/nv-codec-headers.git\n",
        "!cd nv-codec-headers && sudo make install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PWT0m_LGq_GE"
      },
      "source": [
        "### Download FFmpeg source code\n",
        "\n",
        "Next we download the source code of FFmpeg 4. Any version later than 4.1 should work. We use 4.4.2 here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "AKERwitAzfLo"
      },
      "outputs": [],
      "source": [
        "!wget -q https://github.com/FFmpeg/FFmpeg/archive/refs/tags/n4.4.2.tar.gz\n",
        "!tar -xf n4.4.2.tar.gz\n",
        "!mv FFmpeg-n4.4.2 ffmpeg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpWDpxBvDy-b"
      },
      "source": [
        "### Install FFmpeg build and runtime dependencies\n",
        "\n",
        "In the later test, we use H264-encoded MP4 video streamed over HTTPS protocol, so we install the libraries for them here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vn50oz2UvvPP",
        "outputId": "99b07669-45da-47ed-ffb2-a63f2bea8175"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "... Omitted for brevity ...\n",
            "\n",
            "Setting up libx264-dev:amd64 (2:0.152.2854+gite9a5903-2) ...\n",
            "Setting up yasm (1.3.0-2build1) ...\n",
            "Setting up libunbound2:amd64 (1.6.7-1ubuntu2.4) ...\n",
            "Setting up libp11-kit-dev:amd64 (0.23.9-2ubuntu0.1) ...\n",
            "Setting up libtasn1-6-dev:amd64 (4.13-2) ...\n",
            "Setting up libtasn1-doc (4.13-2) ...\n",
            "Setting up libgnutlsxx28:amd64 (3.5.18-1ubuntu1.5) ...\n",
            "Setting up libgnutls-dane0:amd64 (3.5.18-1ubuntu1.5) ...\n",
            "Setting up libgnutls-openssl27:amd64 (3.5.18-1ubuntu1.5) ...\n",
            "Setting up libgmpxx4ldbl:amd64 (2:6.1.2+dfsg-2) ...\n",
            "Setting up libidn2-dev:amd64 (2.0.4-1.1ubuntu0.2) ...\n",
            "Setting up libidn2-0-dev (2.0.4-1.1ubuntu0.2) ...\n",
            "Setting up libgmp-dev:amd64 (2:6.1.2+dfsg-2) ...\n",
            "Setting up nettle-dev:amd64 (3.4.1-0ubuntu0.18.04.1) ...\n",
            "Setting up libgnutls28-dev:amd64 (3.5.18-1ubuntu1.5) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "Processing triggers for libc-bin (2.27-3ubuntu1.3) ...\n",
            "/sbin/ldconfig.real: /usr/local/lib/python3.7/dist-packages/ideep4py/lib/libmkldnn.so.0 is not a symbolic link\n",
            "\n"
          ]
        }
      ],
      "source": [
        "!apt update\n",
        "!apt install -y yasm libx264-dev libgnutls28-dev"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wL3lAZWPTJaN"
      },
      "source": [
        "### Configure FFmpeg build with Nvidia CUDA hardware support\n",
        "\n",
        "Next we configure FFmpeg build. Note the following:\n",
        "\n",
        "1. We provide flags like `-I/usr/local/cuda/include`, `-L/usr/local/cuda/lib64` and `--enable-nvdec` to enable NVDEC. Please check out the Transcoding Guide† for the detail.\n",
        "2. We also provide NVCC flags with compute capability 37.   \n",
        "This is because by default the configuration script verifies NVCC by compiling sample code targeting compute capability 30, which is too old for CUDA 11.\n",
        "3. Many features are disabled to reduce the compilation time.\n",
        "4. We install the library in `/usr/lib/`, which is one of the active search path of the dynamic loader.   \n",
        "Doing so allows the resulting libraries to be found without requiring a restart of the current session. This might be an undesirable location, e.g. when one isn't using a disposable VM.\n",
        "\n",
        "† NVIDIA FFmpeg Transcoding Guide https://developer.nvidia.com/blog/nvidia-ffmpeg-transcoding-guide/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KZr8bMdztRz1",
        "outputId": "b06e28b3-2c26-4073-e18a-13f2ffdeab80"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "env: ccap=37\n",
            "install prefix            /usr/\n",
            "source path               .\n",
            "C compiler                gcc\n",
            "C library                 glibc\n",
            "ARCH                      x86 (generic)\n",
            "big-endian                no\n",
            "runtime cpu detection     yes\n",
            "standalone assembly       yes\n",
            "x86 assembler             yasm\n",
            "MMX enabled               yes\n",
            "MMXEXT enabled            yes\n",
            "3DNow! enabled            yes\n",
            "3DNow! extended enabled   yes\n",
            "SSE enabled               yes\n",
            "SSSE3 enabled             yes\n",
            "AESNI enabled             yes\n",
            "AVX enabled               yes\n",
            "AVX2 enabled              yes\n",
            "AVX-512 enabled           yes\n",
            "XOP enabled               yes\n",
            "FMA3 enabled              yes\n",
            "FMA4 enabled              yes\n",
            "i686 features enabled     yes\n",
            "CMOV is fast              yes\n",
            "EBX available             yes\n",
            "EBP available             yes\n",
            "debug symbols             yes\n",
            "strip symbols             yes\n",
            "optimize for size         no\n",
            "optimizations             yes\n",
            "static                    no\n",
            "shared                    yes\n",
            "postprocessing support    no\n",
            "network support           yes\n",
            "threading support         pthreads\n",
            "safe bitstream reader     yes\n",
            "texi2html enabled         no\n",
            "perl enabled              yes\n",
            "pod2man enabled           yes\n",
            "makeinfo enabled          no\n",
            "makeinfo supports HTML    no\n",
            "\n",
            "External libraries:\n",
            "alsa                    libx264                 lzma\n",
            "bzlib                   libxcb                  zlib\n",
            "gnutls                  libxcb_shape\n",
            "iconv                   libxcb_xfixes\n",
            "\n",
            "External libraries providing hardware acceleration:\n",
            "cuda                    cuvid                   nvdec\n",
            "cuda_llvm               ffnvcodec               nvenc\n",
            "cuda_nvcc               libnpp                  v4l2_m2m\n",
            "\n",
            "Libraries:\n",
            "avcodec                 avformat                swscale\n",
            "avdevice                avutil\n",
            "avfilter                swresample\n",
            "\n",
            "Programs:\n",
            "ffmpeg                  ffprobe\n",
            "\n",
            "Enabled decoders:\n",
            "aac                     hevc                    vc1\n",
            "av1                     mjpeg                   vp8\n",
            "h263                    mpeg1video              vp9\n",
            "h264                    mpeg2video\n",
            "h264_cuvid              mpeg4\n",
            "\n",
            "Enabled encoders:\n",
            "\n",
            "Enabled hwaccels:\n",
            "av1_nvdec               mpeg1_nvdec             vp8_nvdec\n",
            "h264_nvdec              mpeg2_nvdec             vp9_nvdec\n",
            "hevc_nvdec              mpeg4_nvdec             wmv3_nvdec\n",
            "mjpeg_nvdec             vc1_nvdec\n",
            "\n",
            "Enabled parsers:\n",
            "h263                    mpeg4video              vp9\n",
            "\n",
            "Enabled demuxers:\n",
            "mov\n",
            "\n",
            "Enabled muxers:\n",
            "\n",
            "Enabled protocols:\n",
            "file                    tcp\n",
            "https                   tls\n",
            "\n",
            "Enabled filters:\n",
            "aformat                 hflip                   trim\n",
            "anull                   null                    vflip\n",
            "atrim                   scale\n",
            "format                  transpose\n",
            "\n",
            "Enabled bsfs:\n",
            "h264_mp4toannexb        null                    vp9_superframe_split\n",
            "\n",
            "Enabled indevs:\n",
            "\n",
            "Enabled outdevs:\n",
            "\n",
            "License: nonfree and unredistributable\n"
          ]
        }
      ],
      "source": [
        "# NOTE:\n",
        "# When the configure script of FFmpeg 4 checks nvcc, it uses compute\n",
        "# capability of 30 (3.0) by default. CUDA 11, however, does not support \n",
        "# compute capability 30.\n",
        "# Here, we use 37, which is supported by CUDA 11 and both K80 and T4.\n",
        "#\n",
        "# Tesla K80: 37\n",
        "# NVIDIA T4: 75\n",
        "\n",
        "%env ccap=37\n",
        "\n",
        "# NOTE:\n",
        "# We disable most of the features to speed up compilation\n",
        "# The necessary components are\n",
        "# - demuxer: mov\n",
        "# - decoder: h264\n",
        "# - gnutls (HTTPS)\n",
        "\n",
        "!cd ffmpeg && ./configure \\\n",
        "  --prefix='/usr/' \\\n",
        "  --extra-cflags='-I/usr/local/cuda/include' \\\n",
        "  --extra-ldflags='-L/usr/local/cuda/lib64' \\\n",
        "  --nvccflags=\"-gencode arch=compute_${ccap},code=sm_${ccap} -O2\" \\\n",
        "  --disable-doc \\\n",
        "  --disable-static \\\n",
        "  --disable-bsfs \\\n",
        "  --disable-decoders \\\n",
        "  --disable-encoders \\\n",
        "  --disable-filters \\\n",
        "  --disable-demuxers \\\n",
        "  --disable-devices \\\n",
        "  --disable-muxers \\\n",
        "  --disable-parsers \\\n",
        "  --disable-postproc \\\n",
        "  --disable-protocols \\\n",
        "  --enable-decoder=aac \\\n",
        "  --enable-decoder=h264 \\\n",
        "  --enable-decoder=h264_cuvid \\\n",
        "  --enable-demuxer=mov \\\n",
        "  --enable-filter=scale \\\n",
        "  --enable-protocol=file \\\n",
        "  --enable-protocol=https \\\n",
        "  --enable-gnutls \\\n",
        "  --enable-shared \\\n",
        "  --enable-gpl \\\n",
        "  --enable-nonfree \\\n",
        "  --enable-cuda-nvcc \\\n",
        "  --enable-libx264 \\\n",
        "  --enable-libnpp \\\n",
        "  --enable-nvenc \\\n",
        "  --enable-nvdec"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFTCT16ufL-e"
      },
      "source": [
        "### Build and install FFmpeg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ki9PN3V8XbSh",
        "outputId": "baf9edc7-1cb1-4fb1-b4e4-a9ffeeca2f68"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "... Omitted for brevity ...\n",
            "\n",
            "INSTALL\tlibavdevice/libavdevice.so\n",
            "INSTALL\tlibavfilter/libavfilter.so\n",
            "INSTALL\tlibavformat/libavformat.so\n",
            "INSTALL\tlibavcodec/libavcodec.so\n",
            "INSTALL\tlibswresample/libswresample.so\n",
            "INSTALL\tlibswscale/libswscale.so\n",
            "INSTALL\tlibavutil/libavutil.so\n",
            "INSTALL\tffmpeg\n",
            "INSTALL\tffprobe\n"
          ]
        }
      ],
      "source": [
        "!cd ffmpeg && make clean && make -j && make install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l1A9bI0BkWCf"
      },
      "source": [
        "## Check FFmpeg installation\n",
        "\n",
        "Let's do a quick sanity check to confirm that the FFmpeg we built works."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4fu9POO7kbQL",
        "outputId": "d7478489-9dd4-415a-eee9-49e0fd4e1220"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ffprobe version 4.4.2 Copyright (c) 2007-2021 the FFmpeg developers\n",
            "  built with gcc 7 (Ubuntu 7.5.0-3ubuntu1~18.04)\n",
            "  configuration: --prefix=/usr/ --extra-cflags=-I/usr/local/cuda/include --extra-ldflags=-L/usr/local/cuda/lib64 --nvccflags='-gencode arch=compute_37,code=sm_37 -O2' --disable-doc --disable-static --disable-bsfs --disable-decoders --disable-encoders --disable-filters --disable-demuxers --disable-devices --disable-muxers --disable-parsers --disable-postproc --disable-protocols --enable-decoder=aac --enable-decoder=h264 --enable-decoder=h264_cuvid --enable-demuxer=mov --enable-filter=scale --enable-protocol=file --enable-protocol=https --enable-gnutls --enable-shared --enable-gpl --enable-nonfree --enable-cuda-nvcc --enable-libx264 --enable-libnpp --enable-nvenc --enable-nvdec\n",
            "  libavutil      56. 70.100 / 56. 70.100\n",
            "  libavcodec     58.134.100 / 58.134.100\n",
            "  libavformat    58. 76.100 / 58. 76.100\n",
            "  libavdevice    58. 13.100 / 58. 13.100\n",
            "  libavfilter     7.110.100 /  7.110.100\n",
            "  libswscale      5.  9.100 /  5.  9.100\n",
            "  libswresample   3.  9.100 /  3.  9.100\n",
            "Decoders:\n",
            " V..... = Video\n",
            " A..... = Audio\n",
            " S..... = Subtitle\n",
            " .F.... = Frame-level multithreading\n",
            " ..S... = Slice-level multithreading\n",
            " ...X.. = Codec is experimental\n",
            " ....B. = Supports draw_horiz_band\n",
            " .....D = Supports direct rendering method 1\n",
            " ------\n",
            " V....D av1                  Alliance for Open Media AV1\n",
            " V...BD h263                 H.263 / H.263-1996, H.263+ / H.263-1998 / H.263 version 2\n",
            " VFS..D h264                 H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10\n",
            " V..... h264_cuvid           Nvidia CUVID H264 decoder (codec h264)\n",
            " VFS..D hevc                 HEVC (High Efficiency Video Coding)\n",
            " V....D mjpeg                MJPEG (Motion JPEG)\n",
            " V.S.BD mpeg1video           MPEG-1 video\n",
            " V.S.BD mpeg2video           MPEG-2 video\n",
            " VF..BD mpeg4                MPEG-4 part 2\n",
            " V....D vc1                  SMPTE VC-1\n",
            " VFS..D vp8                  On2 VP8\n",
            " VFS..D vp9                  Google VP9\n",
            " A....D aac                  AAC (Advanced Audio Coding)\n"
          ]
        }
      ],
      "source": [
        "!ffprobe -decoders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A--jYQBrfcOr",
        "outputId": "576139f5-b194-45ac-979f-e6b639b4442c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ffprobe version 4.4.2 Copyright (c) 2007-2021 the FFmpeg developers\n",
            "  built with gcc 7 (Ubuntu 7.5.0-3ubuntu1~18.04)\n",
            "  configuration: --prefix=/usr/ --extra-cflags=-I/usr/local/cuda/include --extra-ldflags=-L/usr/local/cuda/lib64 --nvccflags='-gencode arch=compute_37,code=sm_37 -O2' --disable-doc --disable-static --disable-bsfs --disable-decoders --disable-encoders --disable-filters --disable-demuxers --disable-devices --disable-muxers --disable-parsers --disable-postproc --disable-protocols --enable-decoder=aac --enable-decoder=h264 --enable-decoder=h264_cuvid --enable-demuxer=mov --enable-filter=scale --enable-protocol=file --enable-protocol=https --enable-gnutls --enable-shared --enable-gpl --enable-nonfree --enable-cuda-nvcc --enable-libx264 --enable-libnpp --enable-nvenc --enable-nvdec\n",
            "  libavutil      56. 70.100 / 56. 70.100\n",
            "  libavcodec     58.134.100 / 58.134.100\n",
            "  libavformat    58. 76.100 / 58. 76.100\n",
            "  libavdevice    58. 13.100 / 58. 13.100\n",
            "  libavfilter     7.110.100 /  7.110.100\n",
            "  libswscale      5.  9.100 /  5.  9.100\n",
            "  libswresample   3.  9.100 /  3.  9.100\n",
            "Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4':\n",
            "  Metadata:\n",
            "    major_brand     : mp42\n",
            "    minor_version   : 0\n",
            "    compatible_brands: mp42mp41\n",
            "    creation_time   : 2018-07-24T17:40:47.000000Z\n",
            "  Duration: 00:03:26.04, start: 0.000000, bitrate: 10302 kb/s\n",
            "  Stream #0:0(eng): Video: h264 (Main) (avc1 / 0x31637661), yuv420p(tv, bt709), 1920x1080 [SAR 1:1 DAR 16:9], 9958 kb/s, 29.97 fps, 29.97 tbr, 30k tbn, 60k tbc (default)\n",
            "    Metadata:\n",
            "      creation_time   : 2018-07-24T17:40:48.000000Z\n",
            "      handler_name    : ?Mainconcept Video Media Handler\n",
            "      vendor_id       : [0][0][0][0]\n",
            "      encoder         : AVC Coding\n",
            "  Stream #0:1(eng): Audio: aac (LC) (mp4a / 0x6134706D), 48000 Hz, stereo, fltp, 317 kb/s (default)\n",
            "    Metadata:\n",
            "      creation_time   : 2018-07-24T17:40:48.000000Z\n",
            "      handler_name    : #Mainconcept MP4 Sound Media Handler\n",
            "      vendor_id       : [0][0][0][0]\n"
          ]
        }
      ],
      "source": [
        "!ffprobe \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HSCZFqt95GZm"
      },
      "source": [
        "## Benchmark NVDEC with TorchAudio\n",
        "\n",
        "Now that FFmpeg and the resulting libraries are ready to use, we test NVDEC with TorchAudio. For the basics of TorchAudio's streaming API, please refer to [Streaming API tutorial](https://pytorch.org/audio/main/tutorials/streaming_api_tutorial.html).\n",
        "\n",
        "**Note**\n",
        "\n",
        "If you rebuild FFmpeg after importing class StreamReader, you'll need to restart the session to activate the newly built FFmpeg libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wDJ-qdM45fV6",
        "outputId": "80b36159-f4dd-4465-84bf-6f3e773de2ca"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.12.0.dev20220517+cu113\n",
            "0.12.0.dev20220517+cu113\n"
          ]
        }
      ],
      "source": [
        "import time\n",
        "\n",
        "import torch\n",
        "import torchaudio\n",
        "\n",
        "print(torch.__version__)\n",
        "print(torchaudio.__version__)\n",
        "\n",
        "from torchaudio.io import StreamReader"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "\n",
        "pd.set_option('display.max_rows', None)\n",
        "pd.set_option('display.max_columns', None)"
      ],
      "metadata": {
        "id": "g3TMksRk5ara"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -q -O input.mp4 \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4\""
      ],
      "metadata": {
        "id": "d0GYPyGGTz-q"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "go6xfv49doJ9"
      },
      "source": [
        "First, we define the functions we'll use for testing.\n",
        "\n",
        "Funcion `test` decodes the given source from start to end, and it reports the elapsed time, and returns one image frmae as a sample."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "GsmoRAY75T6J"
      },
      "outputs": [],
      "source": [
        "result = torch.zeros((3, 2))\n",
        "samples = [[None, None] for _ in range(3)]\n",
        "\n",
        "\n",
        "def test(src, config, i_sample):\n",
        "  print(\"=\" * 40)\n",
        "  print(\"* Configuration:\", config)\n",
        "  print(\"* Source:\", src)\n",
        "  print(\"=\" * 40)\n",
        "\n",
        "  s = StreamReader(src)\n",
        "  s.add_video_stream(5, **config)\n",
        "\n",
        "  t0 = time.monotonic()\n",
        "  num_frames = 0\n",
        "  for i, (chunk, ) in enumerate(s.stream()):\n",
        "    if i == 0:\n",
        "      print(' - Chunk:', chunk.shape, chunk.device, chunk.dtype)\n",
        "    if i == i_sample:\n",
        "      sample = chunk[0]\n",
        "    num_frames += chunk.shape[0]\n",
        "  elapsed = time.monotonic() - t0\n",
        "\n",
        "  print()\n",
        "  print(f\" - Processed {num_frames} frames.\")\n",
        "  print(f\" - Elapsed: {elapsed} seconds.\")\n",
        "  print()\n",
        "\n",
        "  return elapsed, sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGCQUW9gskBZ"
      },
      "source": [
        "### Decode MP4 from local file\n",
        "\n",
        "For the first test, we compare the time it takes for CPU and NVDEC to decode 250MB of MP4 video."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "mlrA9oVM3qvt"
      },
      "outputs": [],
      "source": [
        "local_src = \"input.mp4\"\n",
        "\n",
        "cpu_conf = {\n",
        "    \"decoder\": \"h264\",  # CPU decoding\n",
        "}\n",
        "cuda_conf = {\n",
        "    \"decoder\": \"h264_cuvid\",  # Use CUDA HW decoder\n",
        "    \"hw_accel\": \"cuda:0\",  # Then keep the memory on CUDA:0\n",
        "}\n",
        "\n",
        "i_sample = 520"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9QScuPXO3_8e"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z6kivDretHtS",
        "outputId": "9180c550-71e8-465d-c9d2-9259611ef075"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 1080, 1920]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 242.20810158899997 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test(local_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "vAEDVY0D37bJ"
      },
      "outputs": [],
      "source": [
        "result[0, 0] = elapsed\n",
        "samples[0][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jn8Wj9TN3-iY"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V6jwiJD-tW0M",
        "outputId": "80827cc0-4f38-4b96-93aa-0e172bcbf0be"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'hw_accel': 'cuda:0'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 1080, 1920]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 45.03753989300003 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test(local_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "yIA6UAaw39VC"
      },
      "outputs": [],
      "source": [
        "result[0, 1] = elapsed\n",
        "samples[0][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sYBrsBKstY0B"
      },
      "source": [
        "### Decode MP4 from network\n",
        "\n",
        "Let's run the same test on the source retrieved via network on-the-fly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "ICa21rL84QVb"
      },
      "outputs": [],
      "source": [
        "network_src = \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4\"\n",
        "i_sample = 750"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RsHY7D84YZ3"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6AGfH8X8tboy",
        "outputId": "9a6f021e-7990-4c10-a6ce-ec0e5e31b9f1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 1080, 1920]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 239.9579033 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test(network_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "4u8glKk14bYT"
      },
      "outputs": [],
      "source": [
        "result[1, 0] = elapsed\n",
        "samples[1][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jRAnSGM4cYS"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "95bky0TqtfMI",
        "outputId": "1dfc5fd2-923c-41e9-f5fc-287f2d1d1f35"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'hw_accel': 'cuda:0'}\n",
            "* Source: https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 1080, 1920]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 42.85994595500006 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test(network_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "kir92rPk4eLr"
      },
      "outputs": [],
      "source": [
        "result[1, 1] = elapsed\n",
        "samples[1][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jbCyeWFsvM6r"
      },
      "source": [
        "### Decoding and resizing\n",
        "\n",
        "In the next test, we add preprocessing. NVDEC supports several preprocessing schemes, which are also performed on the chosen hardware. For CPU, we apply the same kind of software preprocessing through FFmpeg's filter graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "XKpuH_tXvMSF"
      },
      "outputs": [],
      "source": [
        "cpu_conf = {\n",
        "    \"decoder\": \"h264\",  # CPU decoding\n",
        "    \"filter_desc\": \"scale=360:240\",  # Software filter\n",
        "}\n",
        "cuda_conf = {\n",
        "    \"decoder\": \"h264_cuvid\",  # Use CUDA HW decoder\n",
        "    \"decoder_options\": {\n",
        "        \"resize\": \"360x240\",  # Then apply HW preprocessing (resize)\n",
        "    },\n",
        "    \"hw_accel\": \"cuda:0\",  # Then keep the memory on CUDA:0\n",
        "}\n",
        "\n",
        "i_sample = 1085"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRL3C-bm4ijE"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OEtAVQVYwVCy",
        "outputId": "1b301b02-1d19-4e5d-f867-3de35957f6e7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264', 'filter_desc': 'scale=360:240'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 240, 360]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 85.71133870999984 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test(local_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "uzZMz0rW4j73"
      },
      "outputs": [],
      "source": [
        "result[2, 0] = elapsed\n",
        "samples[2][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xkrBNRNP4laK"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "XajhZb-G4mwm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8d95b66d-0338-405b-d56d-24cbed11d59e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'decoder_options': {'resize': '360x240'}, 'hw_accel': 'cuda:0'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 240, 360]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 42.306798207000156 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test(local_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "5mXoy38TwYVx"
      },
      "outputs": [],
      "source": [
        "result[2, 1] = elapsed\n",
        "samples[2][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Results\n",
        "\n",
        "The following table summarizes the time it took to decode the same media with CPU and NVDEC.\n",
        "We see significant speedup with NVDEC."
      ],
      "metadata": {
        "id": "KEA2jrrKqmr4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "res = pd.DataFrame(\n",
        "    result.numpy(),\n",
        "    index=[\"Decoding (local file)\", \"Decoding (network file)\", \"Decoding + Resize\"],\n",
        "    columns=[\"CPU\", \"NVDEC\"],\n",
        ")\n",
        "print(res)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ni0bihmogZAb",
        "outputId": "af159c73-8ed7-4ccd-fdcf-d8b471e32cc4"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "                                CPU      NVDEC\n",
            "Decoding (local file)    242.208099  45.037540\n",
            "Decoding (network file)  239.957901  42.859947\n",
            "Decoding + Resize         85.711342  42.306797\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As a reference, the following result was obtained with NVidia T4.\n",
        "\n",
        "```\n",
        "                                CPU      NVDEC\n",
        "Decoding (local file)    177.580841  13.604422\n",
        "Decoding (network file)  175.980652  17.819304\n",
        "Decoding + Resize         71.580780  10.200425\n",
        "```"
      ],
      "metadata": {
        "id": "jRcZOe0VXEC2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following code shows some frames generated by CPU decoding and NVDEC. They produce seemingly identical results."
      ],
      "metadata": {
        "id": "i71SJL8PmZfE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def yuv_to_rgb(img):\n",
        "  img = img.cpu().to(torch.float)\n",
        "  y = img[..., 0, :, :]\n",
        "  u = img[..., 1, :, :]\n",
        "  v = img[..., 2, :, :]\n",
        "\n",
        "  y /= 255\n",
        "  u = u / 255 - 0.5\n",
        "  v = v / 255 - 0.5\n",
        "\n",
        "  r = y + 1.14 * v\n",
        "  g = y + -0.396 * u - 0.581 * v\n",
        "  b = y + 2.029 * u\n",
        "\n",
        "  rgb = torch.stack([r, g, b], -1)\n",
        "  rgb = (rgb * 255).clamp(0, 255).to(torch.uint8)\n",
        "  return rgb.numpy()\n",
        "\n",
        "\n",
        "f, axs = plt.subplots(3, 2, figsize=[12.8, 14.4])\n",
        "for i in range(3):\n",
        "  for j in range(2):\n",
        "    axs[i][j].imshow(yuv_to_rgb(samples[i][j]))\n",
        "    axs[i][j].set_title(\n",
        "        f\"{'CPU' if j == 0 else 'NVDEC'}{' with resize' if i == 2 else ''}\")\n",
        "plt.plot(block=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 833
        },
        "id": "5AL9u6_xmRQa",
        "outputId": "d22e733e-469b-4bb2-afe9-8b281abfcb7a"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[]"
            ]
          },
          "metadata": {},
          "execution_count": 30
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 921.6x1036.8 with 6 Axes>"
            ],
            "image/png": "[encoded data, 550024 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Conclusion\n",
        "\n",
        "We looked at how to build FFmpeg libraries with NVDEC support and use it from TorchAudio. NVDEC provides significant speed up."
      ],
      "metadata": {
        "id": "yuYzPpV9zc4M"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "S8sX7UqrSPck",
        "VcKsO-zurtWM",
        "PWT0m_LGq_GE",
        "TpWDpxBvDy-b",
        "MFTCT16ufL-e"
      ],
      "name": "TorchAudio_Streaming_API_HW_Acceleration_Tutorial.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
