onnx_to_tf_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow.keras as keras\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keras.backend.set_image_data_format('channels_first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def block(last, i):\n",
    "    x = keras.layers.Conv2D(256, (3, 3), padding=\"same\", name=\"block{}/conv{}\".format(i, 1), activation=\"relu\")(last)\n",
    "    x = keras.layers.Conv2D(256, (3, 3), padding=\"same\", name=\"block{}/conv{}\".format(i, 2), activation=\"linear\")(x)\n",
    "    add = keras.layers.Add()([x, last])\n",
    "    rel = keras.layers.ReLU()(add)\n",
    "    return rel, add\n",
    "\n",
    "\n",
    "inputs = keras.Input((112, 8, 8))\n",
    "last = keras.layers.Conv2D(256, (3, 3), padding=\"same\", name=\"inputconv\")(inputs)\n",
    "last = keras.layers.ReLU()(last)\n",
    "inter_outputs = [last]\n",
    "for b in range(0, 20):\n",
    "    last, inter = block(last, b)\n",
    "    inter_outputs.append(inter)\n",
    "\n",
    "policy_conv = keras.layers.Conv2D(32, (1, 1), activation=\"relu\", padding=\"same\", name=\"policy_conv\")(last)\n",
    "flatten = keras.layers.Flatten(data_format=\"channels_last\")(policy_conv)\n",
    "policy_output = keras.layers.Dense(1858, activation=\"linear\", name=\"policy_dense\")(flatten)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = keras.Model(inputs, policy_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "\n",
    "from onnx import numpy_helper\n",
    "_model = onnx.load(\"32390.onnx\")\n",
    "INTIALIZERS = _model.graph.initializer\n",
    "weights = {}\n",
    "for initializer in INTIALIZERS:\n",
    "    W = numpy_helper.to_array(initializer)\n",
    "    if len(W.shape) > 2:\n",
    "        weights[initializer.name[1:]] = np.moveaxis(W, [0, 1], [-1, -2])\n",
    "    else:\n",
    "        weights[initializer.name[1:]] = W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights[\"policy/dense/matmul/w\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights[\"output/policy/w\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.moveaxis(weights[\"inputconv/w/kernel\"], [0, 1], [-1, -2]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights[\"policy/conv/w/kernel\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for layer in model.layers:\n",
    "    if \"policy_conv\" in layer.name:\n",
    "        layer.set_weights([weights[\"policy/conv/w/kernel\"], weights[\"policy/conv/w/bias\"]])\n",
    "    elif \"policy_dense\" in layer.name:\n",
    "        layer.set_weights([weights[\"policy/dense/matmul/w\"], weights[\"output/policy/w\"]])\n",
    "    elif \"conv\" in layer.name:\n",
    "        layer.set_weights([weights[layer.name + \"/w/kernel\"], weights[layer.name + \"/w/bias\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verifying that it gives the same outputs as the onnx model...\n",
    "import numpy as np\n",
    "np.random.seed(123)\n",
    "random_inputs = np.random.random((4, 112, 8, 8)).astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_outputs = np.array(model(random_inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime as rt\n",
    "\n",
    "sess = rt.InferenceSession(\"32390.onnx\", providers=['CUDAExecutionProvider'])\n",
    "input_name = sess.get_inputs()[0].name\n",
    "output_names = [[o.name for o in sess.get_outputs()][0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_name = sess.get_inputs()[0].name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = np.squeeze(np.array(sess.run(output_names, {input_name: random_inputs})))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_outputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.allclose(tf_outputs, pred, atol=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(\"policy_model\")"
   ]
  }
 ],
 "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.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}