torchscript.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "from networks.models import build_model"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "import torch\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "model = build_model(\"FBA.pth\").cuda().eval()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "modifying input layer to accept 11 channels\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "image = torch.randn((1, 3, 128, 128)).cuda()\n",
    "two_chan_trimap=  torch.randn((1, 2, 128, 128)).cuda()\n",
    "image_n=torch.randn((1, 3, 128, 128)).cuda()\n",
    "trimap_transformed =torch.randn((1, 6, 128, 128)).cuda()\n",
    "output = model(image, two_chan_trimap, image_n, trimap_transformed)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/home/marco/anaconda3/lib/python3.8/site-packages/torch/nn/functional.py:693: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /opt/conda/conda-bld/pytorch_1623448278899/work/c10/core/TensorImpl.h:1156.)\n",
      "  return torch._C._nn.max_pool2d_with_indices(input, kernel_size, stride, padding, dilation, ceil_mode)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "output = model(image, two_chan_trimap, image_n, trimap_transformed)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "jitmodel = torch.jit.trace(model, [image, two_chan_trimap, image_n, trimap_transformed])\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/home/marco/anaconda3/lib/python3.8/site-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.\n",
      "To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /opt/conda/conda-bld/pytorch_1623448278899/work/aten/src/ATen/native/BinaryOps.cpp:467.)\n",
      "  return torch.floor_divide(self, other)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "output_jit = jitmodel(image, two_chan_trimap, image_n, trimap_transformed)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "print(torch.mean( (output - output_jit)**2).item())"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "0.0\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "%timeit _ = jitmodel(image, two_chan_trimap, image_n, trimap_transformed); torch.cuda.synchronize()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "5.73 ms ± 96.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "%timeit _ = model(image, two_chan_trimap, image_n, trimap_transformed); torch.cuda.synchronize()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "7.72 ms ± 102 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.8.5",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.5 64-bit ('base': conda)"
  },
  "interpreter": {
   "hash": "c81b9cd8e0a9008ef2c2b50032556bb47b7837e23c42e37fcb1c71e2e49ce391"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}