{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3BxT9HiUjMyv"
      },
      "source": [
        "# GPU video decoder/encoder\n",
        "\n",
        "**Author**: [Moto Hira](moto@meta.com)\n",
        "\n",
        "This tutorial shows how to use NVIDIA's hardware video decoder (NVDEC) and encoder (NVENC) with TorchAudio.\n",
        "\n",
        "Using hardware encoder/decoder improves the speed of loading and saving certain types of videos.\n",
        "\n",
        "Using them in TorchAduio requires FFmpeg built with NVENC/NVDEC support. Google Colab has such FFmpeg pre-installed, so you can [run this tutorial on Google Colab](https://colab.research.google.com/drive/1DDah_IaGULEO66CfQWltRqaVheBkiXdN#sandboxMode=true).\n",
        "\n",
        "If you want to enable GPU decoding/encoding, please refer to [Enabling GPU video decoder/encoder](./build.ffmpeg.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Update PyTorch and TorchAudio with nightly builds\n",
        "\n",
        "Until PyTorch 2.0 is released, we need to use the nightly builds of PyTorch and TorchAudio."
      ],
      "metadata": {
        "id": "8EI-DwaeQbjp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fACIAl1dgYt-",
        "outputId": "94c153e5-5260-4b83-9bc4-64d40c23e25f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/, https://download.pytorch.org/whl/nightly/cu117\n",
            "Collecting torch\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu117/torch-2.0.0.dev20230209%2Bcu117-cp38-cp38-linux_x86_64.whl (1837.7 MB)\n",
            "Collecting torchaudio\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cu117/torchaudio-2.0.0.dev20230208%2Bcu117-cp38-cp38-linux_x86_64.whl (4.4 MB)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from torch) (3.9.0)\n",
            "Collecting pytorch-triton==2.0.0+0d7e753227\n",
            "  Downloading https://download.pytorch.org/whl/nightly/pytorch_triton-2.0.0%2B0d7e753227-cp38-cp38-linux_x86_64.whl (18.7 MB)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch) (4.4.0)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.8/dist-packages (from torch) (3.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.8/dist-packages (from torch) (1.7.1)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.8/dist-packages (from pytorch-triton==2.0.0+0d7e753227->torch) (3.22.6)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.8/dist-packages (from sympy->torch) (1.2.1)\n",
            "Installing collected packages: pytorch-triton, torch, torchaudio\n",
            "Successfully installed pytorch-triton-2.0.0+0d7e753227 torch-2.0.0.dev20230209+cu117 torchaudio-2.0.0.dev20230208+cu117\n"
          ]
        }
      ],
      "source": [
        "!pip uninstall -y -q torch torchaudio torchvision torchtext\n",
        "!pip install --progress-bar off --pre torch torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu117 2> /dev/null"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install third party libraries and download assets"
      ],
      "metadata": {
        "id": "uC1wdJP35Z7l"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lbA3S6p_Z2lI",
        "outputId": "d8282d99-3090-4cea-b66f-91da45b21764"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting boto3\n",
            "  Downloading boto3-1.26.67-py3-none-any.whl (132 kB)\n",
            "Collecting botocore<1.30.0,>=1.29.67\n",
            "  Downloading botocore-1.29.67-py3-none-any.whl (10.4 MB)\n",
            "Collecting jmespath<2.0.0,>=0.7.1\n",
            "  Downloading jmespath-1.0.1-py3-none-any.whl (20 kB)\n",
            "Collecting s3transfer<0.7.0,>=0.6.0\n",
            "  Downloading s3transfer-0.6.0-py3-none-any.whl (79 kB)\n",
            "Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /usr/local/lib/python3.8/dist-packages (from botocore<1.30.0,>=1.29.67->boto3) (2.8.2)\n",
            "Collecting urllib3<1.27,>=1.25.4\n",
            "  Downloading urllib3-1.26.14-py2.py3-none-any.whl (140 kB)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil<3.0.0,>=2.1->botocore<1.30.0,>=1.29.67->boto3) (1.15.0)\n",
            "Installing collected packages: urllib3, jmespath, botocore, s3transfer, boto3\n",
            "  Attempting uninstall: urllib3\n",
            "    Found existing installation: urllib3 1.24.3\n",
            "    Uninstalling urllib3-1.24.3:\n",
            "      Successfully uninstalled urllib3-1.24.3\n",
            "Successfully installed boto3-1.26.67 botocore-1.29.67 jmespath-1.0.1 s3transfer-0.6.0 urllib3-1.26.14\n"
          ]
        }
      ],
      "source": [
        "!pip3 install --progress-bar off boto3 2> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "d0GYPyGGTz-q"
      },
      "outputs": [],
      "source": [
        "!wget -q -O input.mp4 \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\""
      ]
    },
    {
      "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": "7415cdd3-d8e8-476d-da13-27fa87267497"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Thu Feb  9 15:54:05 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |\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 T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   58C    P0    28W /  70W |      0MiB / 15360MiB |      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": "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": 55,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4fu9POO7kbQL",
        "outputId": "3514bcc9-17e7-47a2-bb7e-a8162ccb1eb4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " V..... h264_cuvid           Nvidia CUVID H264 decoder (codec h264)\n",
            " V..... hevc_cuvid           Nvidia CUVID HEVC decoder (codec hevc)\n",
            " V..... mjpeg_cuvid          Nvidia CUVID MJPEG decoder (codec mjpeg)\n",
            " V..... mpeg1_cuvid          Nvidia CUVID MPEG1VIDEO decoder (codec mpeg1video)\n",
            " V..... mpeg2_cuvid          Nvidia CUVID MPEG2VIDEO decoder (codec mpeg2video)\n",
            " V..... mpeg4_cuvid          Nvidia CUVID MPEG4 decoder (codec mpeg4)\n",
            " V..... vc1_cuvid            Nvidia CUVID VC1 decoder (codec vc1)\n",
            " V..... vp8_cuvid            Nvidia CUVID VP8 decoder (codec vp8)\n",
            " V..... vp9_cuvid            Nvidia CUVID VP9 decoder (codec vp9)\n"
          ]
        }
      ],
      "source": [
        "!ffmpeg -hide_banner -decoders | grep -i cuvid"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "499UCyKtFPrO",
        "outputId": "06fe7564-5a17-453c-efbf-b04136652586"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " V..... h264_nvenc           NVIDIA NVENC H.264 encoder (codec h264)\n",
            " V..... nvenc                NVIDIA NVENC H.264 encoder (codec h264)\n",
            " V..... nvenc_h264           NVIDIA NVENC H.264 encoder (codec h264)\n",
            " V..... nvenc_hevc           NVIDIA NVENC hevc encoder (codec hevc)\n",
            " V..... hevc_nvenc           NVIDIA NVENC hevc encoder (codec hevc)\n"
          ]
        }
      ],
      "source": [
        "!ffmpeg -hide_banner -encoders | grep -i nvenc"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following command fetches video from remote server, decode with NVDEC (cuvid) and re-encode with NVENC. If this command does not work, then there is an issue with FFmpeg installation, and TorchAudio would not be able to use them either."
      ],
      "metadata": {
        "id": "yZunKfgUb-i_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A--jYQBrfcOr",
        "outputId": "ba32eea4-52c5-4864-e7ba-45d2200b43ed"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "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_small.mp4':\n",
            "  Metadata:\n",
            "    major_brand     : mp42\n",
            "    minor_version   : 512\n",
            "    compatible_brands: mp42iso2avc1mp41\n",
            "    encoder         : Lavf58.76.100\n",
            "  Duration: 00:03:26.04, start: 0.000000, bitrate: 1294 kb/s\n",
            "    Stream #0:0(eng): Video: h264 (High) (avc1 / 0x31637661), yuv420p(tv, bt709), 960x540 [SAR 1:1 DAR 16:9], 1156 kb/s, 29.97 fps, 29.97 tbr, 30k tbn, 59.94 tbc (default)\n",
            "    Metadata:\n",
            "      handler_name    : ?Mainconcept Video Media Handler\n",
            "    Stream #0:1(eng): Audio: aac (LC) (mp4a / 0x6134706D), 48000 Hz, stereo, fltp, 128 kb/s (default)\n",
            "    Metadata:\n",
            "      handler_name    : #Mainconcept MP4 Sound Media Handler\n",
            "Stream mapping:\n",
            "  Stream #0:0 -> #0:0 (h264 (h264_cuvid) -> h264 (h264_nvenc))\n",
            "  Stream #0:1 -> #0:1 (copy)\n",
            "Press [q] to stop, [?] for help\n",
            "Output #0, mp4, to 'test.mp4':\n",
            "  Metadata:\n",
            "    major_brand     : mp42\n",
            "    minor_version   : 512\n",
            "    compatible_brands: mp42iso2avc1mp41\n",
            "    encoder         : Lavf58.29.100\n",
            "    Stream #0:0(eng): Video: h264 (h264_nvenc) (Main) (avc1 / 0x31637661), cuda, 360x240 [SAR 1:1 DAR 3:2], q=-1--1, 5000 kb/s, 29.97 fps, 30k tbn, 29.97 tbc (default)\n",
            "    Metadata:\n",
            "      handler_name    : ?Mainconcept Video Media Handler\n",
            "      encoder         : Lavc58.54.100 h264_nvenc\n",
            "    Side data:\n",
            "      cpb: bitrate max/min/avg: 0/0/5000000 buffer size: 10000000 vbv_delay: -1\n",
            "    Stream #0:1(eng): Audio: aac (LC) (mp4a / 0x6134706D), 48000 Hz, stereo, fltp, 128 kb/s (default)\n",
            "    Metadata:\n",
            "      handler_name    : #Mainconcept MP4 Sound Media Handler\n",
            "frame= 6175 fps=1673 q=8.0 Lsize=   41278kB time=00:03:26.17 bitrate=1640.1kbits/s speed=55.9x    \n",
            "video:37869kB audio:3234kB subtitle:0kB other streams:0kB global headers:0kB muxing overhead: 0.424815%\n"
          ]
        }
      ],
      "source": [
        "!ffmpeg -hide_banner -y -vsync 0 -hwaccel cuvid -hwaccel_output_format cuda -c:v h264_cuvid -resize 360x240 -i \"https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\" -c:a copy -c:v h264_nvenc -b:v 5M test.mp4"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Check the TorchAudio / FFmpeg integration"
      ],
      "metadata": {
        "id": "YV-gmE6o6pRi"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wDJ-qdM45fV6",
        "outputId": "3cd9990b-a910-4428-ad3b-408bb69d0808"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.0.0.dev20230209+cu117\n",
            "2.0.0.dev20230208+cu117\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torchaudio\n",
        "\n",
        "print(torch.__version__)\n",
        "print(torchaudio.__version__)\n",
        "\n",
        "from torchaudio.io import StreamReader, StreamWriter"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torchaudio.utils import ffmpeg_utils\n",
        "\n",
        "print(\"Library versions:\")\n",
        "print(ffmpeg_utils.get_versions())\n",
        "print(\"\\nBuild config:\")\n",
        "print(ffmpeg_utils.get_build_config())\n",
        "print(\"\\nDecoders:\")\n",
        "print([k for k in ffmpeg_utils.get_video_decoders().keys() if \"cuvid\" in k])\n",
        "print(\"\\nEncoders:\")\n",
        "print([k for k in ffmpeg_utils.get_video_encoders().keys() if \"nvenc\" in k])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6hBVOJuo6yWQ",
        "outputId": "559d6dd9-39f8-4fc2-c00f-1500f277b596"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Library versions:\n",
            "{'libavutil': (56, 31, 100), 'libavcodec': (58, 54, 100), 'libavformat': (58, 29, 100), 'libavfilter': (7, 57, 100), 'libavdevice': (58, 8, 100)}\n",
            "\n",
            "Build config:\n",
            "--prefix=/usr --extra-version=0ubuntu0.1 --toolchain=hardened --libdir=/usr/lib/x86_64-linux-gnu --incdir=/usr/include/x86_64-linux-gnu --arch=amd64 --enable-gpl --disable-stripping --enable-avresample --disable-filter=resample --enable-avisynth --enable-gnutls --enable-ladspa --enable-libaom --enable-libass --enable-libbluray --enable-libbs2b --enable-libcaca --enable-libcdio --enable-libcodec2 --enable-libflite --enable-libfontconfig --enable-libfreetype --enable-libfribidi --enable-libgme --enable-libgsm --enable-libjack --enable-libmp3lame --enable-libmysofa --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-libpulse --enable-librsvg --enable-librubberband --enable-libshine --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libssh --enable-libtheora --enable-libtwolame --enable-libvidstab --enable-libvorbis --enable-libvpx --enable-libwavpack --enable-libwebp --enable-libx265 --enable-libxml2 --enable-libxvid --enable-libzmq --enable-libzvbi --enable-lv2 --enable-omx --enable-openal --enable-opencl --enable-opengl --enable-sdl2 --enable-libdc1394 --enable-libdrm --enable-libiec61883 --enable-nvenc --enable-chromaprint --enable-frei0r --enable-libx264 --enable-shared\n",
            "\n",
            "Decoders:\n",
            "['h264_cuvid', 'hevc_cuvid', 'mjpeg_cuvid', 'mpeg1_cuvid', 'mpeg2_cuvid', 'mpeg4_cuvid', 'vc1_cuvid', 'vp8_cuvid', 'vp9_cuvid']\n",
            "\n",
            "Encoders:\n",
            "['h264_nvenc', 'nvenc', 'nvenc_h264', 'nvenc_hevc', 'hevc_nvenc']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYzJWphMG9nh"
      },
      "source": [
        "## Benchmarking GPU Encoding and Decoding\n",
        "\n",
        "Now that FFmpeg and the resulting libraries are ready to use, we test NVDEC/NVENC with TorchAudio. For the basics of TorchAudio's streaming APIs, please refer to [Media I/O tutorials](https://pytorch.org/audio/main/tutorials.io.html)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "zKMACIJNIXns"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "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)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HSCZFqt95GZm"
      },
      "source": [
        "## Benchmark NVDEC with `StreamReader`\n",
        "\n",
        "First we test hardware decoding, and we fetch video from multiple locations (local file, network file, AWS S3) and use NVDEC to decod them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g3TMksRk5ara",
        "outputId": "98830754-dbe1-4246-9fc6-e4016b382a89"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.26.67\n"
          ]
        }
      ],
      "source": [
        "import boto3\n",
        "from botocore import UNSIGNED\n",
        "from botocore.config import Config\n",
        "\n",
        "print(boto3.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "go6xfv49doJ9"
      },
      "source": [
        "First, we define the functions we'll use for testing.\n",
        "\n",
        "Funcion `test_decode` 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": 11,
      "metadata": {
        "id": "GsmoRAY75T6J"
      },
      "outputs": [],
      "source": [
        "result = torch.zeros((4, 2))\n",
        "samples = [[None, None] for _ in range(4)]\n",
        "\n",
        "\n",
        "def test_decode(src, config, i_sample):\n",
        "  print(\"=\" * 40)\n",
        "  print(\"* Configuration:\", config)\n",
        "  print(\"* Source:\", src)\n",
        "  print(\"=\" * 40)\n",
        "\n",
        "  t0 = time.monotonic()\n",
        "  s = StreamReader(src)\n",
        "  s.add_video_stream(5, **config)\n",
        "\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": 12,
      "metadata": {
        "id": "mlrA9oVM3qvt"
      },
      "outputs": [],
      "source": [
        "local_src = \"input.mp4\"\n",
        "\n",
        "i_sample = 520"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9QScuPXO3_8e"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z6kivDretHtS",
        "outputId": "6ac2d3d4-3526-4946-bb47-a9016b253ba2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: input.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 43.65067997 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cpu_conf = {\n",
        "    \"decoder\": \"h264\",  # CPU decoding\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "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": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V6jwiJD-tW0M",
        "outputId": "62d36f34-969b-42fb-f961-7b2a525a0dd1"
      },
      "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, 540, 960]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 5.754925530000008 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "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",
        "elapsed, sample = test_decode(local_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "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": 17,
      "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_small.mp4\"\n",
        "i_sample = 750"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RsHY7D84YZ3"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6AGfH8X8tboy",
        "outputId": "54a6bfee-cf08-430b-cd7a-b07f3056974a"
      },
      "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_small.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 33.74701378400002 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test_decode(network_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "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": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "95bky0TqtfMI",
        "outputId": "99d1c9c6-6034-4cda-a4a6-47973e7517cf"
      },
      "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_small.mp4\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 15.769149663000007 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "elapsed, sample = test_decode(network_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "kir92rPk4eLr"
      },
      "outputs": [],
      "source": [
        "result[1, 1] = elapsed\n",
        "samples[1][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "owPsmnyTMiQm"
      },
      "source": [
        "### Decode MP4 directly from S3\n",
        "\n",
        "Using file-like object input, we can fetch a video stored on AWS S3 and decode it without saving it on local file system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "XBYc5AiGORUU"
      },
      "outputs": [],
      "source": [
        "bucket = \"pytorch\"\n",
        "key = \"torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4\"\n",
        "\n",
        "s3_client = boto3.client(\"s3\", config=Config(signature_version=UNSIGNED))\n",
        "i_sample = 115"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BG1fHJWFND4C"
      },
      "source": [
        "#### Defining Helper class\n",
        "\n",
        "StreamReader supports file-like objects with `read` method. In addition to this,\n",
        "if the file-like object has `seek` method, StreamReader attempts to use it for more reliable detection of media formats.\n",
        "\n",
        "However, the seek method of `boto3`'s S3 client response object only raises errors to let users know that seek operation is not supported. Therefore we wrap it with a class that does not have `seek` method. This way, StreamReader won't try to use the `seek` method.\n",
        "\n",
        "**Note**\n",
        "\n",
        "Due to the nature of streaming, when using file-like object without seek method, some formats are not supported. For example, MP4 formats contain metadata at the beginning of file or at the end. If metadata is located at the end, without `seek` method, StreamReader cannot decode streams."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "Odq5pMJ4M2ph"
      },
      "outputs": [],
      "source": [
        "# Wrapper to hide the native `seek` method of boto3, which\n",
        "# only raises an error.\n",
        "class UnseekableWrapper:\n",
        "  def __init__(self, obj):\n",
        "    self.obj = obj\n",
        "\n",
        "  def read(self, n):\n",
        "    return self.obj.read(n)\n",
        "\n",
        "  def __str__(self):\n",
        "    return str(self.obj)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ftfQQLIpOBgR"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "rjDH0g0qOAhh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d7c303bb-4a40-4259-f4ef-b967027590e1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264'}\n",
            "* Source: <botocore.response.StreamingBody object at 0x7fcc9b0b83a0>\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cpu torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 29.465254898000012 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "response = s3_client.get_object(Bucket=bucket, Key=key)\n",
        "src = UnseekableWrapper(response[\"Body\"])\n",
        "elapsed, sample = test_decode(src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "JfjvXqBqQGyu"
      },
      "outputs": [],
      "source": [
        "result[2, 0] = elapsed\n",
        "samples[2][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lh3_FqdSODsM"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "17-0MVlZOm-f",
        "outputId": "36891a0f-8fa2-4b99-b27a-4ac93fa19588"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'hw_accel': 'cuda:0'}\n",
            "* Source: <botocore.response.StreamingBody object at 0x7fcc9b0b8730>\n",
            "========================================\n",
            " - Chunk: torch.Size([5, 3, 540, 960]) cuda:0 torch.uint8\n",
            "\n",
            " - Processed 6175 frames.\n",
            " - Elapsed: 3.475249672000018 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "response = s3_client.get_object(Bucket=bucket, Key=key)\n",
        "src = UnseekableWrapper(response[\"Body\"])\n",
        "elapsed, sample = test_decode(src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "VpDLx494QIN4"
      },
      "outputs": [],
      "source": [
        "result[2, 1] = elapsed\n",
        "samples[2][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": 28,
      "metadata": {
        "id": "XKpuH_tXvMSF"
      },
      "outputs": [],
      "source": [
        "i_sample = 1085"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRL3C-bm4ijE"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OEtAVQVYwVCy",
        "outputId": "ddb62460-d669-4afb-86d9-1c33d8c4be08"
      },
      "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: 17.26762169899996 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cpu_conf = {\n",
        "    \"decoder\": \"h264\",  # CPU decoding\n",
        "    \"filter_desc\": \"scale=360:240\",  # Software filter\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cpu_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "uzZMz0rW4j73"
      },
      "outputs": [],
      "source": [
        "result[3, 0] = elapsed\n",
        "samples[3][0] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xkrBNRNP4laK"
      },
      "source": [
        "#### CUDA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "XajhZb-G4mwm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "aafb9dc0-4dcf-4c37-c7f7-840e832f5aeb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'decoder': 'h264_cuvid', 'decoder_option': {'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: 15.233130482000035 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "cuda_conf = {\n",
        "    \"decoder\": \"h264_cuvid\",  # Use CUDA HW decoder\n",
        "    \"decoder_option\": {\n",
        "        \"resize\": \"360x240\",  # Then apply HW preprocessing (resize)\n",
        "    },\n",
        "    \"hw_accel\": \"cuda:0\",  # Then keep the memory on CUDA:0\n",
        "}\n",
        "\n",
        "elapsed, sample = test_decode(local_src, cuda_conf, i_sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "5mXoy38TwYVx"
      },
      "outputs": [],
      "source": [
        "result[3, 1] = elapsed\n",
        "samples[3][1] = sample"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEA2jrrKqmr4"
      },
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "Ni0bihmogZAb",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ad15bfbb-f153-42b5-f207-bfd35f1f51bd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "                                       CPU      NVDEC\n",
            "Decoding (local file)            43.650681   5.754926\n",
            "Decoding (network file)          33.747013  15.769150\n",
            "Decoding (file-like object, S3)  29.465256   3.475250\n",
            "Decoding + Resize                17.267622  15.233130\n"
          ]
        }
      ],
      "source": [
        "res = pd.DataFrame(\n",
        "    result.numpy(),\n",
        "    index=[\"Decoding (local file)\", \"Decoding (network file)\", \"Decoding (file-like object, S3)\", \"Decoding + Resize\"],\n",
        "    columns=[\"CPU\", \"NVDEC\"],\n",
        ")\n",
        "print(res)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i71SJL8PmZfE"
      },
      "source": [
        "The following code shows some frames generated by CPU decoding and NVDEC. They produce seemingly identical results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "nWTT5ih5v1s6"
      },
      "outputs": [],
      "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()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "5AL9u6_xmRQa",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "a717d7b4-9d30-4dd8-879e-d06722e0b4a3"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[]"
            ]
          },
          "metadata": {},
          "execution_count": 35
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 921.6x1382.4 with 8 Axes>"
            ],
            "image/png": "[encoded data, 674568 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "f, axs = plt.subplots(4, 2, figsize=[12.8, 19.2])\n",
        "for i in range(4):\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 == 3 else ''}\")\n",
        "plt.plot(block=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eXzKSVrHk1vS"
      },
      "source": [
        "## Benchmark NVENC with `StreamWriter`\n",
        "\n",
        "Next, we benchmark encoding speed with StreamWriter and NVENC.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "k5WkyQEp0hjv"
      },
      "outputs": [],
      "source": [
        "def test_encode(data, dst, **config):\n",
        "  print(\"=\" * 40)\n",
        "  print(\"* Configuration:\", config)\n",
        "  print(\"* Destination:\", dst)\n",
        "  print(\"=\" * 40)\n",
        "\n",
        "  t0 = time.monotonic()\n",
        "  s = StreamWriter(dst)\n",
        "  s.add_video_stream(**config)\n",
        "  with s.open():\n",
        "    s.write_video_chunk(0, data)\n",
        "  elapsed = time.monotonic() - t0\n",
        "\n",
        "  print()\n",
        "  print(f\" - Processed {len(data)} frames.\")\n",
        "  print(f\" - Elapsed: {elapsed} seconds.\")\n",
        "  print()\n",
        "  return elapsed\n",
        "\n",
        "result = torch.zeros((3, 3))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We use ``StreamReader`` to generate test data."
      ],
      "metadata": {
        "id": "wx3cQ7YwX880"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "nHynkI-oK9mX"
      },
      "outputs": [],
      "source": [
        "def get_data(frame_rate, height, width, format, duration=15):\n",
        "  src = f\"testsrc2=rate={frame_rate}:size={width}x{height}:duration={duration}\"\n",
        "  s = StreamReader(src=src, format=\"lavfi\")\n",
        "  s.add_basic_video_stream(-1, format=format)\n",
        "  s.process_all_packets()\n",
        "  video, = s.pop_chunks()\n",
        "  return video"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1IQafz-HEQtu"
      },
      "source": [
        "### Encode MP4 - 360P\n",
        "\n",
        "For the first test, we compare the time it takes for CPU and NVENC to encode 15 seconds of video with small resolution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "id": "2dZk8Ui7EZew"
      },
      "outputs": [],
      "source": [
        "pict_config = {\n",
        "    \"height\": 360,\n",
        "    \"width\": 640,\n",
        "    \"frame_rate\": 30000/1001,\n",
        "    \"format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "video = get_data(**pict_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZWLHmf8H3Q5"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "id": "wP648bHCH9p_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7186edac-47de-4bf3-aee4-972be5587aab"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 360, 'width': 640, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'libx264', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 360p_cpu.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 6.311792093000008 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"libx264\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[0, 0] = test_encode(video, \"360p_cpu.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P731iWsIH45U"
      },
      "source": [
        "#### CUDA (from CPU Tensor)\n",
        "\n",
        "Now we test NVENC. This time, the data is sent from CPU memory to GPU memory as part of encoding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "zQ9W_UVTIf-E",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "535e49f3-ba72-4ef9-aa53-006bccee4f0d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 360, 'width': 640, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}}\n",
            "* Destination: 360p_cuda.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.5352325430000064 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",  # Use NVENC\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},  # Run encoding on the cuda:0 device\n",
        "}\n",
        "\n",
        "result[1, 0] = test_encode(video, \"360p_cuda.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QirJIHuSm-WT"
      },
      "source": [
        "#### CUDA (from CUDA Tensor)\n",
        "\n",
        "If the data is already present on CUDA, then we can pass it to GPU encoder directly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "EHjyMYpJnEb3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2138a71c-7e9f-4571-d076-4ef44a1ea02c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 360, 'width': 640, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}, 'hw_accel': 'cuda:0'}\n",
            "* Destination: 360p_cuda_hw.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.7045466739999711 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "device = \"cuda:0\"\n",
        "\n",
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",  # GPU Encoder\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},  # Run encoding on the cuda:0 device\n",
        "    \"hw_accel\": device,  # Data comes from cuda:0 device\n",
        "}\n",
        "\n",
        "result[2, 0] = test_encode(video.to(torch.device(device)), \"360p_cuda_hw.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LBCBSKkRJeiA"
      },
      "source": [
        "### Encode MP4 - 720P\n",
        "\n",
        "Let's run the same tests on video with larger resolution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "CUdzN9BzJlo_"
      },
      "outputs": [],
      "source": [
        "pict_config = {\n",
        "    \"height\": 720,\n",
        "    \"width\": 1280,\n",
        "    \"frame_rate\": 30000/1001,\n",
        "    \"format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "video = get_data(**pict_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wTjAcXGAJsE8"
      },
      "source": [
        "#### CPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FoEM76wFJta5",
        "outputId": "25c88435-06d8-4383-ccfa-a747dc246e46"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 720, 'width': 1280, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'libx264', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 720p_cpu.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 14.017578084000036 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"libx264\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[0, 1] = test_encode(video, \"720p_cpu.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lgbIN67BKARA"
      },
      "source": [
        "#### CUDA (from CPU Tensor)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fgNR79olJ7Yo",
        "outputId": "9bbbacb6-a3d5-485a-9194-26efc9fe6f1e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 720, 'width': 1280, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 720p_cuda.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.9797491379999883 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[1, 1] = test_encode(video, \"720p_cuda.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oD_XUn4-KDod"
      },
      "source": [
        "#### CUDA (from CUDA Tensor)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2NX4SU5xJ7zc",
        "outputId": "c81cc029-e6e0-4552-d3a0-cb5ba20935c8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 720, 'width': 1280, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}, 'hw_accel': 'cuda:0'}\n",
            "* Destination: 720p_cuda_hw.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 0.7259890020000057 seconds.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "device = \"cuda:0\"\n",
        "\n",
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},\n",
        "    \"hw_accel\": device,\n",
        "}\n",
        "\n",
        "result[2, 1] = test_encode(video.to(torch.device(device)), \"720p_cuda_hw.mp4\", **pict_config, **encode_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Encode MP4 - 1080P\n",
        "\n",
        "We make the video with even larger."
      ],
      "metadata": {
        "id": "FoMDEhtYO-xD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pict_config = {\n",
        "    \"height\": 1080,\n",
        "    \"width\": 1920,\n",
        "    \"frame_rate\": 30000/1001,\n",
        "    \"format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "video = get_data(**pict_config)"
      ],
      "metadata": {
        "id": "klcLfXipO5OX"
      },
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CPU"
      ],
      "metadata": {
        "id": "o28eOGRuPD3n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"libx264\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[0, 2] = test_encode(video, \"1080p_cpu.mp4\", **pict_config, **encode_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6UuiAEowOedt",
        "outputId": "06584da3-d1d4-4cd4-f96f-ea002a2361b6"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 1080, 'width': 1920, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'libx264', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 1080p_cpu.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 29.441768007999997 seconds.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CUDA (from CPU Tensor)"
      ],
      "metadata": {
        "id": "ACJ3oEwgPFAh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "}\n",
        "\n",
        "result[1, 2] = test_encode(video, \"1080p_cuda.mp4\", **pict_config, **encode_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dzYNn8HtOeVi",
        "outputId": "47a24824-e3a3-4426-9b6d-414438c7a8d4"
      },
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 1080, 'width': 1920, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p'}\n",
            "* Destination: 1080p_cuda.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 1.5768605540000067 seconds.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CUDA (from CUDA Tensor)"
      ],
      "metadata": {
        "id": "YcEH15BYPOQ_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = \"cuda:0\"\n",
        "\n",
        "encode_config = {\n",
        "    \"encoder\": \"h264_nvenc\",\n",
        "    \"encoder_format\": \"yuv444p\",\n",
        "    \"encoder_option\": {\"gpu\": \"0\"},\n",
        "    \"hw_accel\": device,\n",
        "}\n",
        "\n",
        "result[2, 2] = test_encode(video.to(torch.device(device)), \"1080p_cuda_hw.mp4\", **pict_config, **encode_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cUZOWVllOeRl",
        "outputId": "a0cbae4c-376d-4a8e-f12b-625432ae634f"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "========================================\n",
            "* Configuration: {'height': 1080, 'width': 1920, 'frame_rate': 29.97002997002997, 'format': 'yuv444p', 'encoder': 'h264_nvenc', 'encoder_format': 'yuv444p', 'encoder_option': {'gpu': '0'}, 'hw_accel': 'cuda:0'}\n",
            "* Destination: 1080p_cuda_hw.mp4\n",
            "========================================\n",
            "\n",
            " - Processed 450 frames.\n",
            " - Elapsed: 1.4002963130000126 seconds.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Result\n",
        "\n",
        "Here is the result."
      ],
      "metadata": {
        "id": "Z9lEeklfPY-Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "labels = [\"CPU\", \"CUDA (from CPU Tensor)\", \"CUDA (from CUDA Tensor)\"]\n",
        "columns = [\"360P\", \"720P\", \"1080P\"]\n",
        "res = pd.DataFrame(\n",
        "    result.numpy(),\n",
        "    index=labels,\n",
        "    columns=columns,\n",
        ")\n",
        "print(res)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jDWUQFjVPaUV",
        "outputId": "00409409-bb44-4d54-c381-589c631a6c12"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "                             360P       720P      1080P\n",
            "CPU                      6.311792  14.017578  29.441769\n",
            "CUDA (from CPU Tensor)   0.535233   0.979749   1.576861\n",
            "CUDA (from CUDA Tensor)  0.704547   0.725989   1.400296\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "plt.plot(result.T)\n",
        "plt.legend(labels)\n",
        "plt.xticks([i for i in range(3)], columns)\n",
        "plt.grid(visible=True, axis='y')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "gEccR5zwTYRQ",
        "outputId": "5f0f71e7-6336-416d-9b53-5653d05ef23f"
      },
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "[encoded data, 15960 characters]\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The resulting videos look like the following."
      ],
      "metadata": {
        "id": "mVeuM2gmkKld"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML\n",
        "\n",
        "HTML('''\n",
        "<div>\n",
        "  <video width=360 controls autoplay>\n",
        "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cpu.mp4\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  <video width=360 controls autoplay>\n",
        "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda.mp4\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  <video width=360 controls autoplay>\n",
        "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda_hw.mp4\" type=\"video/mp4\">\n",
        "  </video>\n",
        "</div>\n",
        "''')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 224
        },
        "id": "LNPWCXxOgnqz",
        "outputId": "72aec5f1-3c46-4294-f330-fb0f92e65e94"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "<div>\n",
              "  <video width=360 controls autoplay>\n",
              "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cpu.mp4\" type=\"video/mp4\">\n",
              "  </video>\n",
              "  <video width=360 controls autoplay>\n",
              "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda.mp4\" type=\"video/mp4\">\n",
              "  </video>\n",
              "  <video width=360 controls autoplay>\n",
              "    <source src=\"https://download.pytorch.org/torchaudio/tutorial-assets/streamwriter_360p_cuda_hw.mp4\" type=\"video/mp4\">\n",
              "  </video>\n",
              "</div>\n"
            ]
          },
          "metadata": {},
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yuYzPpV9zc4M"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "We looked at how to build FFmpeg libraries with NVDEC/NVENC support and use them from TorchAudio. NVDEC/NVENC provide significant speed up when saving/loading a video."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
