pyro_ENUMERATION_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorail on Enumeration of Cartesian Products in Pyro"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import the necessary modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import pyro\n",
    "import pyro.poutine as poutine\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer import SVI, Trace_ELBO, TraceEnum_ELBO, config_enumerate\n",
    "\n",
    "smoke_test = ('CI' in os.environ)\n",
    "assert pyro.__version__.startswith('0.3.0')\n",
    "pyro.enable_validation(True) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Normal approach (i.e. no enumeration)\n",
    "\n",
    "Note that:\n",
    "1. in Trace_ELBO we draw \"num_particles\" samples.\n",
    "2. the configuration which is drawn has the shape \"BATCH x N_BOXES\", i.e. both batch and dependent dimensions are expanded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config_normal: torch.Size([3, 4]) tensor([[0., 1., 0., 1.],\n",
      "        [0., 1., 1., 0.],\n",
      "        [0., 1., 1., 0.]])\n",
      "config_normal: torch.Size([3, 4]) tensor([[0., 0., 1., 0.],\n",
      "        [0., 1., 1., 1.],\n",
      "        [1., 0., 1., 1.]])\n",
      "------------------\n",
      "Debug normal model\n",
      "Trace Shapes:    \n",
      " Param Sites:    \n",
      "Sample Sites:    \n",
      "   batch dist   |\n",
      "        value 3 |\n",
      "     log_prob   |\n",
      " n_boxes dist   |\n",
      "        value 4 |\n",
      "     log_prob   |\n",
      "   box_0 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n",
      "   box_1 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n",
      "   box_2 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n",
      "   box_3 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n"
     ]
    }
   ],
   "source": [
    "def model():\n",
    "    p = 0.1*torch.Tensor([[1,2,3,4], [5,5,5,5], [7,7,7,7]])\n",
    "    with pyro.iarange('batch', 3):\n",
    "        values = []\n",
    "        for i in pyro.irange('n_boxes', 4):\n",
    "            s = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i]))\n",
    "            values.append(s)\n",
    "        config = torch.stack(values, dim=1)\n",
    "        \n",
    "def guide():\n",
    "    p = 0.1*torch.Tensor([[1,2,3,4], [5,5,5,5], [7,7,7,7]])\n",
    "    with pyro.iarange('batch', 3):\n",
    "        values = []\n",
    "        for i in pyro.irange('n_boxes', 4):\n",
    "            s = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i]))\n",
    "            values.append(s)\n",
    "        config = torch.stack(values, dim=1)\n",
    "        print(\"config_normal:\",config.shape,config)\n",
    "        \n",
    "elbo = Trace_ELBO(max_plate_nesting=1,num_particles=2)\n",
    "elbo.loss(model,guide);\n",
    "\n",
    "print(\"------------------\")\n",
    "print(\"Debug normal model\")\n",
    "trace = poutine.trace(model).get_trace()\n",
    "trace.compute_log_prob()  \n",
    "print(trace.format_shapes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sequential enumeration\n",
    "\n",
    "Note that:\n",
    "1. we use TraceEnum_ELBO with max_plate_nesting=1 \n",
    "2. we could add \"num_particles\" to draw more samples. This make ssense only if we have both enumerated and sampled variables.\n",
    "3. the guide generates 2^(N_BOXES) configurations of shape \"1 x N_BOXES\", i.e. we do NOT enumerate over the bathc dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config_seq: torch.Size([1, 4]) tensor([[1., 1., 1., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 1., 1., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 1., 0., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 1., 0., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 0., 1., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 0., 1., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 0., 0., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[1., 0., 0., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 1., 1., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 1., 1., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 1., 0., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 1., 0., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 0., 1., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 0., 1., 0.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 0., 0., 1.]])\n",
      "config_seq: torch.Size([1, 4]) tensor([[0., 0., 0., 0.]])\n",
      "----------------------\n",
      "Debug sequential model\n",
      "Trace Shapes:    \n",
      " Param Sites:    \n",
      "Sample Sites:    \n",
      "   batch dist   |\n",
      "        value 3 |\n",
      "     log_prob   |\n",
      " n_boxes dist   |\n",
      "        value 4 |\n",
      "     log_prob   |\n",
      "   box_0 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n",
      "   box_1 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n",
      "   box_2 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n",
      "   box_3 dist 3 |\n",
      "        value 3 |\n",
      "     log_prob 3 |\n"
     ]
    }
   ],
   "source": [
    "def model_seq():\n",
    "    p = 0.1*torch.Tensor([[1,2,3,4], [5,5,5,5], [7,7,7,7]])\n",
    "    with pyro.iarange('batch', 3):\n",
    "        values = []\n",
    "        for i in pyro.irange('n_boxes', 4):\n",
    "            s = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i]),infer={\"enumerate\": \"sequential\"})\n",
    "            values.append(s)\n",
    "        config = torch.stack(values, dim=1)\n",
    "        \n",
    "def guide_seq():\n",
    "    p = 0.1*torch.Tensor([[1,2,3,4], [5,5,5,5], [7,7,7,7]])\n",
    "    with pyro.iarange('batch', 3):\n",
    "        values = []\n",
    "        for i in pyro.irange('n_boxes', 4):\n",
    "            s = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i]),infer={\"enumerate\": \"sequential\"})\n",
    "            values.append(s)\n",
    "        config = torch.stack(values, dim=1)\n",
    "        print(\"config_seq:\",config.shape,config)\n",
    "        \n",
    "elbo = TraceEnum_ELBO(max_plate_nesting=1)\n",
    "elbo.loss(model_seq,guide_seq);      \n",
    "\n",
    "print(\"----------------------\")\n",
    "print(\"Debug sequential model\")\n",
    "trace = poutine.trace(model_seq).get_trace()\n",
    "trace.compute_log_prob()  \n",
    "print(trace.format_shapes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel enumeration\n",
    "\n",
    "Note that:\n",
    "1. we use TraceEnum_ELBO with max_plate_nesting=1 \n",
    "2. we could add \"num_particles\" to draw more samples. This make ssense only if we have both enumerated and sampled variables.\n",
    "3. the guide generates a tensor of size 2 x ... x 1 for each box variable. To obtain the enumeration over the cartesian space you need to multiply the tensors. This is possible b/c the size are brodcastasble.\n",
    "4. Parallel enumeration makes sense if there are \"obs\" statement that depend on some only on some of the variables. If the \"obs\" statement depends on the complete enumeration it make sense to use sequential numeration instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s ->  torch.Size([2, 1]) tensor([[0.],\n",
      "        [1.]])\n",
      "s ->  torch.Size([2, 1, 1]) tensor([[[0.]],\n",
      "\n",
      "        [[1.]]])\n",
      "s ->  torch.Size([2, 1, 1, 1]) tensor([[[[0.]]],\n",
      "\n",
      "\n",
      "        [[[1.]]]])\n",
      "s ->  torch.Size([2, 1, 1, 1, 1]) tensor([[[[[0.]]]],\n",
      "\n",
      "\n",
      "\n",
      "        [[[[1.]]]]])\n",
      "p_boxes.shape torch.Size([2, 2, 2, 2, 1])\n",
      "--------------------\n",
      "Debug parallel model\n",
      "Trace Shapes:            \n",
      " Param Sites:            \n",
      "Sample Sites:            \n",
      "   batch dist           |\n",
      "        value         3 |\n",
      "     log_prob           |\n",
      " n_boxes dist           |\n",
      "        value         4 |\n",
      "     log_prob           |\n",
      "   box_0 dist         3 |\n",
      "        value       2 1 |\n",
      "     log_prob       2 3 |\n",
      "   box_1 dist         3 |\n",
      "        value     2 1 1 |\n",
      "     log_prob     2 1 3 |\n",
      "   box_2 dist         3 |\n",
      "        value   2 1 1 1 |\n",
      "     log_prob   2 1 1 3 |\n",
      "   box_3 dist         3 |\n",
      "        value 2 1 1 1 1 |\n",
      "     log_prob 2 1 1 1 3 |\n"
     ]
    }
   ],
   "source": [
    "def model_par():\n",
    "    p = 0.1*torch.Tensor([[1,2,3,4], [5,5,5,5], [7,7,7,7]])\n",
    "    with pyro.iarange('batch', 3):\n",
    "        values = []\n",
    "        for i in pyro.irange('n_boxes', 4):\n",
    "            s = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i]),infer={\"enumerate\": \"parallel\"})\n",
    "            values.append(s)\n",
    "        #config = torch.stack(values, dim=1)\n",
    "        \n",
    "def guide_par():\n",
    "    p = 0.1*torch.Tensor([[1,2,3,4], [5,5,5,5], [7,7,7,7]])\n",
    "    with pyro.iarange('batch', 3):\n",
    "        values = []\n",
    "        p_boxes = torch.ones(1)\n",
    "        for i in pyro.irange('n_boxes', 4):\n",
    "            s = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i]),infer={\"enumerate\": \"parallel\"})\n",
    "            values.append(s)\n",
    "            p_boxes = p_boxes * s\n",
    "            print(\"s -> \",s.shape,s)\n",
    "        print(\"p_boxes.shape\",p_boxes.shape)\n",
    "\n",
    "elbo = TraceEnum_ELBO(max_plate_nesting=1)\n",
    "elbo.loss(model_par,guide_par);\n",
    "\n",
    "print(\"--------------------\")\n",
    "print(\"Debug parallel model\")    \n",
    "trace = poutine.trace(poutine.enum(model_par, first_available_dim=-2)).get_trace()\n",
    "trace.compute_log_prob()  # optional, but allows printing of log_prob shapes\n",
    "print(trace.format_shapes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try the real deal\n",
    "\n",
    "It is extremely important that the enumerate variables are last.\n",
    "This is b/c in pyro if \"a\" follows \"b\" then \"b\" might depend on \"a\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 1., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[1., 0., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 1., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 1., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 1., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config: torch.Size([1, 4]) tensor([[0., 0., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 0., 1.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 0., 0.]])\n",
      "z torch.Size([3, 4, 50]) tensor([[ 0.6550,  0.5389,  0.0536,  0.4080],\n",
      "        [ 0.2538,  1.3784, -0.5017, -1.5882],\n",
      "        [ 0.0999, -0.1163,  0.2889,  0.7559]])\n",
      "config: torch.Size([1, 4]) tensor([[0., 0., 0., 0.]])\n",
      "--------------------\n",
      "Debug real model\n",
      "z torch.Size([3, 4, 50]) tensor([[-1.2063, -0.3114, -0.4207, -0.7088],\n",
      "        [-1.4023, -1.2109, -0.1666, -2.0044],\n",
      "        [ 0.0248, -1.2643, -1.7743, -0.0443]])\n",
      "config: torch.Size([3, 4]) tensor([[1., 1., 0., 1.],\n",
      "        [1., 1., 1., 0.],\n",
      "        [0., 1., 1., 1.]])\n",
      "Trace Shapes:         \n",
      " Param Sites:         \n",
      "Sample Sites:         \n",
      "   batch dist   |     \n",
      "        value 3 |     \n",
      "     log_prob   |     \n",
      "latent_z dist 3 | 4 50\n",
      "        value 3 | 4 50\n",
      "     log_prob 3 |     \n",
      " n_boxes dist   |     \n",
      "        value 4 |     \n",
      "     log_prob   |     \n",
      "   box_0 dist 3 |     \n",
      "        value 3 |     \n",
      "     log_prob 3 |     \n",
      "   box_1 dist 3 |     \n",
      "        value 3 |     \n",
      "     log_prob 3 |     \n",
      "   box_2 dist 3 |     \n",
      "        value 3 |     \n",
      "     log_prob 3 |     \n",
      "   box_3 dist 3 |     \n",
      "        value 3 |     \n",
      "     log_prob 3 |     \n"
     ]
    }
   ],
   "source": [
    "def model_and_guide():\n",
    "            \n",
    "    batch_size = 3\n",
    "    N_DAPI_MAX = 4\n",
    "    Zwhat_dim = 50\n",
    "    p = torch.rand((batch_size,N_DAPI_MAX,1))\n",
    "    \n",
    "    with pyro.iarange('batch', batch_size):\n",
    "        \n",
    "        z = pyro.sample(\"latent_z\",dist.Normal(0,1).expand([N_DAPI_MAX,Zwhat_dim]).to_event(2))\n",
    "        print(\"z\",z.shape,z[:,:,12])\n",
    "        \n",
    "        # Sample the c values at the end\n",
    "        values = []\n",
    "        for i in pyro.irange('n_boxes', N_DAPI_MAX):\n",
    "            c = pyro.sample('box_%d' % i, dist.Bernoulli(probs=p[:,i,0]),infer={\"enumerate\": \"sequential\"})\n",
    "            values.append(c)\n",
    "        config = torch.stack(values, dim=1)\n",
    "        print(\"config:\",config.shape,config)\n",
    "        \n",
    "\n",
    "\n",
    "elbo = TraceEnum_ELBO(max_plate_nesting=1)\n",
    "elbo.loss(model_and_guide,model_and_guide);\n",
    "\n",
    "print(\"--------------------\")\n",
    "print(\"Debug real model\")    \n",
    "trace = poutine.trace(model_and_guide).get_trace()\n",
    "trace.compute_log_prob()  \n",
    "print(trace.format_shapes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### VAE CLASS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Compositional_VAE(torch.nn.Module):\n",
    "    \n",
    "    def __init__(self, params,inference,generator):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Instantiate the encoder and decoder\n",
    "        self.inference = inference\n",
    "        self.generator = generator\n",
    "        \n",
    "        # Parameters\n",
    "        self.single_image_dim = params['SD.dim_output']\n",
    "        self.EPS_UNIFORM = params['safety_uniform_dist']   \n",
    "        self.x_uniform_dist = 1.0+self.EPS_UNIFORM \n",
    "        self.N_DAPI_MAX = params['YOLO.N_dapi_max']\n",
    "        self.Zwhere_dim = params['UNET.Zwhere_dim']\n",
    "        self.Zwhat_dim  = params['UNET.Zwhat_dim']\n",
    "        self.N_DAPI_EXPECTED = params['expected_number_of_cells']\n",
    "        self.p_bernulli = float(self.N_DAPI_EXPECTED)/self.N_DAPI_MAX\n",
    "        bb_average = 2.0*params['bb_average_UNET_grid'] / params['UNET.W_coarsest']\n",
    "        print(\"Debug bb_avergae\",bb_average)\n",
    "        bb_std = bb_average*params['bb_std_over_mean']\n",
    "        bb_var = bb_std*bb_std\n",
    "        self.bb_gamma_beta  = bb_average/bb_var\n",
    "        self.bb_gamma_alpha = bb_average*self.bb_gamma_beta\n",
    "        print(\"Debug alpha,beta\",self.bb_gamma_alpha,self.bb_gamma_beta)\n",
    "           \n",
    "        # Put everything on the cude if necessary\n",
    "        self.use_cuda = params[\"use_cuda\"]\n",
    "        if self.use_cuda:\n",
    "            self.cuda()  \n",
    "    \n",
    "      \n",
    "    def model(self, imgs=None):\n",
    "        \"\"\" 1.  sample the latent from the prior:\n",
    "            2.  runs the generative model\n",
    "            3.  return the generated image (and if actual data is passed score the generative model against it)\n",
    "        \"\"\"\n",
    "        \n",
    "        #-----------------------#\n",
    "        #--------  Trick -------#\n",
    "        #-----------------------#\n",
    "        if(imgs is None):\n",
    "            observed = False\n",
    "            W = self.generator.W_raw_image\n",
    "            CH = self.generator.CH_raw_image \n",
    "            imgs = torch.zeros(8,CH,W,W)\n",
    "            if(self.use_cuda):\n",
    "                imgs=imgs.cuda()\n",
    "        else:\n",
    "            observed = True\n",
    "        #-----------------------#\n",
    "        #----- Enf of Trick ----#\n",
    "        #-----------------------#\n",
    "            \n",
    "        batch_size = imgs.shape[0]\n",
    "        \n",
    "        # This is a trick to have zeor and one written on the same device as the images\n",
    "        zero = imgs.new_zeros(1)\n",
    "        one = imgs.new_ones(1)\n",
    "        p_bernulli = self.p_bernulli*one\n",
    "        alpha = self.bb_gamma_alpha*one\n",
    "        beta = self.bb_gamma_beta*one\n",
    "        x_uniform_dist = self.x_uniform_dist*one\n",
    "        # The center of the cell can be slightly outside the image, \n",
    "        # therefore I set the support of the distribution to (-1.1,1.1) \n",
    "        \n",
    "        pyro.module(\"generator\", self.generator)\n",
    "        \n",
    "        with pyro.plate(\"batch\", batch_size, dim=-1):\n",
    "            \n",
    "            z = pyro.sample(\"latent_z\",\n",
    "                            dist.Normal(zero,one).expand([self.N_DAPI_MAX,self.Zwhat_dim]).to_event(2))\n",
    "\n",
    "            cxcy = pyro.sample(\"center_bb\",\n",
    "                               dist.Uniform(-x_uniform_dist,x_uniform_dist).expand([self.N_DAPI_MAX,2]).to_event(2))\n",
    "            dxdy = pyro.sample(\"size_bb\",\n",
    "                               dist.Gamma(alpha,beta).expand([self.N_DAPI_MAX,2]).to_event(2))\n",
    "            \n",
    "            # Probability of a cell\n",
    "            values = []\n",
    "            for i in pyro.irange('n_boxes', self.N_DAPI_MAX):\n",
    "                c = pyro.sample('c_%d' % i, dist.Bernoulli(p_bernulli))\n",
    "                values.append(c)\n",
    "            c_cell = torch.stack(values, dim=1)\n",
    "            c_config = c_cell.unsqueeze(2).expand(dxdy.shape[0],-1,-1)\n",
    "            #print(\"c_config.shape\",c_config.shape)\n",
    "            assert((batch_size,self.N_DAPI_MAX,1)==c_config.shape)\n",
    "            \n",
    "            # Invert stuff for easier work by the generator\n",
    "            sxsy = 2.0/dxdy\n",
    "            kxky = -cxcy*sxsy  \n",
    "            x_mu = self.generator.forward(c_config,sxsy,kxky,z,self.detach_single_dapi_decoder)\n",
    "            \n",
    "            if(observed):\n",
    "                x_std = 0.1 * torch.ones_like(x_mu)\n",
    "                pyro.sample(\"obs\",dist.Normal(x_mu.view(-1,self.single_image_dim),x_std.view(-1,self.single_image_dim)).to_event(1),\n",
    "                            obs=imgs.view(-1,self.single_image_dim))\n",
    "            \n",
    "            return x_mu,c_config,dxdy,cxcy\n",
    "            \n",
    "    def guide(self, imgs=None):\n",
    "        \"\"\" The GUIDE takes a mini-batch of images and: \n",
    "            1. run the inference to get z_where,z_what\n",
    "            encoder to get: p,sxsy_mu,kxky_mu,z_mu,z_std\n",
    "            2. Compute cxcy_mu and dxdy_mu\n",
    "            2. sample:\n",
    "                - z ~ N(z_mu, z_std) where each component of z is drawn independently\n",
    "                - c ~ Bernulli(p)\n",
    "                - cxcy ~ N(cxcy_mu,0.1) \n",
    "                - dxdy ~ gamma(dxdy_mu,0.1) \n",
    "        \"\"\"\n",
    "        \n",
    "        #-----------------------#\n",
    "        #--------  Trick -------#\n",
    "        #-----------------------#\n",
    "        if(imgs is None):\n",
    "            observed = False\n",
    "            W = self.generator.W_raw_image\n",
    "            CH = self.generator.CH_raw_image \n",
    "            imgs = torch.zeros(8,CH,W,W)\n",
    "            if(self.use_cuda):\n",
    "                imgs=imgs.cuda()\n",
    "        else:\n",
    "            observed = True    \n",
    "        #-----------------------#\n",
    "        #----- Enf of Trick ----#\n",
    "        #-----------------------#\n",
    "               \n",
    "        batch_size = imgs.shape[0]            \n",
    "        pyro.module(\"inference\",self.inference)\n",
    "        \n",
    "        #--------------------------#\n",
    "        #-- 1. run the inference --#\n",
    "        #--------------------------#\n",
    "        p,sxsy_mu,kxky_mu,z_mu,z_std = self.inference.forward(imgs)        \n",
    "        \n",
    "        # Transform stuff around for easier sampling\n",
    "        dxdy_mu = 2.0/sxsy_mu      # size of the bounding box \n",
    "        cxcy_mu = torch.clamp(-kxky_mu/sxsy_mu, min=-self.x_uniform_dist, max=self.x_uniform_dist) \n",
    "                                  \n",
    "        #---------------#    \n",
    "        #-- 2. Sample --#\n",
    "        #---------------#\n",
    "        with pyro.plate(\"batch\", batch_size, dim =-1 ):   \n",
    "            \n",
    "            # Latent z\n",
    "            pyro.sample(\"latent_z\",dist.Normal(z_mu, z_std).to_event(2))\n",
    "            \n",
    "            # Centers\n",
    "            pyro.sample(\"center_bb\",dist.Normal(cxcy_mu, 0.01*torch.ones_like(cxcy_mu)).to_event(2))\n",
    "            \n",
    "            # Size of bounding boxes\n",
    "            alpha_dxdy = 100*torch.ones_like(dxdy_mu)\n",
    "            beta_dxdy = alpha_dxdy / dxdy_mu\n",
    "            pyro.sample(\"size_bb\",dist.Gamma(alpha_dxdy, beta_dxdy).to_event(2)) \n",
    "            \n",
    "            # Probability of a cell\n",
    "            values = []\n",
    "            for i in pyro.irange('n_boxes', self.N_DAPI_MAX):\n",
    "                c = pyro.sample('c_%d' % i, dist.Bernoulli(probs=p[:,i,0]),infer={\"enumerate\": \"sequential\"})\n",
    "                values.append(c)\n",
    "            c_cell = torch.stack(values, dim=1)\n",
    "            #print(\"c_cell:\",c_cell.shape,c_cell)\n",
    "                \n",
    "                   \n",
    "    def train_one_epoch(self, svi, train_loader, desired_channels,verbose):\n",
    "        self.train()\n",
    "        epoch_loss = 0.\n",
    "        normalizer_train = 0.\n",
    "        for i, x in enumerate(train_loader): #get the images\n",
    "                      \n",
    "            # set the number of channels to desired_channels\n",
    "            B,CH,W,H = x.shape\n",
    "            x = x.view(-1,desired_channels,W,H)\n",
    "            normalizer_train += float(B*CH)/desired_channels\n",
    "            \n",
    "            # if on GPU put mini-batch into CUDA memory\n",
    "            if self.use_cuda:\n",
    "                x = x.cuda()\n",
    "            loss = svi.step(x)\n",
    "            if(verbose):\n",
    "                print(\"i= %3d train_loss=%.5f\" %(i,loss))\n",
    "            epoch_loss += loss\n",
    "\n",
    "        total_epoch_loss_train = epoch_loss / normalizer_train\n",
    "        return total_epoch_loss_train\n",
    "\n",
    "    def evaluate_one_epoch(self, svi, test_loader, desired_channels,verbose):\n",
    "        self.eval()\n",
    "        test_loss = 0.\n",
    "        normalizer_test = 0.\n",
    "        for i, x in enumerate(test_loader): #get the images\n",
    "            \n",
    "            # set the number of channels to desired_channels\n",
    "            B,CH,W,H = x.shape\n",
    "            x = x.view(-1,desired_channels,W,H)\n",
    "            normalizer_test += float(B*CH)/desired_channels\n",
    "            \n",
    "            # if on GPU put mini-batch into CUDA memory\n",
    "            if self.use_cuda:\n",
    "                x = x.cuda()\n",
    "            loss = svi.evaluate_loss(x)\n",
    "            if(verbose):\n",
    "                print(\"i= %3d test_loss=%.5f\" %(i,loss))\n",
    "            test_loss += loss\n",
    "\n",
    "        total_epoch_loss_test = test_loss / normalizer_test\n",
    "        return total_epoch_loss_test \n",
    "    \n",
    "    \n",
    "    def generate_synthetic_data(self, N=100):\n",
    "        \n",
    "        x_mu,c,dxdy,cxcy = self.model()\n",
    "        BATCH, CH, W, H = x_mu.shape    \n",
    "        synthetic_data = torch.zeros((N, CH, W, H), dtype=torch.float32, device='cpu', requires_grad=False) \n",
    "        \n",
    "        l = 0 \n",
    "        while (l<N):\n",
    "            # Compute left and right indeces\n",
    "            r = min(l+BATCH,N)\n",
    "            d = r-l\n",
    "            synthetic_data[l:r,:,:,:]=x_mu[0:d,:,:,:].detach().cpu()\n",
    "            l = r\n",
    "            x_mu,c,dxdy,cxcy = self.model()\n",
    "        return synthetic_data\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate simulation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'safety_uniform_dist': 0.1, 'expected_number_of_cells': 0.9, 'bb_average_UNET_grid': 1.0, 'bb_std_over_mean': 0.1, 'UNET.W_coarsest': 2, 'UNET.N_max_pool': 4, 'UNET.N_up_conv': 0, 'UNET.CH_after_first_two_conv': 8, 'UNET.N_prediction_maps': 1, 'UNET.Zwhat_dim': 50, 'UNET.Zwhere_dim': 5, 'UNET.CH_raw_image': 1, 'UNET.W_raw_image': 32, 'UNET.W_output': 2, 'SD.W': 32, 'SD.output_channel': 1, 'SD.dim_z': 50, 'SD.dim_output': 1024, 'SD.dim_h1': 400, 'SD.dim_h2': -1, 'YOLO.p_threshold': 0.0, 'YOLO.IoU_threshold': 0.2, 'YOLO.N_dapi_max': 4, 'device': device(type='cpu'), 'dtype': torch.float32, 'dtype_tensor': <class 'torch.FloatTensor'>, 'use_cuda': False}\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from low_level_utility import DataSet_for_cell_segmentation, Synthetic_DataSet, dataloader, N_RandomCrop, show_batch \n",
    "from low_level_utility import overlay_image_and_unet_grid, save_model, save_obj, load_model, load_obj, test_model\n",
    "from unet import partial_UNET, CNN_Net\n",
    "from yolo import YOLO_filter\n",
    "from dapi_generator import CompositionalGenerator\n",
    "from simulation_dictionary import SimulationDictionary \n",
    "#from vae import Compositional_VAE\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as tvt\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer import SVI, Trace_ELBO\n",
    "from pyro.optim import Adam, Adamax, SGD\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.data import TensorDataset\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#### Create the simulation parameters dictionary\n",
    "params=SimulationDictionary()\n",
    "\n",
    "# Check the params              \n",
    "params[\"expected_number_of_cells\"]=0.9\n",
    "print(params)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LOAD SYNTHETIC DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'num_workers': 1, 'pin_memory': False}\n",
      "{'num_workers': 1, 'pin_memory': False}\n",
      "torch.Size([16, 1, 32, 32])\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAADOCAYAAAAuRfn0AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAEPdJREFUeJzt3X+s3XV9x/Hni58OxUDhlnUtXWFp/IGZ/LhxMJbFgU50BlgiC41xDavpP27DzUVhJJRG/5BsQZ1xbo0ozcJAqLLWxumaijH7p+PyQ34VBBFLpbZXAXUapy3v/XG+xWu5l3Puvefcc/rl+Uhuvt/v5/s99/vKp+e876efc77nm6pCknT4O2LYASRJ/WFBl6SWsKBLUktY0CWpJSzoktQSFnRJagkLuiS1xLwKepKLkjya5PEkV/UrlCRp9jLXC4uSHAl8C3grsBu4C1hVVQ/3L54kqVdHzeOxbwIer6onAJLcClwCzFjQTz755FqxYsU8TilJLz933333D6pqrNtx8ynoS4GnpmzvBn7v0IOSrAXWAixfvpyJiYl5nFKSXn6SfLeX4+Yzh55p2l40f1NVG6pqvKrGx8a6/oGRJM3RfAr6buDUKdvLgKfnF0eSNFfzKeh3ASuTnJbkGOByYEt/YkmSZmvOc+hVtT/JXwJfBY4EPltVD/UtmSRpVubzpihV9WXgy33KIkmaB68UlaSWsKBLUkvMa8plEK677rq+/r4vfelLL6zfc889s378qlWrXlhfsmRJXzLNxQ033ABAr1f2zqcf169fP+fHHi566cdkuk/m6qBen4sH+3H//v0AHHHEaIwjjz76aAAOHDgw1Bz9vA3oaPSsJGneRm6EPmqOO+64F9aPP/74ISaRpJfmCF2SWsKCLkktYUGXpJawoEtSS1jQJaklLOiS1BIWdElqCQu6JLVE14Ke5LNJ9iV5cErboiTbkjzWLE8cbExJUje9jNBvAi46pO0qYHtVrQS2N9uSpCHqWtCr6hvAM4c0XwJsbNY3Apf2OZckaZbmOod+SlXtAWiWi/sXSZI0FwN/UzTJ2iQTSSYmJycHfTpJetma67ct7k2ypKr2JFkC7JvpwKraAGwAGB8f798X/w7B8uXLAbjiiiuGnESSXmyuI/QtwOpmfTWwuT9xJElz1XWEnuQW4M3AyUl2A+uAjwK3JVkD7AIuG2RISfN3zjnnDDuCBqxrQa+qVTPsurDPWSRJ8+Adi1rqk5/85LAjaMTcddddw46gAfPSf0lqCQu6JLWEUy4t9cwzh17cK6ntHKFLUktY0CWpJSzoktQSzqFLWhBHHWW5GTRH6JLUEiP3J3P9+vXDjiBJhyVH6JLUEhZ0SWoJC7oktYQFXZJaomtBT3JqkjuT7EzyUJIrm/ZFSbYleaxZnjj4uJKkmfQyQt8PfKCqXgecC7wvyeuBq4DtVbUS2N5sS5KGpGtBr6o9VXVPs/4TYCewFLgE2NgcthG4dFAhJUndzWoOPckK4CxgB3BKVe2BTtEHFs/wmLVJJpJMTE5Ozi+tJGlGPRf0JK8CvgC8v6p+3OvjqmpDVY1X1fjY2NhcMkqSetDTlaJJjqZTzG+uqi82zXuTLKmqPUmWAPsGFXKY9u7d+8L6rl27AHjqqaeGFUeSZtTLp1wC3AjsrKobpuzaAqxu1lcDm/sfT5LUq15G6OcD7wEeSHJf0/b3wEeB25KsAXYBlw0m4nBt3br1hfX7779/iEkk6aV1LehV9d9AZth9YX/jSJLmyitFJaklLOiS1BIWdElqCQu6JLWEBV2SWmLkbkGnl69169YNO4J0WHOELkktYUGXpJawoEtSSziHLr1MXHvttX39fR/+8If7+vs0f47QJaklHKFrQfTyCZZBfcpl/fr1A/m9h5uPfOQjff19jtBHjyN0SWoJC7oktUTXKZckrwC+ARzbHL+pqtYlOQ24FVgE3AO8p6p+Mciww3bwjkXD/C/8ddddN7RzSxptvYzQ/w+4oKreCJwJXJTkXOB64GNVtRJ4FlgzuJiSpG66FvTq+N9m8+jmp4ALgE1N+0bg0oEklCT1pKc59CRHNref2wdsA74NPFdV+5tDdgNLZ3js2iQTSSYmJyf7kVmSNI2eCnpVHaiqM4FlwJuA10132AyP3VBV41U1PjY2NvekkqSXNKtPuVTVc8DXgXOBE5IcfFN1GfB0f6NJkmaja0FPMpbkhGb9N4C3ADuBO4F3NYetBjYPKqQkqbterhRdAmxMciSdPwC3VdXWJA8Dtyb5CHAvcOMAc0qSuuha0KvqfuCsadqfoDOfLkkaAV4pKkktYUGXpJawoEtSS1jQJaklLOiS1BIWdElqiZG7Y1HVtN8goFmyH39lrl85bB/2h/24cByhS1JLWNAlqSUs6JLUEhZ0SWoJC7oktYQFXZJaoueC3tyG7t4kW5vt05LsSPJYks8nOWZwMSVJ3cxmhH4lnRtbHHQ98LGqWgk8C6zpZzBJ0uz0epPoZcCfAJ9ptgNcAGxqDtkIXDqIgJKk3vR6pejHgQ8CxzfbJwHPVdX+Zns3sHS6ByZZC6wFWL58edcTdf5WaCa9XnVnP760XvrRPnxph/tz8fnnnx92BKC//dPLPUXfCeyrqrunNk9z6LT/ulW1oarGq2p8bGxsjjElSd30MkI/H7g4yTuAVwCvpjNiPyHJUc0ofRnw9OBiSpK66TpCr6qrq2pZVa0ALge+VlXvBu4E3tUcthrYPLCUkqSu5vM59A8Bf5vkcTpz6jf2J5IkaS5m9fW5VfV14OvN+hPAm/ofSZI0F14pKkktYUGXpJawoEtSS1jQJaklLOiS1BIjd5NoSb/uve9977AjALBjxw4AHnjggSEn0UwcoUtSSzhCl0bc0qXTfu/dgjv++OO7H6ShcoQuSS3hCF2awbp16wZ+jttvvx2Ahx9+eODnUvs5QpeklrCgS1JLWNAlqSUs6JLUEj29KZrkSeAnwAFgf1WNJ1kEfB5YATwJ/FlVPTuYmJKkbmYzQv+jqjqzqsab7auA7VW1EtjebEuShmQ+Uy6XABub9Y3ApfOPI0maq14LegH/leTuJGubtlOqag9As1w83QOTrE0ykWRicnJy/oklSdPq9cKi86vq6SSLgW1JHun1BFW1AdgAMD4+XnPIKEnqQU8j9Kp6ulnuA+6gcy/RvUmWADTLfYMKKUnqrmtBT/LKJMcfXAf+GHgQ2AKsbg5bDWweVEhJUne9TLmcAtyR5ODx/15VX0lyF3BbkjXALuCywcWUpP464ojRuAynqn8z0V0LelU9AbxxmvYfAhf2LYkkaV5G40+UJGneLOiS1BIWdElqCQu6JLWEBV2SWsKCLkktYUGXpJawoEtSS1jQJaklLOiS1BIWdElqiV6/D1162Vm3bt3Az3HGGWcA8NBDD71o3/r16wd+frWLI3RJagkLuiS1RE8FPckJSTYleSTJziTnJVmUZFuSx5rliYMOK0maWa8j9E8AX6mq19L5bvSdwFXA9qpaCWxvtiVJQ9LLLeheDfwhcCNAVf2iqp4DLgE2NodtBC4dVEhJUne9jNBPByaBzyW5N8lnmnuLnlJVewCa5eLpHpxkbZKJJBOTk5N9Cy5J+nW9FPSjgLOBT1fVWcBPmcX0SlVtqKrxqhofGxubY0xJUje9FPTdwO6q2tFsb6JT4PcmWQLQLPcNJqIkqRddC3pVfR94KslrmqYLgYeBLcDqpm01sHkgCSVJPen1StG/Am5OcgzwBHAFnT8GtyVZA+wCLhtMRElSL3oq6FV1HzA+za4L+xtHkjRXfpeLNIOf//znAz/HgQMHuh7jd7qoV176L0kt4QhdmsH1118/7AjSrDhCl6SWsKBLUktY0CWpJSzoktQSFnRJagkLuiS1hAVdklrCgi5JLWFBl6SWsKBLUkv0ck/R1yS5b8rPj5O8P8miJNuSPNYsT1yIwJKk6fVyg4tHq+rMqjoTOAf4GXAHndvQba+qlcB2ZnFbOklS/812yuVC4NtV9V3gEmBj074RuLSfwSRJszPbgn45cEuzfkpV7QFoloune0CStUkmkkxMTk7OPakk6SX1XNCb289dDNw+mxNU1YaqGq+q8bGxsdnmkyT1aDYj9LcD91TV3mZ7b5IlAM1yX7/DSZJ6N5uCvopfTbcAbAFWN+urgc39CiVJmr1UVfeDkuOAp4DTq+pHTdtJwG3AcmAXcFlVPfNSv2d8fLwmJibmHVqSXk6S3F1V492O6+kWdFX1M+CkQ9p+SOdTL5KkEeCVopLUEhZ0SWoJC7oktYQFXZJawoIuSS1hQZeklrCgS1JL9HRhUd9OlkwCPwV+sGAn7Y+TMfNCORxzm3lhHI6ZoT+5f7uqun4Z1oIWdIAkE71c8TRKzLxwDsfcZl4Yh2NmWNjcTrlIUktY0CWpJYZR0DcM4ZzzZeaFczjmNvPCOBwzwwLmXvA5dEnSYDjlIkktYUGXpJZYsIKe5KIkjyZ5PMlVC3Xe2UpyapI7k+xM8lCSK5v2RUm2JXmsWZ447KyHSnJkknuTbG22T0uyo8n8+ea+sCMjyQlJNiV5pOnv80a9n5P8TfO8eDDJLUleMYr9nOSzSfYleXBK27R9m45/al6b9yc5e4Qy/0Pz/Lg/yR1JTpiy7+om86NJ3jYqmafs+7skleTkZnvg/bwgBT3JkcCn6NyX9PXAqiSvX4hzz8F+4ANV9TrgXOB9TdargO1VtRLY3myPmiuBnVO2rwc+1mR+FlgzlFQz+wTwlap6LfBGOtlHtp+TLAX+GhivqjcARwKXM5r9fBNw0SFtM/Xt24GVzc9a4NMLlPFQN/HizNuAN1TV7wLfAq4GaF6TlwNnNI/556bOLLSbeHFmkpwKvJXO3dwOGnw/V9XAf4DzgK9O2b4auHohzt2H7Jubf5hHgSVN2xLg0WFnOyTnMjov0guArUDoXJ121HT/BsP+AV4NfIfmjfkp7SPbz8BSOrdiXETnbl9bgbeNaj8DK4AHu/Ut8K/AqumOG3bmQ/b9KXBzs/5rNQT4KnDeqGQGNtEZpDwJnLxQ/bxQUy4HXwgH7W7aRlqSFcBZwA7glKraA9AsFw8v2bQ+DnwQeL7ZPgl4rqr2N9uj1uenA5PA55ppos8keSUj3M9V9T3gH+mMuvYAPwLuZrT7eaqZ+vZweX3+BfCfzfrIZk5yMfC9qvrmIbsGnnmhCnqmaRvpz0smeRXwBeD9VfXjYed5KUneCeyrqrunNk9z6Cj1+VHA2cCnq+osOt/xMzLTK9Np5pwvAU4Dfgt4JZ3/Rh9qlPq5F6P+XCHJNXSmQ28+2DTNYUPPnOQ44Brg2ul2T9PW18wLVdB3A6dO2V4GPL1A5561JEfTKeY3V9UXm+a9SZY0+5cA+4aVbxrnAxcneRK4lc60y8eBE5IcvBH4qPX5bmB3Ve1otjfRKfCj3M9vAb5TVZNV9Uvgi8DvM9r9PNVMfTvSr88kq4F3Au+uZq6C0c38O3T+4H+zeT0uA+5J8pssQOaFKuh3ASubTwMcQ+fNjC0LdO5ZSRLgRmBnVd0wZdcWYHWzvprO3PpIqKqrq2pZVa2g07dfq6p3A3cC72oOG7XM3weeSvKapulC4GFGuJ/pTLWcm+S45nlyMPPI9vMhZurbLcCfN5/COBf40cGpmWFLchHwIeDiqvrZlF1bgMuTHJvkNDpvNP7PMDJOVVUPVNXiqlrRvB53A2c3z/fB9/MCvnHwDjrvUn8buGYYb170mPMP6Pw36H7gvubnHXTmpLcDjzXLRcPOOkP+NwNbm/XT6TzJHwduB44ddr5Dsp4JTDR9/R/AiaPez8B64BHgQeDfgGNHsZ+BW+jM8/+STlFZM1Pf0pkK+FTz2nyAzqd4RiXz43TmnQ++Fv9lyvHXNJkfBd4+KpkP2f8kv3pTdOD97KX/ktQSXikqSS1hQZeklrCgS1JLWNAlqSUs6JLUEhZ0SWoJC7oktcT/A7+SwPVZP0OvAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#BATCH_SIZE = 256\n",
    "BATCH_SIZE = 16\n",
    "\n",
    "synthetic_data_train = load_obj(\"synthetic_data_train_v4\")  \n",
    "synthetic_data_test = load_obj(\"synthetic_data_test_v4\")  \n",
    "            \n",
    "kwargs = {'num_workers': 1, 'pin_memory': params[\"use_cuda\"]}\n",
    "\n",
    "train_data = Synthetic_DataSet(synthetic_data_train)\n",
    "test_data  = Synthetic_DataSet(synthetic_data_test)\n",
    "trainloader = dataloader(train_data,BATCH_SIZE,shuffle=True, **kwargs)\n",
    "testloader =  dataloader(test_data,BATCH_SIZE,shuffle=False, **kwargs)\n",
    "\n",
    "imgs=trainloader.load(CHANNELS=1)\n",
    "show_batch(imgs[:8],nrow=4,npadding=4)\n",
    "print(imgs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check UNET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#unet = CNN_Net(params)\n",
    "#unet.show_receptive_field()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#image_and_grid = overlay_image_and_unet_grid(imgs[0],unet.W_list)\n",
    "#L = int(image_and_grid.shape[0]/3)\n",
    "#show_batch(image_and_grid[0:L],nrow=4,npadding=10)\n",
    "#show_batch(image_and_grid[L:2*L],nrow=4,npadding=10)\n",
    "#show_batch(image_and_grid[2*L:],nrow=5,npadding=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Debug bb_avergae 1.0\n",
      "Debug alpha,beta 99.99999999999999 99.99999999999999\n",
      "--------------------\n",
      "Debug real model\n",
      " Trace Shapes:         \n",
      "  Param Sites:         \n",
      " Sample Sites:         \n",
      "    batch dist   |     \n",
      "         value 8 |     \n",
      "      log_prob   |     \n",
      " latent_z dist 8 | 4 50\n",
      "         value 8 | 4 50\n",
      "      log_prob 8 |     \n",
      "center_bb dist 8 | 4  2\n",
      "         value 8 | 4  2\n",
      "      log_prob 8 |     \n",
      "  size_bb dist 8 | 4  2\n",
      "         value 8 | 4  2\n",
      "      log_prob 8 |     \n",
      "  n_boxes dist   |     \n",
      "         value 4 |     \n",
      "      log_prob   |     \n",
      "      c_0 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "      c_1 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "      c_2 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "      c_3 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "--------------------\n",
      "Debug real guide\n",
      "                                       Trace Shapes:                   \n",
      "                                        Param Sites:                   \n",
      "            inference$$$0.cnn.0.double_conv.0.weight   8   1 3   3     \n",
      "            inference$$$0.cnn.0.double_conv.1.weight             8     \n",
      "              inference$$$0.cnn.0.double_conv.1.bias             8     \n",
      "            inference$$$0.cnn.0.double_conv.3.weight   8   8 3   3     \n",
      "            inference$$$0.cnn.0.double_conv.4.weight             8     \n",
      "              inference$$$0.cnn.0.double_conv.4.bias             8     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.0.weight  16   8 3   3     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.1.weight            16     \n",
      "  inference$$$0.cnn.1.double_conv.double_conv.1.bias            16     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.3.weight  16  16 3   3     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.4.weight            16     \n",
      "  inference$$$0.cnn.1.double_conv.double_conv.4.bias            16     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.0.weight  32  16 3   3     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.1.weight            32     \n",
      "  inference$$$0.cnn.2.double_conv.double_conv.1.bias            32     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.3.weight  32  32 3   3     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.4.weight            32     \n",
      "  inference$$$0.cnn.2.double_conv.double_conv.4.bias            32     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.0.weight  64  32 3   3     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.1.weight            64     \n",
      "  inference$$$0.cnn.3.double_conv.double_conv.1.bias            64     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.3.weight  64  64 3   3     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.4.weight            64     \n",
      "  inference$$$0.cnn.3.double_conv.double_conv.4.bias            64     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.0.weight 128  64 3   3     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.1.weight           128     \n",
      "  inference$$$0.cnn.4.double_conv.double_conv.1.bias           128     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.3.weight 128 128 3   3     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.4.weight           128     \n",
      "  inference$$$0.cnn.4.double_conv.double_conv.4.bias           128     \n",
      "    inference$$$0.cnn.5.compute_zwhere.comp_p.weight   1 128 1   1     \n",
      "      inference$$$0.cnn.5.compute_zwhere.comp_p.bias             1     \n",
      " inference$$$0.cnn.5.compute_zwhere.comp_txty.weight   2 128 1   1     \n",
      "   inference$$$0.cnn.5.compute_zwhere.comp_txty.bias             2     \n",
      " inference$$$0.cnn.5.compute_zwhere.comp_sxsy.weight   2 128 1   1     \n",
      "   inference$$$0.cnn.5.compute_zwhere.comp_sxsy.bias             2     \n",
      "  inference$$$0.cnn.5.compute_zwhat.comp_z_mu.weight  50 128 1   1     \n",
      "    inference$$$0.cnn.5.compute_zwhat.comp_z_mu.bias            50     \n",
      " inference$$$0.cnn.5.compute_zwhat.comp_z_std.weight  50 128 1   1     \n",
      "   inference$$$0.cnn.5.compute_zwhat.comp_z_std.bias            50     \n",
      "                                       Sample Sites:                   \n",
      "                                          batch dist             |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob             |     \n",
      "                                       latent_z dist         8   | 4 50\n",
      "                                               value         8   | 4 50\n",
      "                                            log_prob         8   |     \n",
      "                                      center_bb dist         8   | 4  2\n",
      "                                               value         8   | 4  2\n",
      "                                            log_prob         8   |     \n",
      "                                        size_bb dist         8   | 4  2\n",
      "                                               value         8   | 4  2\n",
      "                                            log_prob         8   |     \n",
      "                                        n_boxes dist             |     \n",
      "                                               value         4   |     \n",
      "                                            log_prob             |     \n",
      "                                            c_0 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n",
      "                                            c_1 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n",
      "                                            c_2 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n",
      "                                            c_3 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n"
     ]
    }
   ],
   "source": [
    "# Set up pyro environment\n",
    "pyro.clear_param_store()\n",
    "pyro.set_rng_seed(0)\n",
    "TEST_FREQUENCY = 5\n",
    "\n",
    "#smoke_test = ('CI' in os.environ)\n",
    "assert pyro.__version__.startswith('0.3.0')\n",
    "smoke_test= True\n",
    "if(smoke_test):\n",
    "    pyro.enable_validation(True)\n",
    "    pyro.distributions.enable_validation(True)\n",
    "    NUM_EPOCHS = 21\n",
    "else:\n",
    "    pyro.enable_validation(False)\n",
    "    pyro.distributions.enable_validation(False)\n",
    "    NUM_EPOCHS = 1001\n",
    "    \n",
    "\n",
    "# setup the ENCODER, DECODER and VAE \n",
    "inference = torch.nn.Sequential(CNN_Net(params), YOLO_filter(params))\n",
    "#generator = CompositionalGenerator(params,'DAPI')\n",
    "#generator = CompositionalGenerator(params,'DISK')\n",
    "generator = CompositionalGenerator(params,'SQUARE')\n",
    "vae       = Compositional_VAE(params, inference, generator)\n",
    "\n",
    "vae.detach_single_dapi_decoder = True\n",
    "\n",
    "elbo = TraceEnum_ELBO(max_plate_nesting=1)\n",
    "elbo.loss(vae.model,vae.guide);\n",
    "\n",
    "print(\"--------------------\")\n",
    "print(\"Debug real model\")    \n",
    "trace = poutine.trace(vae.model).get_trace()\n",
    "trace.compute_log_prob()  \n",
    "print(trace.format_shapes())\n",
    "\n",
    "print(\"--------------------\")\n",
    "print(\"Debug real guide\")    \n",
    "trace = poutine.trace(vae.guide).get_trace()\n",
    "trace.compute_log_prob()  \n",
    "print(trace.format_shapes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate synthetic data if necessary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAADOCAYAAAAuRfn0AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAEB1JREFUeJzt3X+MZXV9xvH3A7haQQMrs7BloYvNBrWmAp1YKE3TgihaAzSCAU07aTfZ/mFbbW0UStLEpH9o2vijibXdiLAQiiBCd0OslqyYpkmDDIKILLiICCvr7mhBrZoi7qd/3LNxXGe4587cX3t4v5LJvefcc+c8+e7cZ85+7z1zUlVIkg5/R0w6gCRpOCx0SeoIC12SOsJCl6SOsNAlqSMsdEnqCAtdkjpiVYWe5IIkDyd5JMkVwwolSRpcVnpiUZIjga8B5wN7gLuBy6vqweHFkyS1ddQqnvta4JGqehQgySeBi4BlC/3444+vjRs3rmKXkvT8c88993ynqmb6bbeaQj8JeGLR8h7gNw/dKMkWYAvAKaecwvz8/Cp2KUnPP0m+2Wa71cyhZ4l1vzB/U1Vbq2q2qmZnZvr+gpEkrdBqCn0PcPKi5Q3Ak6uLI0laqdUU+t3ApiSnJlkDXAbsGE4sSdKgVjyHXlXPJvkz4HPAkcAnquqrQ0smSRrIat4Upao+A3xmSFkkSavgmaKS1BEWuiR1xKqmXEYhWerTkN104MCBgZ/TdnyeT+O4Em3OkHYMn1vbs8zbjONKXgujNDc3B8D1118/8n0N8zKgHqFLUkdY6JLUERa6JHWEhS5JHWGhS1JHWOiS1BEWuiR1hIUuSR1hoUtSR1joktQRfQs9ySeS7E/ywKJ1a5PckWR3c3vcaGNKkvppc4R+LXDBIeuuAHZW1SZgZ7MsSZqgvoVeVf8J/M8hqy8CtjX3twEXDzmXJGlAK51DP6Gq9gI0t+uGF0mStBIjf1M0yZYk80nmFxYWRr07SXreWmmh70uyHqC53b/chlW1tapmq2p2ZmZmhbuTJPWz0kLfAcw19+eA7cOJI0laqTYfW7wR+G/gtCR7kmwG3g+cn2Q3cH6zLEmaoL6XoKuqy5d56LwhZ5EkrYJnikpSR1joktQRFrokdUTfOXSNzhFHDP77tKpGkERSF3iELkkdYaFLUkdY6JLUERa6JHWEhS5JHWGhS1JHWOiS1BEWuiR1hCcWSZq4o48+etIRfs4zzzwz6Qgr4hG6JHVE3yP0JCcD1wEnAgeArVX1kSRrgZuAjcBjwFur6qnRRZXUVT/+8Y8nHaET2hyhPwu8u6peCZwFvCPJq4ArgJ1VtQnY2SxLkiakb6FX1d6q+lJz/wfALuAk4CJgW7PZNuDiUYWUJPU30Bx6ko3AGcBdwAlVtRd6pQ+sW+Y5W5LMJ5lfWFhYXVpJ0rJaF3qSY4BPA++qqu+3fV5Vba2q2aqanZmZWUlGSVILrQo9yQvolfkNVXVrs3pfkvXN4+uB/aOJKElqo2+hJwlwNbCrqj646KEdwFxzfw7YPvx4kqS22pxYdA7wh8BXktzXrPsb4P3AzUk2A48Dl44moqTV2LVr16QjaEz6FnpV/ReQZR4+b7hxJEkr5an/Useddtppk46gMfHUf0nqCAtdkjrCQpekjrDQJakjLHRJ6ggLXZI6wkKXpI6w0CWpIyx0SeoIC12SOsJCl6SOsNAlqSMsdEnqiDYXuHhRki8m+XKSryZ5X7P+1CR3Jdmd5KYka0YfV5K0nDZH6P8HnFtVrwFOBy5IchbwAeBDVbUJeArYPLqYkqR++hZ69fxvs/iC5quAc4FbmvXbgItHklCS1Erbi0Qf2Vx+bj9wB/B14OmqerbZZA9w0jLP3ZJkPsn8wsLCMDJLkpbQqtCr6qdVdTqwAXgt8MqlNlvmuVuraraqZmdmZlaeVJL0nAb6lEtVPQ18ATgLODbJwUvYbQCeHG40SdIg2nzKZSbJsc39XwJeB+wC7gQuaTabA7aPKqQkqb82F4leD2xLciS9XwA3V9XtSR4EPpnk74B7gatHmFOS1EffQq+q+4Ezllj/KL35dEnSFPBMUUnqCAtdkjrCQpekjrDQJakj2nzKRdJhbN++fSP5vieeeOJIvq9WziN0SeqIqTtCr1ryLwhoQI7j6jmGw+E4jo9H6JLUERa6JHWEhS5JHWGhS1JHWOiS1BEWuiR1ROtCby5Dd2+S25vlU5PclWR3kpuSrBldTElSP4Mcob+T3oUtDvoA8KGq2gQ8BWweZjBJ0mDaXiR6A/D7wMeb5QDnArc0m2wDLh5FQElSO23PFP0w8B7gJc3yy4Cnq+rZZnkPcNJST0yyBdgCcMopp/TdUe93hZbT9qw7x/G5tRlHx/C5+bM4HMM8k7bNNUXfDOyvqnsWr14q11LPr6qtVTVbVbMzMzMrjClJ6qfNEfo5wIVJ3gS8CHgpvSP2Y5Mc1RylbwCeHF1MSVI/fY/Qq+rKqtpQVRuBy4DPV9XbgTuBS5rN5oDtI0spSeprNZ9Dfy/wV0keoTenfvVwIkmSVmKgP59bVV8AvtDcfxR47fAjSZJWwjNFJakjLHRJ6ggLXZI6wkKXpI6w0CWpIyx0SeoIC12SOsJCl6SOsNAlqSMsdEnqCAtdkjrCQpekjhjoj3Pp8HHgwIGRfN8jjvAYQJpWvjolqSNaHaEneQz4AfBT4Nmqmk2yFrgJ2Ag8Bry1qp4aTUxJUj+DHKH/XlWdXlWzzfIVwM6q2gTsbJYlSROymimXi4Btzf1twMWrjyNJWqm2hV7AfyS5J8mWZt0JVbUXoLldt9QTk2xJMp9kfmFhYfWJJUlLavspl3Oq6skk64A7kjzUdgdVtRXYCjA7O1sryChJaqHVEXpVPdnc7gduo3ct0X1J1gM0t/tHFVKS1F/fQk9ydJKXHLwPvB54ANgBzDWbzQHbRxVSktRfmymXE4Dbkhzc/l+r6rNJ7gZuTrIZeBy4dHQx9XxyzDHHALBmzZoJJ5EOL30LvaoeBV6zxPrvAueNIpQkaXCe+j9iB0+V95T59q655hoA3vKWt0w4iXR4sWUkqSM8Qh+x6667DoC3ve1tE04iqes8QpekjrDQJakjLHRJ6ggLXZI6wkKXpI6w0CWpIyx0SeoIC12SOsJCl6SOsNAlqSMsdEnqiFaFnuTYJLckeSjJriRnJ1mb5I4ku5vb40YdVpK0vLZH6B8BPltVr6D3t9F3AVcAO6tqE7CzWZYkTUibS9C9FPgd4GqAqnqmqp4GLgK2NZttAy4eVUhJUn9tjtBfDiwA1yS5N8nHm2uLnlBVewGa23VLPTnJliTzSeYXFhaGFlyS9PPaFPpRwJnAx6rqDOCHDDC9UlVbq2q2qmZnZmZWGFOS1E+bQt8D7Kmqu5rlW+gV/L4k6wGa2/2jiShJaqNvoVfVt4EnkpzWrDoPeBDYAcw16+aA7SNJKElqpe0l6P4cuCHJGuBR4I/p/TK4Oclm4HHg0tFElCS10arQq+o+YHaJh84bbhxJ0kp5keiOuuSSSyYdQdKYeeq/JHWER+gddeutt046gqQx8whdkjrCQpekjrDQJakjLHRJ6ggLXZI6wkKXpI6w0CWpIyx0SeoITyzS1Ln00tH/nbeqGvk+pHHzCF2SOqLNNUVPS3Lfoq/vJ3lXkrVJ7kiyu7k9bhyBJUlLa3OBi4er6vSqOh34DeBHwG30LkO3s6o2ATsZ4LJ0kqThG3TK5Tzg61X1TeAiYFuzfhtw8TCDSZIGM2ihXwbc2Nw/oar2AjS365Z6QpItSeaTzC8sLKw8qSTpObUu9ObycxcCnxpkB1W1tapmq2p2ZmZm0HySpJYGOUJ/I/ClqtrXLO9Lsh6gud0/7HCSpPYGKfTL+dl0C8AOYK65PwdsH1YoSdLgWp1YlOTFwPnAny5a/X7g5iSbgceBoZwN4gkfw+E4rp5jOByO4/i0KvSq+hHwskPWfZfep14kSVPAM0UlqSMsdEnqCAtdkjrCQpekjrDQJakjLHRJ6ggLXZI6IuP80H+SBeCHwHfGttPhOB4zj8vhmNvM43E4Zobh5P6Vqur7x7DGWugASearanasO10lM4/P4ZjbzONxOGaG8eZ2ykWSOsJCl6SOmEShb53APlfLzONzOOY283gcjplhjLnHPocuSRoNp1wkqSMsdEnqiLEVepILkjyc5JEkV4xrv4NKcnKSO5PsSvLVJO9s1q9NckeS3c3tcZPOeqgkRya5N8ntzfKpSe5qMt/UXBd2aiQ5NsktSR5qxvvsaR/nJH/Z/Fw8kOTGJC+axnFO8okk+5M8sGjdkmObnn9sXpv3JzlzijL/ffPzcX+S25Icu+ixK5vMDyd5w7RkXvTYXyepJMc3yyMf57EUepIjgY/Suy7pq4DLk7xqHPtegWeBd1fVK4GzgHc0Wa8AdlbVJmBnszxt3gnsWrT8AeBDTeangM0TSbW8jwCfrapXAK+hl31qxznJScBfALNV9WrgSOAypnOcrwUuOGTdcmP7RmBT87UF+NiYMh7qWn4x8x3Aq6vq14GvAVcCNK/Jy4Bfa57zT03PjNu1/GJmkpxM7ypvjy9aPfpxrqqRfwFnA59btHwlcOU49j2E7Nubf5iHgfXNuvXAw5POdkjODfRepOcCtwOhd3baUUv9G0z6C3gp8A2aN+YXrZ/acQZOAp4A1tK72tftwBumdZyBjcAD/cYW+Bfg8qW2m3TmQx77A+CG5v7PdQjwOeDsackM3ELvIOUx4PhxjfO4plwOvhAO2tOsm2pJNgJnAHcBJ1TVXoDmdt3kki3pw8B7gAPN8suAp6vq2WZ52sb85cACcE0zTfTxJEczxeNcVd8C/oHeUdde4HvAPUz3OC+23NgeLq/PPwH+vbk/tZmTXAh8q6q+fMhDI888rkLPEuum+vOSSY4BPg28q6q+P+k8zyXJm4H9VXXP4tVLbDpNY34UcCbwsao6g97f+Jma6ZWlNHPOFwGnAr8MHE3vv9GHmqZxbmPaf1ZIchW96dAbDq5aYrOJZ07yYuAq4G+XeniJdUPNPK5C3wOcvGh5A/DkmPY9sCQvoFfmN1TVrc3qfUnWN4+vB/ZPKt8SzgEuTPIY8El60y4fBo5NcvBC4NM25nuAPVV1V7N8C72Cn+Zxfh3wjapaqKqfALcCv8V0j/Niy43tVL8+k8wBbwbeXs1cBdOb+Vfp/cL/cvN63AB8KcmJjCHzuAr9bmBT82mANfTezNgxpn0PJEmAq4FdVfXBRQ/tAOaa+3P05tanQlVdWVUbqmojvbH9fFW9HbgTuKTZbNoyfxt4IslpzarzgAeZ4nGmN9VyVpIXNz8nBzNP7TgfYrmx3QH8UfMpjLOA7x2cmpm0JBcA7wUurKofLXpoB3BZkhcmOZXeG41fnETGxarqK1W1rqo2Nq/HPcCZzc/76Md5jG8cvIneu9RfB66axJsXLXP+Nr3/Bt0P3Nd8vYnenPROYHdzu3bSWZfJ/7vA7c39l9P7IX8E+BTwwknnOyTr6cB8M9b/Bhw37eMMvA94CHgAuB544TSOM3AjvXn+n9Arlc3LjS29qYCPNq/Nr9D7FM+0ZH6E3rzzwdfiPy/a/qom88PAG6cl8yGPP8bP3hQd+Th76r8kdYRnikpSR1joktQRFrokdYSFLkkdYaFLUkdY6JLUERa6JHXE/wNNVoIeY1UMOAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x_mu,c,dxdy,cxcy = vae.model()\n",
    "#print(dxdy)\n",
    "show_batch(x_mu, nrow=4,npadding=4)\n",
    "#print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#synthetic_data_train = vae.generate_synthetic_data(N=500)\n",
    "#synthetic_data_test = vae.generate_synthetic_data(N=10)\n",
    "#save_obj(synthetic_data_train,\"synthetic_data_train_v4\")    \n",
    "#save_obj(synthetic_data_test,\"synthetic_data_test_v4\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the parts of the generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 1, 32, 32])\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAADOCAYAAAAuRfn0AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAECFJREFUeJzt3X+s3XV9x/HnSxCdoimFW+xaWGFpUGcmsBsHY1k2kInOQEnEQAyrW5P+45xuEi0jWWKyPzRbUJY4t0aUsnT8EGVtiNM1FbMsWSq3goiUWkSESqVXActm3Kx974/zbbxe7+Wce+/5db88H8nNOd/v+Z77feXDOa9++Zzzvd9UFZKk5e8low4gSeoPC12SWsJCl6SWsNAlqSUsdElqCQtdklrCQpekllhSoSe5LMn+JI8m2dKvUJKkhctiTyxKcgLwLeBS4CBwH3BNVT3cv3iSpF6duITnvgl4tKoeA0hyO3AFMG+hn3baabVu3bol7FKSXnz27t37g6qa6LbdUgp9DfDkjOWDwG/P3ijJZmAzwJlnnsnU1NQSdilJLz5JvtvLdkuZQ88c635p/qaqtlbVZFVNTkx0/QdGkrRISyn0g8AZM5bXAk8tLY4kabGWUuj3AeuTnJXkJOBqYGd/YkmSFmrRc+hVdTTJnwFfAk4APl1V3+xbMknSgizlQ1Gq6gvAF/qURZK0BJ4pKkktYaFLUkssacplEJK5vg2p43o9s9dxfGG9jOOoL8+4Zs0aAA4dOjTSHPPxtdgf/XydeYQuSS1hoUtSS1joktQSFroktYSFLkktYaFLUktY6JLUEha6JLWEhS5JLWGhS1JLdC30JJ9OcjjJQzPWrUyyK8mB5vaUwcaUJHXTyxH6LcBls9ZtAXZX1Xpgd7MsSRqhroVeVf8BPDNr9RXAtub+NmBDn3NJkhZosXPop1fVIYDmdlX/IkmSFmPgH4om2ZxkKsnU9PT0oHcnSS9aiy30p5OsBmhuD8+3YVVtrarJqpqcmJhY5O4kSd0sttB3Ahub+xuBHf2JI0larF6+tngb8F/AOUkOJtkEfAS4NMkB4NJmWZI0Ql0vQVdV18zz0CV9ziJJWgLPFJWklrDQJaklLHRJagkLXZJawkKXpJaw0CWpJSx0SWoJC12SWsJCl6SWsNAlqSUsdElqCQtdklqi6x/nkqSlWLFixdD3eezYMQCOHDky9H2PkkfoktQSXY/Qk5wB3Aq8BjgGbK2qm5KsBO4A1gGPA++sqmcHF1XScvTMM7OvMT94Bw4cAOCcc84Z+r5HqZcj9KPAB6rqdcAFwHuSvB7YAuyuqvXA7mZZkjQiXQu9qg5V1dea+88D+4A1wBXAtmazbcCGQYWUJHW3oDn0JOuA84A9wOlVdQg6pQ+smuc5m5NMJZmanp5eWlpJ0rx6LvQkJwOfA95fVT1/dFxVW6tqsqomJyYmFpNRktSDnr62mOSldMp8e1V9vln9dJLVVXUoyWrg8KBCSqPwkpf4JTAtL11fsUkC3Azsq6obZzy0E9jY3N8I7Oh/PElSr3o5Qr8IuBb4RpIHmnV/BXwEuDPJJuAJ4KrBRJQk9aJroVfVfwKZ5+FL+htHkrRYThJKUktY6JLUEha6JLWEhS5JLWGhS1JLWOiS1BIWuiS1hIUuSS1hoUtSS1joktQSFroktYSFLkktYaFLUktY6JLUEr1c4OLlSb6a5OtJvpnkw836s5LsSXIgyR1JThp8XEnSfHo5Qv9f4OKqeiNwLnBZkguAjwIfq6r1wLPApsHFlCR107XQq+O/m8WXNj8FXAzc1azfBmwYSEJJUk96mkNPckJz+bnDwC7g28BzVXW02eQgsGae525OMpVkanp6uh+ZJUlz6KnQq+pnVXUusBZ4E/C6uTab57lbq2qyqiYnJiYWn1SS9IIW9C2XqnoO+ApwAbAiyfFrkq4FnupvNEnSQnS9SHSSCeCnVfVckl8B3kznA9F7gXcAtwMbgR2DDCppebr22muHvs/nn39+6PscB10LHVgNbEtyAp0j+jur6p4kDwO3J/kb4H7g5gHmlCR10bXQq+pB4Lw51j9GZz5dkjQGejlCl6RF2759+6gjvGh46r8ktYSFLkktYaFLUktY6JLUEha6JLXE2H3LpWrOvyCgBXIcl84x7A/HcXg8QpeklrDQJaklLHRJagkLXZJawkKXpJaw0CWpJXou9OYydPcnuadZPivJniQHktyR5KTBxZQkdbOQI/T3AftmLH8U+FhVrQeeBTb1M5gkaWF6vUj0WuCPgE81ywEuBu5qNtkGbBhEQElSb3o9U/TjwAeBVzXLpwLPVdXRZvkgsGauJybZDGwGOPPMM7vuqPNvhebT61l3juML62UcHcMX5muxP/p5Jm3XI/QkbwcOV9XemavnyjXX86tqa1VNVtXkxMTEImNKkrrp5Qj9IuDyJG8DXg68ms4R+4okJzZH6WuBpwYXU5LUTdcj9Kq6vqrWVtU64Grgy1X1LuBe4B3NZhuBHQNLKUnqainfQ/8Q8JdJHqUzp35zfyJJkhZjQX8+t6q+Anyluf8Y8Kb+R5IkLYZnikpSS1joktQSFroktYSFLkktYaFLUktY6JLUEha6JLWEhS5JLWGhS1JLLOhMUS0fx44d6+vvu/XWWwF497vf3dffK6l/PEKXpJaw0CWpJSx0SWoJC12SWqKnD0WTPA48D/wMOFpVk0lWAncA64DHgXdW1bODiSlJ6mYhR+h/UFXnVtVks7wF2F1V64HdzbIkaUSWMuVyBbCtub8N2LD0OJKkxeq10Av49yR7k2xu1p1eVYcAmttVcz0xyeYkU0mmpqenl55YkjSnXk8suqiqnkqyCtiV5JFed1BVW4GtAJOTk7WIjJKkHvR0hF5VTzW3h4G76VxL9OkkqwGa28ODCilJ6q5roSd5ZZJXHb8P/CHwELAT2NhsthHYMaiQkqTueplyOR24O8nx7f+lqr6Y5D7gziSbgCeAqwYXU5LUTddCr6rHgDfOsf6HwCWDCCVJWjjPFJWklrDQJaklLHRJagkLXZJawkKXpJaw0CWpJSx0SWoJC12SWsJCl6SWsNAlqSUsdElqCQtdklrCQpeklrDQJakleir0JCuS3JXkkST7klyYZGWSXUkONLenDDqsJGl+vR6h3wR8sapeS+dvo+8DtgC7q2o9sLtZliSNSC+XoHs18HvAzQBV9X9V9RxwBbCt2WwbsGFQISVJ3fVyhH42MA18Jsn9ST7VXFv09Ko6BNDcrprryUk2J5lKMjU9Pd234JKkX9TLNUVPBM4H3ltVe5LcxAKmV6pqK7AVYHJyshaVUlLrnHzyyaOOMK+f/OQnABw9enTESRamlyP0g8DBqtrTLN9Fp+CfTrIaoLk9PJiIkqRe9HKR6O8neTLJOVW1n86FoR9ufjYCH2ludww0qRbkxhtv7Ovvu++++/r6+6QjR46MOsK8rrzySgB27FhetdbLlAvAe4HtSU4CHgP+hM7R/Z1JNgFPAFcNJqIkqRc9FXpVPQBMzvHQJf2NI0larF6P0LXMXHfddaOOIGnIPPVfklrCQpeklrDQJaklLHRJagkLXZJawkKXpJaw0CWpJSx0SWoJC12SWsJCl6SWsNAlqSUsdElqiV6uKXpOkgdm/BxJ8v4kK5PsSnKguT1lGIElSXPrWuhVtb+qzq2qc4HfAn4M3E3nMnS7q2o9sJsFXJZOktR/C51yuQT4dlV9F7gC2Nas3wZs6GcwSdLCLLTQrwZua+6fXlWHAJrbVXM9IcnmJFNJpqanpxefVJL0gnou9Obyc5cDn13IDqpqa1VNVtXkxMTEQvNJknq0kCP0twJfq6qnm+Wnk6wGaG4P9zucJKl3Cyn0a/j5dAvATmBjc38jsLwujy1JLZOq6r5R8grgSeDsqvpRs+5U4E7gTOAJ4KqqeuaFfs/k5GRNTU0tObQkvZgk2VtVk9226+ki0VX1Y+DUWet+SOdbL5KkMeCZopLUEha6JLWEhS5JLWGhS1JLWOiS1BIWuiS1hIUuSS3R04lFfdtZMg38D/CDoe20P07DzMOyHHObeTiWY2boT+5fq6qufwxrqIUOkGSqlzOexomZh2c55jbzcCzHzDDc3E65SFJLWOiS1BKjKPStI9jnUpl5eJZjbjMPx3LMDEPMPfQ5dEnSYDjlIkktYaFLUksMrdCTXJZkf5JHk2wZ1n4XKskZSe5Nsi/JN5O8r1m/MsmuJAea21NGnXW2JCckuT/JPc3yWUn2NJnvaK4LOzaSrEhyV5JHmvG+cNzHOclfNK+Lh5LcluTl4zjOST6d5HCSh2asm3Ns0/H3zXvzwSTnj1Hmv21eHw8muTvJihmPXd9k3p/kLeOSecZj1yWpJKc1ywMf56EUepITgE/QuS7p64Frkrx+GPtehKPAB6rqdcAFwHuarFuA3VW1HtjdLI+b9wH7Zix/FPhYk/lZYNNIUs3vJuCLVfVa4I10so/tOCdZA/w5MFlVbwBOAK5mPMf5FuCyWevmG9u3Auubn83AJ4eUcbZb+OXMu4A3VNVvAt8Crgdo3pNXA7/RPOcfmp4Ztlv45cwkOQO4lM7V3I4b/DhX1cB/gAuBL81Yvh64fhj77kP2Hc1/mP3A6mbdamD/qLPNyrmWzpv0YuAeIHTOTjtxrv8Go/4BXg18h+aD+Rnrx3acgTV0LsW4ks7Vvu4B3jKu4wysAx7qNrbAPwHXzLXdqDPPeuxKYHtz/xc6BPgScOG4ZAbuonOQ8jhw2rDGeVhTLsffCMcdbNaNtSTrgPOAPcDpVXUIoLldNbpkc/o48EHgWLN8KvBcVR1tlsdtzM8GpoHPNNNEn0rySsZ4nKvqe8Df0TnqOgT8CNjLeI/zTPON7XJ5f/4p8G/N/bHNnORy4HtV9fVZDw0887AKPXOsG+vvSyY5Gfgc8P6qOjLqPC8kyduBw1W1d+bqOTYdpzE/ETgf+GRVnUfnb/yMzfTKXJo55yuAs4BfBV5J53+jZxunce7FuL9WSHIDnenQ7cdXzbHZyDMneQVwA/DXcz08x7q+Zh5WoR8EzpixvBZ4akj7XrAkL6VT5tur6vPN6qeTrG4eXw0cHlW+OVwEXJ7kceB2OtMuHwdWJDl+IfBxG/ODwMGq2tMs30Wn4Md5nN8MfKeqpqvqp8Dngd9hvMd5pvnGdqzfn0k2Am8H3lXNXAXjm/nX6fyD//Xm/bgW+FqS1zCEzMMq9PuA9c23AU6i82HGziHte0GSBLgZ2FdVN854aCewsbm/kc7c+lioquuram1VraMztl+uqncB9wLvaDYbt8zfB55Mck6z6hLgYcZ4nOlMtVyQ5BXN6+R45rEd51nmG9udwB8338K4APjR8amZUUtyGfAh4PKq+vGMh3YCVyd5WZKz6HzQ+NVRZJypqr5RVauqal3zfjwInN+83gc/zkP84OBtdD6l/jZwwyg+vOgx5+/S+d+gB4EHmp+30ZmT3g0caG5XjjrrPPl/H7inuX82nRf5o8BngZeNOt+srOcCU81Y/ytwyriPM/Bh4BHgIeCfgZeN4zgDt9GZ5/8pnVLZNN/Y0pkK+ETz3vwGnW/xjEvmR+nMOx9/L/7jjO1vaDLvB946LplnPf44P/9QdODj7Kn/ktQSnikqSS1hoUtSS1joktQSFroktYSFLkktYaFLUktY6JLUEv8PFi6bx46ayvkAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show input images\n",
    "imgs_in=trainloader.load(CHANNELS=1)\n",
    "show_batch(imgs_in[0:8], nrow=4,npadding=4)\n",
    "print(imgs_in.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 1, 32, 32])\n",
      "tensor(1., grad_fn=<SumBackward0>)\n",
      "tensor(0.9597)\n",
      "tensor(1.0581)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAADOCAYAAAAuRfn0AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAEPhJREFUeJzt3W2MXOV5xvH/FWNI84LAsAaDoQZi5aVRA3SVmlKhFJNC0gioFCogpFaL4i9pIG2qBIpUQOqHREV5qZWmtSDBqiiEkFBbKE1qGSLUL5Q1bwEMsUMIODZ4k0BIiFQg3P0wx83i7DKzuzOz48P/J63OnDNnzrn17My1zz4zZ55UFZKk/d/rFroASVJ/GOiS1BIGuiS1hIEuSS1hoEtSSxjoktQSBroktcS8Aj3JWUkeTbIjyWX9KkqSNHuZ64VFSRYB3wPeC+wE7gYuqKqH+1eeJKlXB8zjse8GdlTVYwBJbgLOAWYM9MMPP7xWrFgxj1NK0mvP1q1bf1xVY932m0+gHw08OWV9J/D7++6UZC2wFuDYY49lYmJiHqeUpNeeJD/sZb/5jKFnmm2/MX5TVeuraryqxsfGuv6BkSTN0XwCfSdwzJT15cCu+ZUjSZqr+QT63cDKJMclORA4H9jUn7IkSbM15zH0qnopyV8B3wYWAV+uqof6VpkkaVbm86YoVfVN4Jt9qkWSNA9eKSpJLWGgS1JLzGvIZRCS6T4Nuf868sgj+3q83bt397TfJZdc0tfz9mLdunVDP+dc9XKFdNuei3u9/PLLfTlOr+3T1nbsl35OA2oPXZJaYuR66G2za5cfzZc0HPbQJaklDHRJagkDXZJawkCXpJYw0CWpJQx0SWoJA12SWsJAl6SW6BroSb6cZE+SB6dsW5Jkc5LtzfLQwZYpSeqmlx769cBZ+2y7DNhSVSuBLc26JGkBdQ30qroT+Ok+m88BNjS3NwDn9rkuSdIszXUM/Yiq2g3QLJf2ryRJ0lwM/E3RJGuTTCSZmJycHPTpJOk1a66B/nSSZQDNcs9MO1bV+qoar6rxsbGxOZ5OktTNXAN9E7Cmub0G2NifciRJc9X1+9CT3Ai8Bzg8yU7gSuDTwM1JLgaeAM4bZJGS+uf222/vy3FWr17dl+Oof7oGelVdMMNd/jYlaYQ4Y9Fr2EUXXQTAokWL+nK85557bk6P27BhQ/ed1DdnnHFGX47Tz7kw++3CCy8EYOXKlQtcScc111wDwPPPPz/Q83jpvyS1hIEuSS3hkEtLveUtb+m6z0knnQTAAQf0/jRYt27djPcdfPDBPR9HUv/ZQ5eklrCH3lI7duzous/27dsBWLx4cV+PK2lh2EOXpJYw0CWpJQx0SWoJA12SWsJAl6SWMNAlqSUMdElqCQNdklqia6AnOSbJHUm2JXkoyaXN9iVJNifZ3iwPHXy5kqSZ9NJDfwn4RFW9HVgFfDTJO4DLgC1VtRLY0qxLkhZI10Cvqt1VdU9z++fANuBo4Bxg7xdZbwDOHVSRkqTuZjWGnmQFcBJwF3BEVe2GTugDS2d4zNokE0kmJicn51etJGlGPQd6kjcBXwc+XlU9T01TVeuraryqxsfGxuZSoySpBz1922KSxXTC/Iaq+kaz+ekky6pqd5JlwJ5BFbk/u/baa/t6vI985CN9PZ6k9ujlUy4BrgO2VdVnp9y1CVjT3F4DbOx/eZKkXvXSQz8V+DDw3ST3Ndv+Dvg0cHOSi4EngPMGU+L+be3atX09nj10aWGdcMIJc37ML37xi36X8wpdA72q/hvIDHev7m85kqS5csYiSZqFiy66aNaPeeqppwB4/vnn+13OK3jpvyS1hIEuSS1hoEtSSxjoktQSBroktYSBLkktYaBLUksY6JLUEga6JLWEgS5JLeGl/y21bt26rvu8+OKLACxevHjQ5UhDdf/99wOwc+fOvh/7tNNOm/VjXnjhhb7XMR176JLUEga6JLVE1yGXJK8H7gQOava/paquTHIccBOwBLgH+HBVDef/CvXFpk2bAOjMYSK1x0MPPTSwY995550DO/Z89dJD/1/g9Kp6F3AicFaSVcBngM9V1UrgGeDiwZUpSeqmlwkuCtg7zcbi5qeA04ELm+0bgKuAL/W/RA3Krl27FroESX3U0xh6kkXN9HN7gM3A94Fnq+qlZpedwNEzPHZtkokkE5OTk/2oWZI0jZ4Cvap+VVUnAsuBdwNvn263GR67vqrGq2p8bGxs7pVKkl7VrD7lUlXPAt8BVgGHJNk7ZLMc8P93SVpAXQM9yViSQ5rbvwWcAWwD7gA+2Oy2Btg4qCIlSd31cqXoMmBDkkV0/gDcXFW3JXkYuCnJPwD3AtcNsE5JUhe9fMrlAeCkabY/Rmc8XZI0ArxSVJJawkCXpJYw0CWpJQx0SWoJA12SWsJAl6SWGLkZizrfBab5sh3nzzbsD9txeOyhS1JLGOiS1BIGuiS1hIEuSS1hoEtSSxjoktQSPQd6Mw3dvUlua9aPS3JXku1JvprkwMGVKUnqZjY99EvpTGyx12eAz1XVSuAZ4OJ+FiZJmp1eJ4leDvwJcG2zHuB04JZmlw3AuYMoUJLUm16vFP088Engzc36YcCzVfVSs74TOHq6ByZZC6wFOPbYY7ueqPO3QjPp9ao72/HV9dKOtuGr6/W5uHv3bgCOPPLIQZbTV6973fDeXuznlbS9zCn6AWBPVW2dunm6uqZ7fFWtr6rxqhofGxubY5mSpG566aGfCpyd5P3A64GD6fTYD0lyQNNLXw7sGlyZkqRuuvbQq+ryqlpeVSuA84Hbq+pDwB3AB5vd1gAbB1alJKmr+QwUfQr4myQ76IypX9efkiRJczGrr8+tqu8A32luPwa8u/8lSZLmwitFJaklDHRJagkDXZJawkCXpJYw0CWpJQx0SWoJA12SWsJAl6SWMNAlqSUMdElqCQNdklrCQJekljDQJaklDHRJaomevj43yePAz4FfAS9V1XiSJcBXgRXA48CfVdUzgylTr0WrVq1a6BKk/cpseuh/VFUnVtV4s34ZsKWqVgJbmnVJ0gKZ1QQX+zgHeE9zewOdiS8+Nc96pP935plnLnQJ0n6l1x56Af+VZGuStc22I6pqN0CzXDrdA5OsTTKRZGJycnL+FUuSptVrD/3UqtqVZCmwOckjvZ6gqtYD6wHGx8drDjVKknrQUw+9qnY1yz3ArXTmEn06yTKAZrlnUEVKkrrrGuhJ3pjkzXtvA38MPAhsAtY0u60BNg6qSElSd70MuRwB3Jpk7/7/XlXfSnI3cHOSi4EngPMGV6YkqZuugV5VjwHvmmb7T4DVgyhKkjR7XikqSS1hoEtSSxjoktQSBroktcR8Lv2XpK6OOuqohS7hNcMeuiS1hIEuSS1hoEtSSxjoktQSBroktYSBLkktYaBLUksY6JLUEga6JLVET4Ge5JAktyR5JMm2JKckWZJkc5LtzfLQQRcrSZpZrz30LwDfqqq30flu9G3AZcCWqloJbGnWJUkLpJcp6A4GTgOuA6iqF6rqWeAcYEOz2wbg3EEVKUnqrpce+vHAJPCVJPcmubaZW/SIqtoN0CyXTvfgJGuTTCSZmJyc7FvhkqRX6uXbFg8ATgY+VlV3JfkCsxheqar1wHqA8fHxmlOVGmlXXnnlSBz36quvHkgd0v6ilx76TmBnVd3VrN9CJ+CfTrIMoFnuGUyJkqRedA30qnoKeDLJW5tNq4GHgU3AmmbbGmDjQCqUJPWk1wkuPgbckORA4DHgL+j8Mbg5ycXAE8B5gylRktSLngK9qu4Dxqe5a3V/y5EkzZVXikpSSxjoktQSBroktYSBLkktYaBLUksY6JLUEga6JLWEgS5JLWGgS1JL9HrpvzR0s/32xNnsf9VVV82yGmn02UOXpJYw0CWpJQx0SWqJXuYUfWuS+6b8PJfk40mWJNmcZHuzPHQYBUuSptfLBBePVtWJVXUi8HvAL4Fb6UxDt6WqVgJbmMW0dJKk/pvtkMtq4PtV9UPgHGBDs30DcG4/C5Mkzc5sA/184Mbm9hFVtRugWS6d7gFJ1iaZSDIxOTk590olSa+q50Bvpp87G/jabE5QVeuraryqxsfGxmZbnySpR7Ppob8PuKeqnm7Wn06yDKBZ7ul3cZKk3s0m0C/g18MtAJuANc3tNcDGfhUlSZq9VFX3nZI3AE8Cx1fVz5pthwE3A8cCTwDnVdVPX+044+PjNTExMe+iJem1JMnWqhrvtl9P3+VSVb8EDttn20/ofOpFkjQCvFJUklrCQJekljDQJaklDHRJagkDXZJawkCXpJYw0CWpJXq6sKhvJ0smgeeBHw/tpP1xONY8LPtj3dY8HPtjzdCfun+7qrp+GdZQAx0gyUQvVzyNEmsenv2xbmsejv2xZhhu3Q65SFJLGOiS1BILEejrF+Cc82XNw7M/1m3Nw7E/1gxDrHvoY+iSpMFwyEWSWsJAl6SWGFqgJzkryaNJdiS5bFjnna0kxyS5I8m2JA8lubTZviTJ5iTbm+WhC13rvpIsSnJvktua9eOS3NXU/NVmXtiRkeSQJLckeaRp71NGvZ2T/HXzvHgwyY1JXj+K7Zzky0n2JHlwyrZp2zYd/9S8Nh9IcvII1fyPzfPjgSS3Jjlkyn2XNzU/muTMUal5yn1/m6SSHN6sD7ydhxLoSRYBX6QzL+k7gAuSvGMY556Dl4BPVNXbgVXAR5taLwO2VNVKYEuzPmouBbZNWf8M8Lmm5meAixekqpl9AfhWVb0NeBed2ke2nZMcDVwCjFfVO4FFwPmMZjtfD5y1z7aZ2vZ9wMrmZy3wpSHVuK/r+c2aNwPvrKrfBb4HXA7QvCbPB36necw/NzkzbNfzmzWT5BjgvXRmc9tr8O1cVQP/AU4Bvj1l/XLg8mGcuw+1b2x+MY8Cy5pty4BHF7q2fepcTudFejpwGxA6V6cdMN3vYKF/gIOBH9C8MT9l+8i2M3A0nakYl9CZ7es24MxRbWdgBfBgt7YF/hW4YLr9Frrmfe77U+CG5vYrMgT4NnDKqNQM3EKnk/I4cPiw2nlYQy57Xwh77Wy2jbQkK4CTgLuAI6pqN0CzXLpwlU3r88AngZeb9cOAZ6vqpWZ91Nr8eGAS+EozTHRtkjcywu1cVT8CrqHT69oN/AzYymi381Qzte3+8vr8S+A/m9sjW3OSs4EfVdX9+9w18JqHFeiZZttIf14yyZuArwMfr6rnFrqeV5PkA8Ceqto6dfM0u45Smx8AnAx8qapOovMdPyMzvDKdZsz5HOA44CjgjXT+jd7XKLVzL0b9uUKSK+gMh96wd9M0uy14zUneAFwB/P10d0+zra81DyvQdwLHTFlfDuwa0rlnLcliOmF+Q1V9o9n8dJJlzf3LgD0LVd80TgXOTvI4cBOdYZfPA4ck2TsR+Ki1+U5gZ1Xd1azfQifgR7mdzwB+UFWTVfUi8A3gDxjtdp5qprYd6ddnkjXAB4APVTNWwejWfAKdP/j3N6/H5cA9SY5kCDUPK9DvBlY2nwY4kM6bGZuGdO5ZSRLgOmBbVX12yl2bgDXN7TV0xtZHQlVdXlXLq2oFnba9vao+BNwBfLDZbdRqfgp4Mslbm02rgYcZ4XamM9SyKskbmufJ3ppHtp33MVPbbgL+vPkUxirgZ3uHZhZakrOATwFnV9Uvp9y1CTg/yUFJjqPzRuP/LESNU1XVd6tqaVWtaF6PO4GTm+f74Nt5iG8cvJ/Ou9TfB65YiDcveqzzD+n8G/QAcF/z8346Y9JbgO3NcslC1zpD/e8BbmtuH0/nSb4D+Bpw0ELXt0+tJwITTVv/B3DoqLczcDXwCPAg8G/AQaPYzsCNdMb5X6QTKhfP1LZ0hgK+2Lw2v0vnUzyjUvMOOuPOe1+L/zJl/yuamh8F3jcqNe9z/+P8+k3Rgbezl/5LUkt4pagktYSBLkktYaBLUksY6JLUEga6JLWEgS5JLWGgS1JL/B99HLecapI2DAAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# generate picture according to prior\n",
    "x_mu,c,dxdy,cxcy = vae.model(imgs=None)\n",
    "print(x_mu.shape)\n",
    "show_batch(x_mu,nrow=4,npadding=4)\n",
    "\n",
    "print(torch.sum(c[0,:,0]))\n",
    "print(torch.mean(dxdy[0,:,0]))\n",
    "print(torch.mean(dxdy[0,:,1]))\n",
    "#print(dxdy[0,:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DEBUG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----MODEL ---------\n",
      " Trace Shapes:         \n",
      "  Param Sites:         \n",
      " Sample Sites:         \n",
      "    batch dist   |     \n",
      "         value 8 |     \n",
      "      log_prob   |     \n",
      " latent_z dist 8 | 4 50\n",
      "         value 8 | 4 50\n",
      "      log_prob 8 |     \n",
      "center_bb dist 8 | 4  2\n",
      "         value 8 | 4  2\n",
      "      log_prob 8 |     \n",
      "  size_bb dist 8 | 4  2\n",
      "         value 8 | 4  2\n",
      "      log_prob 8 |     \n",
      "  n_boxes dist   |     \n",
      "         value 4 |     \n",
      "      log_prob   |     \n",
      "      c_0 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "      c_1 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "      c_2 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "      c_3 dist 8 |     \n",
      "         value 8 |     \n",
      "      log_prob 8 |     \n",
      "----GUIDE ---------\n",
      "                                       Trace Shapes:                   \n",
      "                                        Param Sites:                   \n",
      "            inference$$$0.cnn.0.double_conv.0.weight   8   1 3   3     \n",
      "            inference$$$0.cnn.0.double_conv.1.weight             8     \n",
      "              inference$$$0.cnn.0.double_conv.1.bias             8     \n",
      "            inference$$$0.cnn.0.double_conv.3.weight   8   8 3   3     \n",
      "            inference$$$0.cnn.0.double_conv.4.weight             8     \n",
      "              inference$$$0.cnn.0.double_conv.4.bias             8     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.0.weight  16   8 3   3     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.1.weight            16     \n",
      "  inference$$$0.cnn.1.double_conv.double_conv.1.bias            16     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.3.weight  16  16 3   3     \n",
      "inference$$$0.cnn.1.double_conv.double_conv.4.weight            16     \n",
      "  inference$$$0.cnn.1.double_conv.double_conv.4.bias            16     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.0.weight  32  16 3   3     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.1.weight            32     \n",
      "  inference$$$0.cnn.2.double_conv.double_conv.1.bias            32     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.3.weight  32  32 3   3     \n",
      "inference$$$0.cnn.2.double_conv.double_conv.4.weight            32     \n",
      "  inference$$$0.cnn.2.double_conv.double_conv.4.bias            32     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.0.weight  64  32 3   3     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.1.weight            64     \n",
      "  inference$$$0.cnn.3.double_conv.double_conv.1.bias            64     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.3.weight  64  64 3   3     \n",
      "inference$$$0.cnn.3.double_conv.double_conv.4.weight            64     \n",
      "  inference$$$0.cnn.3.double_conv.double_conv.4.bias            64     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.0.weight 128  64 3   3     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.1.weight           128     \n",
      "  inference$$$0.cnn.4.double_conv.double_conv.1.bias           128     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.3.weight 128 128 3   3     \n",
      "inference$$$0.cnn.4.double_conv.double_conv.4.weight           128     \n",
      "  inference$$$0.cnn.4.double_conv.double_conv.4.bias           128     \n",
      "    inference$$$0.cnn.5.compute_zwhere.comp_p.weight   1 128 1   1     \n",
      "      inference$$$0.cnn.5.compute_zwhere.comp_p.bias             1     \n",
      " inference$$$0.cnn.5.compute_zwhere.comp_txty.weight   2 128 1   1     \n",
      "   inference$$$0.cnn.5.compute_zwhere.comp_txty.bias             2     \n",
      " inference$$$0.cnn.5.compute_zwhere.comp_sxsy.weight   2 128 1   1     \n",
      "   inference$$$0.cnn.5.compute_zwhere.comp_sxsy.bias             2     \n",
      "  inference$$$0.cnn.5.compute_zwhat.comp_z_mu.weight  50 128 1   1     \n",
      "    inference$$$0.cnn.5.compute_zwhat.comp_z_mu.bias            50     \n",
      " inference$$$0.cnn.5.compute_zwhat.comp_z_std.weight  50 128 1   1     \n",
      "   inference$$$0.cnn.5.compute_zwhat.comp_z_std.bias            50     \n",
      "                                       Sample Sites:                   \n",
      "                                          batch dist             |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob             |     \n",
      "                                       latent_z dist         8   | 4 50\n",
      "                                               value         8   | 4 50\n",
      "                                            log_prob         8   |     \n",
      "                                      center_bb dist         8   | 4  2\n",
      "                                               value         8   | 4  2\n",
      "                                            log_prob         8   |     \n",
      "                                        size_bb dist         8   | 4  2\n",
      "                                               value         8   | 4  2\n",
      "                                            log_prob         8   |     \n",
      "                                        n_boxes dist             |     \n",
      "                                               value         4   |     \n",
      "                                            log_prob             |     \n",
      "                                            c_0 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n",
      "                                            c_1 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n",
      "                                            c_2 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n",
      "                                            c_3 dist         8   |     \n",
      "                                               value         8   |     \n",
      "                                            log_prob         8   |     \n"
     ]
    }
   ],
   "source": [
    "import pyro.poutine as poutine\n",
    "print(\"----MODEL ---------\")\n",
    "trace = poutine.trace(vae.model).get_trace()\n",
    "trace.compute_log_prob()  # optional, but allows printing of log_prob shapes\n",
    "print(trace.format_shapes())\n",
    "\n",
    "print(\"----GUIDE ---------\")\n",
    "trace = poutine.trace(vae.guide).get_trace()\n",
    "trace.compute_log_prob()  # optional, but allows printing of log_prob shapes\n",
    "print(trace.format_shapes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RUN THE INFERENCE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "#vae\n",
    "# setup the optimizer\n",
    "adam_args = {\"lr\": 1.0e-4}\n",
    "optimizer = Adamax(adam_args)\n",
    "#sgd_params = {\"lr\": 1.0e-4}\n",
    "#optimizer = SGD(sgd_params)\n",
    "\n",
    "# setup the Variational Inference \n",
    "svi = SVI(vae.model, vae.guide, optimizer, loss=TraceEnum_ELBO(max_plate_nesting=1))\n",
    "\n",
    "train_loss_history = list()\n",
    "test_loss_history = list()\n",
    "min_test_loss = 999999"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i=   0 train_loss=22682.77115\n",
      "i=   1 train_loss=17995.15929\n",
      "i=   2 train_loss=18137.59307\n",
      "i=   3 train_loss=20289.50101\n",
      "i=   4 train_loss=14881.73544\n",
      "i=   5 train_loss=16804.86028\n",
      "i=   6 train_loss=13296.22499\n",
      "i=   7 train_loss=13597.25716\n",
      "i=   8 train_loss=14098.80969\n",
      "i=   9 train_loss=14641.83242\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-19-b0e9c253b8fd>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mWRITE_FREQUENCY\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m20\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mNUM_EPOCHS\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m     \u001b[0mtotal_epoch_loss_train\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvae\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_one_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msvi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainloader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdesired_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepoch\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     12\u001b[0m     \u001b[0mtrain_loss_history\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtotal_epoch_loss_train\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"[epoch %03d] average training loss: %.4f\"\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mepoch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtotal_epoch_loss_train\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-6-83c327f3883b>\u001b[0m in \u001b[0;36mtrain_one_epoch\u001b[0;34m(self, svi, train_loader, desired_channels, verbose)\u001b[0m\n\u001b[1;32m    178\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muse_cuda\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    179\u001b[0m                 \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 180\u001b[0;31m             \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msvi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    181\u001b[0m             \u001b[0;32mif\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mverbose\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    182\u001b[0m                 \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"i= %3d train_loss=%.5f\"\u001b[0m \u001b[0;34m%\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/infer/svi.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     97\u001b[0m         \u001b[0;31m# get loss and compute gradients\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     98\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mpoutine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam_only\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mparam_capture\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m             \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_and_grads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mguide\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m         params = set(site[\"value\"].unconstrained()\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/infer/traceenum_elbo.py\u001b[0m in \u001b[0;36mloss_and_grads\u001b[0;34m(self, model, guide, *args, **kwargs)\u001b[0m\n\u001b[1;32m    363\u001b[0m         \"\"\"\n\u001b[1;32m    364\u001b[0m         \u001b[0melbo\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 365\u001b[0;31m         \u001b[0;32mfor\u001b[0m \u001b[0mmodel_trace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mguide_trace\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_traces\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mguide\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    366\u001b[0m             \u001b[0melbo_particle\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_compute_dice_elbo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_trace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mguide_trace\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    367\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mis_identically_zero\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0melbo_particle\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/infer/traceenum_elbo.py\u001b[0m in \u001b[0;36m_get_traces\u001b[0;34m(self, model, guide, *args, **kwargs)\u001b[0m\n\u001b[1;32m    306\u001b[0m             \u001b[0mq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mput\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpoutine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    307\u001b[0m             \u001b[0;32mwhile\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mempty\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 308\u001b[0;31m                 \u001b[0;32myield\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mguide\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    309\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    310\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mguide\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/infer/traceenum_elbo.py\u001b[0m in \u001b[0;36m_get_trace\u001b[0;34m(self, model, guide, *args, **kwargs)\u001b[0m\n\u001b[1;32m    260\u001b[0m         \"\"\"\n\u001b[1;32m    261\u001b[0m         model_trace, guide_trace = get_importance_trace(\n\u001b[0;32m--> 262\u001b[0;31m             \"flat\", self.max_plate_nesting, model, guide, *args, **kwargs)\n\u001b[0m\u001b[1;32m    263\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mis_validation_enabled\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/infer/enum.py\u001b[0m in \u001b[0;36mget_importance_trace\u001b[0;34m(graph_type, max_plate_nesting, model, guide, *args, **kwargs)\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0magainst\u001b[0m \u001b[0mit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     41\u001b[0m     \"\"\"\n\u001b[0;32m---> 42\u001b[0;31m     \u001b[0mguide_trace\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpoutine\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mguide\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgraph_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgraph_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     43\u001b[0m     model_trace = poutine.trace(poutine.replay(model, trace=guide_trace),\n\u001b[1;32m     44\u001b[0m                                 graph_type=graph_type).get_trace(*args, **kwargs)\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/trace_messenger.py\u001b[0m in \u001b[0;36mget_trace\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    167\u001b[0m         \u001b[0mCalls\u001b[0m \u001b[0mthis\u001b[0m \u001b[0mpoutine\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mreturns\u001b[0m \u001b[0mits\u001b[0m \u001b[0mtrace\u001b[0m \u001b[0minstead\u001b[0m \u001b[0mof\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mfunction\u001b[0m\u001b[0;31m'\u001b[0m\u001b[0ms\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    168\u001b[0m         \"\"\"\n\u001b[0;32m--> 169\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    170\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmsngr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/trace_messenger.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    145\u001b[0m                                       args=args, kwargs=kwargs)\n\u001b[1;32m    146\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m                 \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mValueError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m                 \u001b[0mexc_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_value\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraceback\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexc_info\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/handlers.py\u001b[0m in \u001b[0;36m_fn\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    464\u001b[0m                                        escape_fn=functools.partial(escape_fn,\n\u001b[1;32m    465\u001b[0m                                                                    next_trace)))\n\u001b[0;32m--> 466\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mftr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    467\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0mNonlocalExit\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0msite_container\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    468\u001b[0m                     \u001b[0msite_container\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreset_stack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/trace_messenger.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    145\u001b[0m                                       args=args, kwargs=kwargs)\n\u001b[1;32m    146\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m                 \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mValueError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m                 \u001b[0mexc_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_value\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraceback\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexc_info\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/messenger.py\u001b[0m in \u001b[0;36m_wraps\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0m_wraps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0m_wraps\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmsngr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_wraps\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/messenger.py\u001b[0m in \u001b[0;36m_wraps\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0m_wraps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0m_wraps\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmsngr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_wraps\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/pyro/poutine/messenger.py\u001b[0m in \u001b[0;36m_wraps\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0m_wraps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0m_wraps\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmsngr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_wraps\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-6-83c327f3883b>\u001b[0m in \u001b[0;36mguide\u001b[0;34m(self, imgs)\u001b[0m\n\u001b[1;32m    133\u001b[0m         \u001b[0;31m#-- 1. run the inference --#\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    134\u001b[0m         \u001b[0;31m#--------------------------#\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 135\u001b[0;31m         \u001b[0mp\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msxsy_mu\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mkxky_mu\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mz_mu\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mz_std\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minference\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimgs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    136\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    137\u001b[0m         \u001b[0;31m# Transform stuff around for easier sampling\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     90\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 92\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     93\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/DAPI_unsupervised/spacetx-research/Full_enumeration/unet.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, verbose)\u001b[0m\n\u001b[1;32m    251\u001b[0m             \u001b[0;32mif\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mverbose\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    252\u001b[0m                 \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"before level \"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\" shape \"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 253\u001b[0;31m             \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdown\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    254\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    255\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/DAPI_unsupervised/spacetx-research/Full_enumeration/unet.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     50\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax_pool_layer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 52\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdouble_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     53\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mup_block\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mModule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/DAPI_unsupervised/spacetx-research/Full_enumeration/unet.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m         \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdouble_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     29\u001b[0m         \u001b[0;32massert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misnan\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m         \u001b[0;32massert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misnan\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     90\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 92\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     93\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    487\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    488\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    490\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    491\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     74\u001b[0m             \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrunning_mean\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrunning_var\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     75\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrack_running_stats\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 76\u001b[0;31m             exponential_average_factor, self.eps)\n\u001b[0m\u001b[1;32m     77\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     78\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/pyro/lib/python3.6/site-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mbatch_norm\u001b[0;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[1;32m   1621\u001b[0m     return torch.batch_norm(\n\u001b[1;32m   1622\u001b[0m         \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrunning_mean\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrunning_var\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1623\u001b[0;31m         \u001b[0mtraining\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmomentum\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0meps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackends\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcudnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menabled\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1624\u001b[0m     )\n\u001b[1;32m   1625\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Save the parameters beofre starting\n",
    "save_obj(params,\"params_model_SQUARE_v2\")\n",
    "\n",
    "# load saved model\n",
    "#load_model(vae,\"compositional_vae_v2_0\")\n",
    "\n",
    "# training loop\n",
    "TEST_FREQUENCY = 5\n",
    "WRITE_FREQUENCY = 20\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    total_epoch_loss_train = vae.train_one_epoch(svi, trainloader, desired_channels=1, verbose=(epoch == 0))\n",
    "    train_loss_history.append(total_epoch_loss_train)\n",
    "    print(\"[epoch %03d] average training loss: %.4f\" % (epoch, total_epoch_loss_train))\n",
    "\n",
    "    if((epoch % TEST_FREQUENCY) == 0):\n",
    "        # report test diagnostics\n",
    "        total_epoch_loss_test = vae.evaluate_one_epoch(svi, testloader, desired_channels=1, verbose=(epoch == 0))\n",
    "        test_loss_history.append(total_epoch_loss_test)\n",
    "        print(\"[epoch %03d] average test loss: %.4f\" % (epoch, total_epoch_loss_test))\n",
    "        \n",
    "        # Save if you got a new record\n",
    "        if(test_loss_history[-1] < min_test_loss):\n",
    "            print(\"New record!\")\n",
    "            min_test_loss = test_loss_history[-1]\n",
    "            save_model(vae,\"compositional_vae_v3_\"+str(epoch))       \n",
    "            save_obj(test_loss_history,\"test_elbo_history_v3_\"+str(epoch))\n",
    "            save_obj(train_loss_history,\"train_elbo_history_v3_\"+str(epoch))\n",
    "            #save_obj(params,\"params_v1_\",+str(epoch))\n",
    "            \n",
    "    if((epoch % 20) == 0):\n",
    "        save_model(vae,\"compositional_vae_v3_\"+str(epoch))       \n",
    "        save_obj(test_loss_history,\"test_loss_history_v3_\"+str(epoch))\n",
    "        save_obj(train_loss_history,\"train_loss_history_v3_\"+str(epoch))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SAVE EVERYTHING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_model(vae,\"cVAE_square_final_MODEL_v2\") \n",
    "save_obj(params,\"cVAE_square_final_PARAMS_v2\")\n",
    "save_obj(train_loss_history ,\"cVAE_square_final_TRAINLOSS_v2\")\n",
    "save_obj(test_loss_history ,\"cVAE_square_final_TESTLOSS_v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pyro]",
   "language": "python",
   "name": "conda-env-pyro-py"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}