GPU.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPU Tensors.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vM54r6jlKTII",
        "colab_type": "text"
      },
      "source": [
        "# Install detectron2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9_FzH13EjseR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 86
        },
        "outputId": "6c3946e9-d4ae-4a37-f767-0a2e5ddc4313"
      },
      "source": [
        "# install dependencies: (use cu100 because colab is on CUDA 10.0)\n",
        "!pip install -U torch==1.4+cu100 -f https://download.pytorch.org/whl/torch_stable.html \n",
        "import torch\n",
        "print(torch.__version__)\n",
        "print(\"cuda:\", torch.cuda.is_available())"
      ],
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
            "Requirement already up-to-date: torch==1.4+cu100 in /usr/local/lib/python3.6/dist-packages (1.4.0+cu100)\n",
            "1.4.0+cu100\n",
            "cuda: True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4RFplrOoMcKD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create a model and assign some simple weights.\n",
        "fc = torch.nn.Linear(2, 2, bias=False)\n",
        "with torch.no_grad():\n",
        "  fc.eval()\n",
        "  fc.weight.copy_(torch.tensor([[1.0,2.0],[3.0,4.0]]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7xQiH94yM0IA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Send the model to the GPU.\n",
        "fc_gpu = fc.to(\"cuda\")\n",
        "fc_gpu.eval()\n",
        "# I don't really understand why this next line is necessary,\n",
        "# but tracing fails without it.\n",
        "fc_gpu.weight.requires_grad = False"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7d3KxiHO_0gb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Trace it directly.\n",
        "inp = torch.tensor([11.0,111.0]).to(\"cuda\")\n",
        "with torch.no_grad():\n",
        "  traced_model = torch.jit.trace(fc_gpu, inp)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "36utZ79WcgWx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# This doesn't work because the input is not CUDA.\n",
        "# traced_model(torch.tensor([10.0,100.0]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XAjElWsDc10g",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create a wrapper than handles the inter-device copies.\n",
        "def wrapper(inp):\n",
        "  return fc_gpu(inp.to(\"cuda\").detach()).to(\"cpu\")\n",
        "\n",
        "# Trace that wrapper.\n",
        "with torch.no_grad():\n",
        "  traced_for_cpu = torch.jit.trace(wrapper, torch.tensor([10.0,100.0]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZP9mthcrdQoS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Save the model.\n",
        "traced_for_cpu.save(\"model.pt1\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B7FPFKAzeVhP",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "87147a5f-ea84-4b70-b0e6-53fee0af64aa"
      },
      "source": [
        "# This part should work fine in Java, and I believe it is running on GPU.\n",
        "loaded = torch.jit.load(\"model.pt1\")\n",
        "loaded(torch.tensor([10.0,100.0]))"
      ],
      "execution_count": 79,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([210., 430.])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 79
        }
      ]
    }
  ]
}

It exactly not work without setting all the parameters as requires_grad=False. But, please how to set this?

in my test, it doesn’t running on GPU in java.