pytorch-tf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import numpy as np \n",
    "import tensorflow as tf\n",
    "import torch\n",
    "\n",
    "import torchvision.models as models\n",
    "import torchvision\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = models.squeezenet1_1(pretrained=True)\n",
    "destination_py = 'squeezenet.py'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type_lookups = {}\n",
    "outfp = open(destination_py,'w')\n",
    "outfp.write('import tensorflow as tf\\n\\n')\n",
    "out_s = ''\n",
    "def conv2d(c,**kwargs):\n",
    "    padding = 'VALID' if c.padding[0] is 0 else 'SAME'\n",
    "    filters = c.out_channels\n",
    "    size = c.kernel_size\n",
    "    parameters = [p for p in c.parameters()]\n",
    "    W = parameters[0].data.numpy()\n",
    "    if len(parameters) > 1:\n",
    "        b = parameters[1].data.numpy()\n",
    "\n",
    "    W = np.transpose(W,[2,3,1,0])\n",
    "\n",
    "    wi = tf.constant_initializer(W)\n",
    "    if len(parameters) > 1:\n",
    "        bi = tf.constant_initializer(b)\n",
    "    Wt = tf.get_variable('weights',shape=W.shape,initializer=wi)#,\n",
    "    if 'print' not in kwargs or kwargs['print'] == True:\n",
    "        outfp.write(out_s + 'W = tf.get_variable(\"weights\",shape=[{},{},{},{}])\\n'.format(*list(W.shape)))\n",
    "\n",
    "    if len(parameters) > 1:\n",
    "        bt = tf.get_variable('bias',shape=b.shape,initializer=bi)#,\n",
    "        if 'print' not in kwargs or kwargs['print'] == True:\n",
    "            outfp.write(out_s + 'b = tf.get_variable(\"bias\",shape=[{}])\\n'.format(b.shape[0]))\n",
    "    x = tf.nn.conv2d(kwargs['inp'],Wt,[1,c.stride[0],c.stride[1],1],padding)\n",
    "    if 'print' not in kwargs or kwargs['print'] == True:\n",
    "        outfp.write(out_s + 'x = tf.nn.conv2d(x,W,[1,{},{},1],\"{}\")\\n'.format(c.stride[0],c.stride[1],padding))\n",
    "    if len(parameters) > 1:\n",
    "        x = tf.nn.bias_add(x,bt)\n",
    "        if 'print' not in kwargs or kwargs['print'] == True:\n",
    "            outfp.write(out_s + 'x = tf.nn.bias_add(x,b)\\n')\n",
    "\n",
    "    return x\n",
    "\n",
    "def relu(c,**kwargs):\n",
    "    outfp.write(out_s + \"x = tf.nn.relu(x)\\n\")\n",
    "    return tf.nn.relu(kwargs['inp'])\n",
    "def max_pool(c,**kwargs):\n",
    "    padding = 'VALID' if c.padding is 0 else 'SAME'\n",
    "    outfp.write(out_s + \"x = tf.nn.max_pool(x,[1,{0},{0},1],strides=[1,{1},{1},1],padding='{2}')\\n\".format(\n",
    "        c.kernel_size,c.stride,padding))\n",
    "    x = tf.nn.max_pool(kwargs['inp'],[1,c.kernel_size,c.kernel_size,1],strides=[1,c.stride,c.stride,1],padding=padding)\n",
    "    return x\n",
    "def avg_pool(c,**kwargs):\n",
    "    padding = 'VALID' if c.padding is 0 else 'SAME'\n",
    "    outfp.write(out_s + \"x = tf.nn.avg_pool(x,[1,{0},{0},1],strides=[1,{1},{1},1],padding='{2}')\\n\".format(\n",
    "        c.kernel_size,c.stride,padding))\n",
    "    x = tf.nn.avg_pool(kwargs['inp'],[1,c.kernel_size,c.kernel_size,1],strides=[1,c.stride,c.stride,1],padding=padding)\n",
    "    return x\n",
    "def dropout(c,**kwargs):\n",
    "    outfp.write(out_s + 'x = x\\n')\n",
    "    return kwargs['inp']\n",
    "def fire_module(c,**kwargs):\n",
    "    global out_s\n",
    "\n",
    "    # couldn't figure out how to\n",
    "    # automatically unravel it\n",
    "    outfp.write(out_s + \"x = fire_module(x,{0},{1},{2},{3})\\n\".format(\n",
    "        c.squeeze.in_channels,c.squeeze.out_channels,c.expand1x1.out_channels,c.expand3x3.out_channels))\n",
    "    with tf.variable_scope(\"fire\"):\n",
    "        with tf.variable_scope(\"squeeze\"):\n",
    "            s = conv2d(c.squeeze,inp=kwargs['inp'],print=False)\n",
    "            s = tf.nn.relu(s)\n",
    "        with tf.variable_scope(\"e11\"):\n",
    "            e11 = conv2d(c.expand1x1,inp=s,print=False)\n",
    "            e11 = tf.nn.relu(e11)\n",
    "        with tf.variable_scope(\"e33\"):\n",
    "            e33 = conv2d(c.expand3x3,inp=s,print=False)\n",
    "            e33 = tf.nn.relu(e33)\n",
    "    x = tf.concat([e11,e33],3)\n",
    "    return x\n",
    "\n",
    "def seq_container(c,**kwargs):\n",
    "    global out_s\n",
    "    x = kwargs['inp']\n",
    "    for c2 in enumerate(c.children()):\n",
    "        c2_class = c2[1].__class__\n",
    "        if c2_class in type_lookups:\n",
    "            outfp.write(out_s + \"with tf.variable_scope('{}'):\\n\".format('layer' + str(c2[0])))\n",
    "            with tf.variable_scope('layer' + str(c2[0])):\n",
    "                out_s = out_s + '    '\n",
    "                x = type_lookups[c2_class](c2[1],inp = x)\n",
    "                name = kwargs['name'] if 'name' in kwargs else ''\n",
    "                outfp.write(out_s + \"self.layers.append(x)\\n\".format(name + str(c2[0])))\n",
    "\n",
    "                out_s = out_s[:-4]\n",
    "        else:\n",
    "            unknown_class(c2[1])\n",
    "            print(c2_class)\n",
    "    return x\n",
    "def batch_norm(c,**kwargs):\n",
    "    print('batch_norm')\n",
    "    return kwargs['inp']\n",
    "type_lookups[torch.nn.modules.conv.Conv2d] = conv2d\n",
    "type_lookups[torch.nn.modules.activation.ReLU] = relu\n",
    "type_lookups[torch.nn.modules.container.Sequential] = seq_container\n",
    "type_lookups[torch.nn.modules.pooling.MaxPool2d] = max_pool\n",
    "type_lookups[torch.nn.modules.pooling.AvgPool2d] = avg_pool\n",
    "type_lookups[torch.nn.modules.dropout.Dropout] = dropout\n",
    "type_lookups[torchvision.models.squeezenet.Fire] = fire_module\n",
    "type_lookups[torch.nn.modules.batchnorm.BatchNorm2d] = batch_norm\n",
    "tf.reset_default_graph()\n",
    "input_image = tf.placeholder('float',shape=[None,None,None,3],name='input_image')\n",
    "\n",
    "if True:\n",
    "    outfp.write('def fire_module(x,inp,sp,e11p,e33p):\\n')\n",
    "    outfp.write('    with tf.variable_scope(\"fire\"):\\n')\n",
    "    outfp.write('        with tf.variable_scope(\"squeeze\"):\\n')\n",
    "    outfp.write('            W = tf.get_variable(\"weights\",shape=[1,1,inp,sp])\\n')\n",
    "    outfp.write('            b = tf.get_variable(\"bias\",shape=[sp])\\n')\n",
    "    outfp.write('            s = tf.nn.conv2d(x,W,[1,1,1,1],\"VALID\")+b\\n')\n",
    "    outfp.write('            s = tf.nn.relu(s)\\n')\n",
    "    outfp.write('        with tf.variable_scope(\"e11\"):\\n')\n",
    "    outfp.write('            W = tf.get_variable(\"weights\",shape=[1,1,sp,e11p])\\n')\n",
    "    outfp.write('            b = tf.get_variable(\"bias\",shape=[e11p])\\n')\n",
    "    outfp.write('            e11 = tf.nn.conv2d(s,W,[1,1,1,1],\"VALID\")+b\\n')\n",
    "    outfp.write('            e11 = tf.nn.relu(e11)\\n')\n",
    "    outfp.write('        with tf.variable_scope(\"e33\"):\\n')\n",
    "    outfp.write('            W = tf.get_variable(\"weights\",shape=[3,3,sp,e33p])\\n')\n",
    "    outfp.write('            b = tf.get_variable(\"bias\",shape=[e33p])\\n')\n",
    "    outfp.write('            e33 = tf.nn.conv2d(s,W,[1,1,1,1],\"SAME\")+b\\n')\n",
    "    outfp.write('            e33 = tf.nn.relu(e33)\\n')\n",
    "    outfp.write('        return tf.concat([e11,e33],3) \\n\\n')\n",
    "\n",
    "\n",
    "if len([_ for _ in model.children()]) == 2:\n",
    "    outfp.write('class SqueezeNet:\\n')\n",
    "    out_s += '    '\n",
    "    outfp.write(out_s + 'def __init__(self):\\n')\n",
    "    \n",
    "    for idx,c in enumerate(model.children()):\n",
    "        out_s = out_s + '    '\n",
    "\n",
    "        if idx is 0:\n",
    "            outfp.write(out_s+\"self.image = tf.placeholder('float',shape=[None,None,None,3],name='input_image')\\n\")\n",
    "            outfp.write(out_s+\"self.layers = []\\n\")\n",
    "\n",
    "            outfp.write(out_s+'x = self.image\\n')\n",
    "            outfp.write(out_s+\"with tf.variable_scope('features'):\\n\")\n",
    "            with tf.variable_scope('features'):\n",
    "                out_s = out_s + '    '\n",
    "                features = type_lookups[c.__class__](c,inp=input_image)\n",
    "                out_s = out_s[:-4]\n",
    "\n",
    "            outfp.write(out_s+'self.features = x\\n')\n",
    "\n",
    "        elif idx is 1:\n",
    "            outfp.write(out_s+\"with tf.variable_scope('classifier'):\\n\")\n",
    "            with tf.variable_scope('classifier'):\n",
    "                out_s = out_s + '    '\n",
    "                classifier = type_lookups[c.__class__](c,inp=features)\n",
    "                classifier = tf.reshape(classifier,[-1,1000])\n",
    "                out_s = out_s[:-4]\n",
    "\n",
    "            outfp.write(out_s+'self.classifier = tf.reshape(x,[-1,1000])\\n')\n",
    "            outfp.write('\\n\\n')\n",
    "        out_s = out_s[:-4]\n",
    "\n",
    "\n",
    "else:\n",
    "    x = input_image\n",
    "    for idx,c in enumerate(model.children()):\n",
    "        x = type_lookups[c.__class__](c,inp=x)\n",
    "outfp.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier.get_shape(),classifier.name,input_image.name,features.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from scipy.misc import imresize\n",
    "import os\n",
    "\n",
    "with open('labels.txt') as fp:\n",
    "    labels = [c[:-2].split(':')[1] for c in fp.readlines()]\n",
    "def get_img(filename):\n",
    "    vec = np.array(Image.open(filename))\n",
    "    vec = imresize(vec,(224,224)).astype(np.float32)/255.0\n",
    "    mean = np.array([0.485, 0.456, 0.406])\n",
    "    std = np.array([0.229, 0.224, 0.225])\n",
    "    vec = (vec-mean)/std\n",
    "    return vec\n",
    "    \n",
    "img_dir = '.'\n",
    "img_names = [x for x in os.listdir(img_dir) if 'jpeg' in x.lower()]\n",
    "imgs = [get_img(os.path.join(img_dir,x)) for x in img_names]\n",
    "\n",
    "saver = tf.train.Saver()\n",
    "\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "scores = sess.run(classifier,feed_dict={input_image:np.array(imgs).reshape([-1,224,224,3])})\n",
    "for idx,s in enumerate(np.argmax(scores,1)):\n",
    "    print(img_names[idx],labels[s])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saver.save(sess, 'squeezenet.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from torch.autograd import Variable\n",
    "input_data = torch.FloatTensor(np.transpose(np.array(imgs),[0,3,1,2]))\n",
    "model.eval()\n",
    "pyt_scores = model(Variable(input_data))\n",
    "scores_ref = pyt_scores.data.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rel_error(x, y):\n",
    "    return np.max(np.abs(x - y) / (np.maximum(1e-8, np.abs(x) + np.abs(y))))\n",
    "print(rel_error(scores,scores_ref))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}