VAE Architecture Test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.zeros(1, 2, 64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(self, n_params=32):\n",
    "        super(VAE, self).__init__()\n",
    "\n",
    "        # Encoder\n",
    "        self.conv1 = nn.Conv1d(2, 3, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv2 = nn.Conv1d(3, 32, kernel_size=2, stride=2, padding=0)\n",
    "        self.conv3 = nn.Conv1d(32, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv4 = nn.Conv1d(32, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.fc1 = nn.Linear(1024, 128)\n",
    "\n",
    "        # Latent space\n",
    "        self.fc21 = nn.Linear(128, 2)\n",
    "        self.fc22 = nn.Linear(128, 2)\n",
    "\n",
    "        # Decoder\n",
    "        self.fc3 = nn.Linear(2, 128)\n",
    "        self.fc4 = nn.Linear(128, 1024)\n",
    "        self.deconv1 = nn.ConvTranspose1d(32, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.deconv2 = nn.ConvTranspose1d(32, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.deconv3 = nn.ConvTranspose1d(32, 32, kernel_size=2, stride=2, padding=0)\n",
    "        self.conv5 = nn.Conv1d(32, 2, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.relu = nn.ReLU()\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def encode(self, x):\n",
    "        out = self.relu(self.conv1(x))\n",
    "        out = self.relu(self.conv2(out))\n",
    "        out = self.relu(self.conv3(out))\n",
    "        out = self.relu(self.conv4(out))\n",
    "        print(out.size())\n",
    "        out = out.view(out.size(0), -1)\n",
    "        print(out.size())\n",
    "        h1 = self.relu(self.fc1(out))\n",
    "        return self.fc21(h1), self.fc22(h1)\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        if self.training:\n",
    "            std = logvar.mul(0.5).exp_()\n",
    "            eps = Variable(std.data.new(std.size()).normal_())\n",
    "            return eps.mul(std).add_(mu)\n",
    "        else:\n",
    "            return mu\n",
    "\n",
    "    def decode(self, z):\n",
    "        h3 = self.relu(self.fc3(z))\n",
    "        print(h3.size())\n",
    "        out = self.relu(self.fc4(h3))\n",
    "        print(out.size())\n",
    "        # import pdb; pdb.set_trace()\n",
    "        out = out.view(out.size(0), 32, 32)\n",
    "        out = self.relu(self.deconv1(out))\n",
    "        out = self.relu(self.deconv2(out))\n",
    "        out = self.relu(self.deconv3(out))\n",
    "        out = self.sigmoid(self.conv5(out))\n",
    "        return out\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(z), mu, logvar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 32])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 2, 64]) torch.Size([1, 2]) torch.Size([1, 2])\n"
     ]
    }
   ],
   "source": [
    "model = VAE()\n",
    "x_, mu, logvar = model(x)\n",
    "print(x_.size(), mu.size(), logvar.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:dlenv]",
   "language": "python",
   "name": "conda-env-dlenv-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
}