RegNet 模型使用/导出ONNX - Colab 版本

regnet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "RegNet.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "xPRFUnd11L3W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime → \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b8pCEaIn1OWQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yya2ycdj1Qy-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "path = '/content/drive/My Drive/Colab Notebooks/RegNet'\n",
        "if not os.path.exists(path):\n",
        "    os.makedirs(path)\n",
        "os.chdir(path)\n",
        "print(\"Current Working Directory \" , os.getcwd())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kiSRmLND1RR-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install -r requirements.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mf2fs1lu3Y_M",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "os.chdir(path)\n",
        "print(\"Current Working Directory \" , os.getcwd())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rbaiao043d3u",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from pycls.models.regnet import RegNet\n",
        "from pycls.core.config import cfg\n",
        "import torch\n",
        "\n",
        "# modify regnet to equal efficientnet-b5\n",
        "cfg.REGNET.SE_ON = True\n",
        "cfg.REGNET.DEPTH = 17\n",
        "cfg.REGNET.W0 = 192\n",
        "cfg.REGNET.WA = 76.82\n",
        "cfg.REGNET.WM = 2.19\n",
        "cfg.REGNET.GROUP_W = 56\n",
        "# fit our dataset\n",
        "cfg.MODEL.NUM_CLASSES = 8\n",
        "\n",
        "model = RegNet()\n",
        "print(model)\n",
        "\n",
        "# for onnx format export\n",
        "model.eval()\n",
        "dummy_input = torch.randn(1, 3, 456, 456)\n",
        "torch.onnx.export(model, dummy_input, \"regnet.onnx\", input_names=['input'], output_names=['output'], verbose=True)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}

requirements.txt

numpy
opencv-python
setuptools
simplejson
yacs