Convert2ONNX.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "V100"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!wget https://github.com/CAIC-AD/YOLOPv2/releases/download/V0.0.1/yolopv2.pt -O yolopv2.pt"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m4EQeWDUDAP5",
        "outputId": "4ce6f238-8865-4ce0-b164-740d6ce0c44c"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-07-07 02:03:12--  https://github.com/CAIC-AD/YOLOPv2/releases/download/V0.0.1/yolopv2.pt\n",
            "Resolving github.com (github.com)... 140.82.121.3\n",
            "Connecting to github.com (github.com)|140.82.121.3|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/527915303/7bbddd11-355a-4a1d-b175-6a62c72dbad9?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230707%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230707T020312Z&X-Amz-Expires=300&X-Amz-Signature=b2d40342d4a9145d2d67c9a2b807e94955d914e4409e7ab63ee95deb75c0ab8a&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=527915303&response-content-disposition=attachment%3B%20filename%3Dyolopv2.pt&response-content-type=application%2Foctet-stream [following]\n",
            "--2023-07-07 02:03:12--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/527915303/7bbddd11-355a-4a1d-b175-6a62c72dbad9?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230707%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230707T020312Z&X-Amz-Expires=300&X-Amz-Signature=b2d40342d4a9145d2d67c9a2b807e94955d914e4409e7ab63ee95deb75c0ab8a&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=527915303&response-content-disposition=attachment%3B%20filename%3Dyolopv2.pt&response-content-type=application%2Foctet-stream\n",
            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.110.133, 185.199.111.133, ...\n",
            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 156380200 (149M) [application/octet-stream]\n",
            "Saving to: ‘yolopv2.pt’\n",
            "\n",
            "yolopv2.pt          100%[===================>] 149.14M  45.8MB/s    in 3.6s    \n",
            "\n",
            "2023-07-07 02:03:16 (40.9 MB/s) - ‘yolopv2.pt’ saved [156380200/156380200]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install onnx onnxruntime"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8tJhG3a7DB24",
        "outputId": "a269f670-6a50-469b-8d96-80e803c7fbaf"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting onnx\n",
            "  Downloading onnx-1.14.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (14.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m14.6/14.6 MB\u001b[0m \u001b[31m32.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting onnxruntime\n",
            "  Downloading onnxruntime-1.15.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.9 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.9/5.9 MB\u001b[0m \u001b[31m78.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from onnx) (1.22.4)\n",
            "Requirement already satisfied: protobuf>=3.20.2 in /usr/local/lib/python3.10/dist-packages (from onnx) (3.20.3)\n",
            "Requirement already satisfied: typing-extensions>=3.6.2.1 in /usr/local/lib/python3.10/dist-packages (from onnx) (4.6.3)\n",
            "Collecting coloredlogs (from onnxruntime)\n",
            "  Downloading coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: flatbuffers in /usr/local/lib/python3.10/dist-packages (from onnxruntime) (23.5.26)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from onnxruntime) (23.1)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from onnxruntime) (1.11.1)\n",
            "Collecting humanfriendly>=9.1 (from coloredlogs->onnxruntime)\n",
            "  Downloading humanfriendly-10.0-py2.py3-none-any.whl (86 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m10.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->onnxruntime) (1.3.0)\n",
            "Installing collected packages: onnx, humanfriendly, coloredlogs, onnxruntime\n",
            "Successfully installed coloredlogs-15.0.1 humanfriendly-10.0 onnx-1.14.0 onnxruntime-1.15.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "device = 'cpu'\n",
        "model_path = 'yolopv2.pt'\n",
        "\n",
        "model = torch.load(model_path)\n",
        "model = model.to(device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MjpTot2gDGkJ",
        "outputId": "ac891e83-5b9e-4c07-964b-232da414b6b2"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/serialization.py:799: UserWarning: 'torch.load' received a zip file that looks like a TorchScript archive dispatching to 'torch.jit.load' (call 'torch.jit.load' directly to silence this warning)\n",
            "  warnings.warn(\"'torch.load' received a zip file that looks like a TorchScript archive\"\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def convert_to_onnx(net, file_name='output.onnx', input_shape=(640, 640), device='cpu', verbose=False):\n",
        "    input = torch.randn(1, 3, input_shape[1], input_shape[0]).to(device)\n",
        "\n",
        "    torch.onnx.export(\n",
        "        net,\n",
        "        input,\n",
        "        file_name,\n",
        "        input_names=['input_image'],\n",
        "        dynamic_axes={\n",
        "            'input_image': {2: 'height', 3:'width'}\n",
        "        },\n",
        "        verbose=verbose,\n",
        "    )"
      ],
      "metadata": {
        "id": "JRJ82RtXDLGW"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "convert_to_onnx(model, 'YOLOPv2.onnx')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NwT0KETPDmwu",
        "outputId": "09d7148f-b41e-4547-8100-31a9c2b3a09a"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py:825: UserWarning: no signature found for <torch.ScriptMethod object at 0x7fa0fae3efc0>, skipping _decide_input_format\n",
            "  warnings.warn(f\"{e}, skipping _decide_input_format\")\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============= Diagnostic Run torch.onnx.export version 2.0.1+cu118 =============\n",
            "verbose: False, log level: Level.ERROR\n",
            "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
            "\n"
          ]
        }
      ]
    }
  ]
}