main_PurifiedVAE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.distributions import kl_divergence, Normal\n",
    "import numpy as np\n",
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "from tensorboardX import SummaryWriter\n",
    "from torch import optim\n",
    "import pypianoroll"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PurifiedVAE and utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class PurifiedVAE(nn.Module):\n",
    "    def __init__(self, input_dims, hidden_dims, z1_dims, n_step, k=1000, num_layers=1, is_bidirection=True, is_training=True, is_cuda=False):\n",
    "        # z1: chroma+root\n",
    "        super(PurifiedVAE, self).__init__()\n",
    "        self.input_dims = input_dims # only chroma\n",
    "        self.hidden_dims = hidden_dims\n",
    "        self.z1_dims = z1_dims\n",
    "        self.n_step = n_step\n",
    "        self.eps = 1\n",
    "        self.sample = None\n",
    "        self.iteration = 0\n",
    "        self.k = torch.FloatTensor([k])\n",
    "\n",
    "        self.training = is_training\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "        self.hidden_factor = (2 if is_bidirection else 1) * num_layers\n",
    "        self.gru_0 = nn.GRU(input_dims, hidden_dims, batch_first=True, bidirectional=is_bidirection)\n",
    "        self.grucell_0 = nn.GRUCell(input_dims + z1_dims, hidden_dims)\n",
    "        self.mu = nn.Linear(hidden_dims * self.hidden_factor, z1_dims)\n",
    "        self.var = nn.Linear(hidden_dims * self.hidden_factor, z1_dims)\n",
    "        self.linear_init_0 = nn.Linear(z1_dims, hidden_dims)\n",
    "        self.linear_out_0 = nn.Linear(hidden_dims, input_dims)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "\n",
    "    def _sampling(self, x):\n",
    "        idx = x.max(1)[1]\n",
    "        x = torch.zeros_like(x)\n",
    "        arange = torch.arange(x.size(0)).long()\n",
    "        if torch.cuda.is_available():\n",
    "            arange = arange.cuda()\n",
    "        x[arange, idx] = 1\n",
    "        return x\n",
    "\n",
    "    def encoder(self, x):\n",
    "        _, x = self.gru_0(x)\n",
    "        x = x.transpose_(0,1).contiguous()\n",
    "        x = x.view(x.size(0), -1)\n",
    "        mean = self.mu(x)\n",
    "        stddev = (self.var(x) * 0.5).exp_()\n",
    "        return Normal(mean, stddev)\n",
    "\n",
    "    def decoder(self, z):\n",
    "        out = torch.zeros((z.size(0), self.input_dims)) # GRUcell's input and output\n",
    "        #print(out.shape)\n",
    "        out[:, -1] = 1 # ?\n",
    "        x = [] # final output\n",
    "        #print(z.size())\n",
    "        t = torch.tanh(self.linear_init_0(z))\n",
    "        hx = t\n",
    "        if self.is_cuda and torch.cuda.is_available():\n",
    "            out = out.cuda()\n",
    "        for i in range(self.n_step):\n",
    "            out = torch.cat([out, z], 1) # batch_size * (input_dims + z1_dims)\n",
    "            #print(out.shape)\n",
    "            hx = self.grucell_0(out, hx)\n",
    "            out = self.sigmoid(self.linear_out_0(hx)) # batch_size * input_dims\n",
    "            x.append(out)\n",
    "            if self.training:\n",
    "                p = torch.rand(1).item()\n",
    "                if p < self.eps:\n",
    "                    out = self.sample[:, i, :]\n",
    "                else:\n",
    "                    out = self._sampling(out)\n",
    "                self.eps = self.k / (self.k + torch.exp(self.iteration / self.k))\n",
    "                self.iteration += 1\n",
    "            else:\n",
    "                out = self._sampling(out)\n",
    "        return torch.stack(x, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.training:\n",
    "            self.sample = x.clone()\n",
    "        latent = self.encoder(x)\n",
    "        if self.training:\n",
    "            z = latent.rsample()\n",
    "        else:\n",
    "            z = latent.mean\n",
    "        return self.decoder(z), latent.mean, latent.stddev\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def std_normal(shape):\n",
    "    N = Normal(torch.zeros(shape), torch.ones(shape))\n",
    "    if torch.cuda.is_available():\n",
    "        N.loc = N.loc.cuda()\n",
    "        N.scale = N.scale.cuda()\n",
    "    return N\n",
    "\n",
    "def loss_function(recon, target_tensor, distribution, beta=.1):\n",
    "    BCE = F.binary_cross_entropy(recon, target_tensor, reduction='elementwise_mean')\n",
    "    normal = std_normal(distribution.mean.size())\n",
    "    KL = kl_divergence(distribution, normal).mean()\n",
    "    return BCE + beta * KL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DataLoader and Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "NUM_WORKERS = 0\n",
    "NUM_EPOCHS = 200\n",
    "lr = 1e-3\n",
    "input_dims = 12\n",
    "hidden_dims = 1024\n",
    "z1_dims = 128\n",
    "num_steps = 32\n",
    "is_cuda = True\n",
    "decay = 0.9999\n",
    "save_path = './ckpt-PVAE/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MinExponentialLR(ExponentialLR):\n",
    "    def __init__(self, optimizer, gamma, minimum, last_epoch=-1):\n",
    "        self.min = minimum\n",
    "        super(MinExponentialLR, self).__init__(optimizer, gamma, last_epoch=-1)\n",
    "\n",
    "    def get_lr(self):\n",
    "        return [\n",
    "            max(base_lr * self.gamma**self.last_epoch, self.min)\n",
    "            for base_lr in self.base_lrs\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = PurifiedVAE(input_dims, hidden_dims, z1_dims, num_steps, is_cuda=is_cuda)\n",
    "if model.is_cuda:\n",
    "    model.cuda()\n",
    "    #model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "bar8_cp_np = np.load('data/bar8_cp_np.npy')\n",
    "bar8_cp_wq_np = np.load('data/bar8_cp_with_quality_np.npy')\n",
    "bar8_comb = []\n",
    "\n",
    "for i in range(6000):\n",
    "    bar8_comb.append([bar8_cp_np[i], bar8_cp_wq_np[i]]) # data, label\n",
    "    \n",
    "for i in range(6000):\n",
    "    bar8_comb.append([bar8_cp_wq_np[i], bar8_cp_wq_np[i]]) # data, label\n",
    "\n",
    "for i in range(6000, bar8_cp_np.shape[0]):\n",
    "    bar8_comb.append([bar8_cp_np[i], bar8_cp_wq_np[i]]) # data, label\n",
    "    \n",
    "for i in range(6000, bar8_cp_np.shape[0]):\n",
    "    bar8_comb.append([bar8_cp_wq_np[i], bar8_cp_wq_np[i]]) # data, label\n",
    "\n",
    "bar8_cp_tensor = torch.tensor(bar8_cp_np, dtype=torch.float32)\n",
    "bar8_cp_wq_tensor = torch.tensor(bar8_cp_wq_np, dtype=torch.float32)\n",
    "bar8_comb_tensor = torch.tensor(bar8_comb, dtype=torch.float32)\n",
    "\n",
    "basic_trainset = bar8_cp_wq_tensor[:6000]\n",
    "basic_testset = bar8_cp_wq_tensor[6000:]\n",
    "# full_trainset = bar8_comb_tensor[:6000]\n",
    "# full_testset = bar8_comb_tensor[12000:12451]\n",
    "comb_trainset = bar8_comb_tensor[:12000]\n",
    "comb_testset = bar8_comb_tensor[12000:]\n",
    "\n",
    "basic_trainloader = torch.utils.data.DataLoader(basic_trainset, batch_size=BATCH_SIZE,\n",
    "                                        shuffle=True, num_workers=NUM_WORKERS)\n",
    "basic_testloader = torch.utils.data.DataLoader(basic_testset, batch_size=BATCH_SIZE,\n",
    "                                        shuffle=False, num_workers=NUM_WORKERS)\n",
    "# full_trainloader = torch.utils.data.DataLoader(full_trainset, batch_size=BATCH_SIZE,\n",
    "#                                         shuffle=True, num_workers=NUM_WORKERS)\n",
    "# full_testloader = torch.utils.data.DataLoader(full_testset, batch_size=BATCH_SIZE,\n",
    "#                                         shuffle=False, num_workers=NUM_WORKERS)\n",
    "comb_trainloader = torch.utils.data.DataLoader(comb_trainset, batch_size=BATCH_SIZE,\n",
    "                                        shuffle=True, num_workers=NUM_WORKERS)\n",
    "comb_testloader = torch.utils.data.DataLoader(comb_testset, batch_size=BATCH_SIZE,\n",
    "                                        shuffle=False, num_workers=NUM_WORKERS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_final(recon):\n",
    "    return recon.apply_(lambda x: 1 if x>=0.5 else 0)\n",
    "\n",
    "def test_recon(sources, targets):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        recons, _, _ = model(sources.cuda())\n",
    "        recons = get_final(recons.cpu())\n",
    "    size = recons.size()\n",
    "    amount = 1\n",
    "    for x in size:\n",
    "        amount *= x\n",
    "    return (torch.norm(recons-targets, 1)/amount).item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr)\n",
    "scheduler = MinExponentialLR(optimizer, gamma=decay, minimum=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = SummaryWriter('log/{}'.format('VanillaVAE'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "cudnn RNN backward can only be called in training mode",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-21-ca836480282b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m         \u001b[0;31m# Backward and optimize\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\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[0m\u001b[1;32m     18\u001b[0m         \u001b[0;31m#torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\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~/.local/lib/python3.7/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    164\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    165\u001b[0m         \"\"\"\n\u001b[0;32m--> 166\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\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    167\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    168\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\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~/.local/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     97\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     98\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: cudnn RNN backward can only be called in training mode"
     ]
    }
   ],
   "source": [
    "lr = 1e-3\n",
    "for epoch in range(NUM_EPOCHS//4):\n",
    "    for i, targets in enumerate(basic_trainloader):  \n",
    "        # Move tensors to the configured device\n",
    "        if model.cuda:\n",
    "            targets = targets.cuda()\n",
    "            \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Forward pass\n",
    "        recons, means, stddevs = model(targets)\n",
    "        distribution = Normal(means, stddevs)\n",
    "        loss = loss_function(recons, targets, distribution)\n",
    "        \n",
    "        # Backward and optimize\n",
    "        \n",
    "        loss.backward()\n",
    "        #torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        \n",
    "        if i % 50 == 0:\n",
    "            print('batch loss: {:.5f}'.format(loss.item()))\n",
    "        writer.add_scalar('batch_loss', loss.item(), i)\n",
    "        \n",
    "        \n",
    "    \n",
    "    if epoch % 10 == 0:\n",
    "        torch.save(model.cpu().state_dict(), save_path+'stage-1-'+str(epoch)+'-epoch_PurifiedVAE.ckpt')\n",
    "        if torch.cuda.is_available():\n",
    "            model.cuda()\n",
    "        print(epoch, '-epoch Model saved!')\n",
    "        print('train: ', test_recon(bar8_cp_tensor[:6000], bar8_cp_tensor[:6000]))\n",
    "        print('test: ', test_recon(bar8_cp_tensor[6000:], bar8_cp_tensor[6000:]))\n",
    "        model.train()\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for PurifiedVAE:\n\tsize mismatch for gru_0.weight_ih_l0: copying a param with shape torch.Size([2400, 12]) from checkpoint, the shape in current model is torch.Size([3072, 12]).\n\tsize mismatch for gru_0.weight_hh_l0: copying a param with shape torch.Size([2400, 800]) from checkpoint, the shape in current model is torch.Size([3072, 1024]).\n\tsize mismatch for gru_0.bias_ih_l0: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for gru_0.bias_hh_l0: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for gru_0.weight_ih_l0_reverse: copying a param with shape torch.Size([2400, 12]) from checkpoint, the shape in current model is torch.Size([3072, 12]).\n\tsize mismatch for gru_0.weight_hh_l0_reverse: copying a param with shape torch.Size([2400, 800]) from checkpoint, the shape in current model is torch.Size([3072, 1024]).\n\tsize mismatch for gru_0.bias_ih_l0_reverse: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for gru_0.bias_hh_l0_reverse: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for grucell_0.weight_ih: copying a param with shape torch.Size([2400, 140]) from checkpoint, the shape in current model is torch.Size([3072, 140]).\n\tsize mismatch for grucell_0.weight_hh: copying a param with shape torch.Size([2400, 800]) from checkpoint, the shape in current model is torch.Size([3072, 1024]).\n\tsize mismatch for grucell_0.bias_ih: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for grucell_0.bias_hh: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for mu.weight: copying a param with shape torch.Size([128, 1600]) from checkpoint, the shape in current model is torch.Size([128, 2048]).\n\tsize mismatch for var.weight: copying a param with shape torch.Size([128, 1600]) from checkpoint, the shape in current model is torch.Size([128, 2048]).\n\tsize mismatch for linear_init_0.weight: copying a param with shape torch.Size([800, 128]) from checkpoint, the shape in current model is torch.Size([1024, 128]).\n\tsize mismatch for linear_init_0.bias: copying a param with shape torch.Size([800]) from checkpoint, the shape in current model is torch.Size([1024]).\n\tsize mismatch for linear_out_0.weight: copying a param with shape torch.Size([12, 800]) from checkpoint, the shape in current model is torch.Size([12, 1024]).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-23-da31093f2db7>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'ckpt/150-epoch_VanillaVAE.ckpt'\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[0m\u001b[1;32m      2\u001b[0m \u001b[0mmodel\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[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\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~/.local/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mload_state_dict\u001b[0;34m(self, state_dict, strict)\u001b[0m\n\u001b[1;32m    837\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merror_msgs\u001b[0m\u001b[0;34m)\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[1;32m    838\u001b[0m             raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n\u001b[0;32m--> 839\u001b[0;31m                                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n\u001b[0m\u001b[1;32m    840\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_IncompatibleKeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmissing_keys\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0munexpected_keys\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    841\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for PurifiedVAE:\n\tsize mismatch for gru_0.weight_ih_l0: copying a param with shape torch.Size([2400, 12]) from checkpoint, the shape in current model is torch.Size([3072, 12]).\n\tsize mismatch for gru_0.weight_hh_l0: copying a param with shape torch.Size([2400, 800]) from checkpoint, the shape in current model is torch.Size([3072, 1024]).\n\tsize mismatch for gru_0.bias_ih_l0: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for gru_0.bias_hh_l0: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for gru_0.weight_ih_l0_reverse: copying a param with shape torch.Size([2400, 12]) from checkpoint, the shape in current model is torch.Size([3072, 12]).\n\tsize mismatch for gru_0.weight_hh_l0_reverse: copying a param with shape torch.Size([2400, 800]) from checkpoint, the shape in current model is torch.Size([3072, 1024]).\n\tsize mismatch for gru_0.bias_ih_l0_reverse: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for gru_0.bias_hh_l0_reverse: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for grucell_0.weight_ih: copying a param with shape torch.Size([2400, 140]) from checkpoint, the shape in current model is torch.Size([3072, 140]).\n\tsize mismatch for grucell_0.weight_hh: copying a param with shape torch.Size([2400, 800]) from checkpoint, the shape in current model is torch.Size([3072, 1024]).\n\tsize mismatch for grucell_0.bias_ih: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for grucell_0.bias_hh: copying a param with shape torch.Size([2400]) from checkpoint, the shape in current model is torch.Size([3072]).\n\tsize mismatch for mu.weight: copying a param with shape torch.Size([128, 1600]) from checkpoint, the shape in current model is torch.Size([128, 2048]).\n\tsize mismatch for var.weight: copying a param with shape torch.Size([128, 1600]) from checkpoint, the shape in current model is torch.Size([128, 2048]).\n\tsize mismatch for linear_init_0.weight: copying a param with shape torch.Size([800, 128]) from checkpoint, the shape in current model is torch.Size([1024, 128]).\n\tsize mismatch for linear_init_0.bias: copying a param with shape torch.Size([800]) from checkpoint, the shape in current model is torch.Size([1024]).\n\tsize mismatch for linear_out_0.weight: copying a param with shape torch.Size([12, 800]) from checkpoint, the shape in current model is torch.Size([12, 1024])."
     ]
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('ckpt/150-epoch_VanillaVAE.ckpt'))\n",
    "model.cuda()\n",
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n, p in model.named_parameters():\n",
    "    if n in {'grucell_0.weight_ih', 'grucell_0.weight_hh', 'grucell_0.bias_ih', 'grucell_0.bias_hh', 'linear_out_0.weight', 'linear_out_0.bias'}:\n",
    "        p.requires_grad = False\n",
    "    else:\n",
    "        #print(n)\n",
    "        p.requires_grad = True\n",
    "\n",
    "\n",
    "optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr)\n",
    "scheduler = MinExponentialLR(optimizer, gamma=decay, minimum=1e-8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = SummaryWriter('log/{}'.format('PurifiedVAE'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch loss: 0.69914\n",
      "batch loss: 0.35754\n",
      "50 -epoch Model saved!\n",
      "test:  0.1983947902917862\n",
      "batch loss: 0.29896\n",
      "batch loss: 0.29860\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-31-04832623be2e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0;31m# Backward and optimize\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 20\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\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[0m\u001b[1;32m     21\u001b[0m         \u001b[0;31m#torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\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~/.local/lib/python3.7/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    164\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    165\u001b[0m         \"\"\"\n\u001b[0;32m--> 166\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\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    167\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    168\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\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~/.local/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     97\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     98\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "lr = 1e-3\n",
    "for epoch in range(NUM_EPOCHS//4, NUM_EPOCHS):\n",
    "    for i, data in enumerate(comb_trainloader):  \n",
    "        # Move tensors to the configured device\n",
    "        sources = data[:,0,:,:]\n",
    "        targets = data[:,1,:,:]\n",
    "        if model.cuda:\n",
    "            sources = sources.cuda()\n",
    "            targets = targets.cuda()\n",
    "            \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Forward pass\n",
    "        recons, means, stddevs = model(sources)\n",
    "        distribution = Normal(means, stddevs)\n",
    "        loss = loss_function(recons, targets, distribution)\n",
    "        \n",
    "        # Backward and optimize\n",
    "        \n",
    "        loss.backward()\n",
    "        #torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        \n",
    "        if i % 50 == 0:\n",
    "            print('batch loss: {:.5f}'.format(loss.item()))\n",
    "        writer.add_scalar('batch_loss', loss.item(), i)\n",
    "        \n",
    "        \n",
    "    \n",
    "    if epoch % 10 == 0:\n",
    "        torch.save(model.cpu().state_dict(), save_path+'stage-2-'+str(epoch)+'-epoch_PurifiedVAE.ckpt')\n",
    "        if torch.cuda.is_available():\n",
    "            model.cuda()\n",
    "        print(epoch, '-epoch Model saved!')\n",
    "        #print('train: ', test_recon(bar8_comb_tensor[:12000,0,:,:], bar8_comb_tensor[:12000,0,:,:]))\n",
    "        print('test: ', test_recon(bar8_comb_tensor[12000:,0,:,:], bar8_comb_tensor[12000:,0,:,:]))\n",
    "        model.train()\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test:  0.4244791567325592\n"
     ]
    }
   ],
   "source": [
    "print('test: ', test_recon(bar8_comb_tensor[12000:12451][0], bar8_comb_tensor[12000:12451][1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('test: ', test_recon(bar8_comb_tensor[0][6000:], bar8_comb_tensor[1][6000:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_final(recon):\n",
    "    return recon.apply_(lambda x: 1 if x>=0.5 else 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([561, 32, 12])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bar8_cp_tensor[6000:].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "215424"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "561*32*12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "#with torch.no_grad():\n",
    "model.eval()\n",
    "model.cuda()\n",
    "with torch.no_grad():\n",
    "    recons, _, _ = model(bar8_cp_tensor[6000:].cuda())\n",
    "    recons = get_final(recons.cpu())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.22070428729057312"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(torch.abs(recons-bar8_cp_wq_tensor[6000:]).sum()/215424).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_recon(targets):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        recons, _, _ = model(targets.cuda())\n",
    "        recons = get_final(recons.cpu())\n",
    "    size = recons.size()\n",
    "    amount = 1\n",
    "    for x in size:\n",
    "        amount *= x\n",
    "    return (torch.norm(recons-targets, 1)/amount).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6561, 32, 12])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recons.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2519424"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "6561*32*12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 1., 0.]])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recons[1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.],\n",
       "        [0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 1., 0.]])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bar8_cp_tensor[1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "test_recon(bar8_cp_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('ckpt-PVAE/stage-2-90-epoch_PurifiedVAE.ckpt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PurifiedVAE(\n",
       "  (gru_0): GRU(12, 1024, batch_first=True, bidirectional=True)\n",
       "  (grucell_0): GRUCell(140, 1024)\n",
       "  (mu): Linear(in_features=2048, out_features=128, bias=True)\n",
       "  (var): Linear(in_features=2048, out_features=128, bias=True)\n",
       "  (linear_init_0): Linear(in_features=128, out_features=1024, bias=True)\n",
       "  (linear_out_0): Linear(in_features=1024, out_features=12, bias=True)\n",
       "  (sigmoid): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.eval()\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "test_recon() missing 1 required positional argument: 'targets'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-7edf8d13a6ec>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtest_recon\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbar8_cp_tensor\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m6000\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[0m",
      "\u001b[0;31mTypeError\u001b[0m: test_recon() missing 1 required positional argument: 'targets'"
     ]
    }
   ],
   "source": [
    "test_recon(bar8_cp_tensor[:6000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.22756749391555786"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('ckpt/190-epoch_VanillaVAE.ckpt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VanillaVAE(\n",
       "  (gru_0): GRU(12, 20, batch_first=True, bidirectional=True)\n",
       "  (grucell_0): GRUCell(32, 20)\n",
       "  (mu): Linear(in_features=40, out_features=20, bias=True)\n",
       "  (var): Linear(in_features=40, out_features=20, bias=True)\n",
       "  (linear_init_0): Linear(in_features=20, out_features=20, bias=True)\n",
       "  (linear_out_0): Linear(in_features=20, out_features=12, bias=True)\n",
       "  (sigmoid): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.eval()\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "test_recon() missing 1 required positional argument: 'targets'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-7edf8d13a6ec>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtest_recon\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbar8_cp_tensor\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m6000\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[0m",
      "\u001b[0;31mTypeError\u001b[0m: test_recon() missing 1 required positional argument: 'targets'"
     ]
    }
   ],
   "source": [
    "test_recon(bar8_cp_tensor[:6000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4.,\n",
       "       4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4.])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bar8_cp_np[0].sum(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.,\n",
       "       3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bar8_cp_wq_np[0].sum(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "only one element tensors can be converted to Python scalars",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-31-cd15507fedbe>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mget_final\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbar8_cp_tensor\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\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[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\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[0m",
      "\u001b[0;31mValueError\u001b[0m: only one element tensors can be converted to Python scalars"
     ]
    }
   ],
   "source": [
    "get_final(model(bar8_cp_tensor[0:1].cuda())[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True],\n",
      "         [True, True, True, True, True, True, True, True, True, True, True,\n",
      "          True]]])\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    recon = model(bar8_cp_tensor[0:1].cuda())[0].cpu()\n",
    "    print(get_final(recon)==bar8_cp_wq_tensor[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon_np = recon.numpy().reshape(-1,12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.zeros([32,60], dtype=np.float32)\n",
    "b = np.zeros([32,56], dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon_pr = np.concatenate((a,recon_np,b), axis=1)\n",
    "in_pr = np.concatenate((a, bar8_cp_np[0], b), axis=1)\n",
    "in_wq_pr = np.concatenate((a, bar8_cp_wq_np[0], b), axis=1)\n",
    "in2_pr = np.concatenate((a, bar8_cp_np[1], b), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon_track = pypianoroll.Track(recon_pr)\n",
    "in_track = pypianoroll.Track(in_pr)\n",
    "in_wq_track = pypianoroll.Track(in_wq_pr)\n",
    "in2_track = pypianoroll.Track(in2_pr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT = pypianoroll.Multitrack(beat_resolution=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [],
   "source": [
    "MT.append_track(track=recon_track, name='recon')\n",
    "MT.append_track(track=in_track, name='in')\n",
    "MT.append_track(track=in_wq_track, name='in_wq')\n",
    "MT.append_track(track=in2_track, name='in2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [],
   "source": [
    "pypianoroll.write(MT,'output/output.mid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}