Generative_M2_VAE_debugging.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import time\n",
    "\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import OrderedDict\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import RobustScaler, StandardScaler\n",
    "from sklearn.metrics import confusion_matrix, accuracy_score, f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.inference import Trainer, plot_loss\n",
    "from utils.models import DNN\n",
    "import utils.datasets as d\n",
    "import utils.layers as layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.debugger import set_trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "TMP_X = None\n",
    "TMP_y= None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Potřebné funkce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class One_Hot(object):\n",
    "    def __init__(self, n_classes):\n",
    "        self.n_classes = n_classes\n",
    "        self.class_matrix = torch.diag(torch.ones(n_classes))\n",
    "\n",
    "    def __call__(self, p):\n",
    "        return self.class_matrix[p]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_nll(y_true, mu, sigma, reduction=\"mean\"):\n",
    "    \"\"\"\n",
    "    Negative log likelihood (loss function) of gaussian random variable\n",
    "\n",
    "    : param y_true: \ttarget value\n",
    "    : param mu: \t\tmean of distribution ... size = (batch_size, output_dim)\n",
    "    : param sigma: \t\tstandard deviation of distribution  ... size (batch_size, 1) <– same variance\n",
    "\n",
    "    returns mean loss per sample (not per point)\n",
    "    \"\"\"\n",
    "    dim = mu.shape[1]/2\n",
    "    var = sigma.pow(2).squeeze()\n",
    "\n",
    "    if reduction==\"mean\":\n",
    "        return - (torch.mean(torch.sum((y_true-mu).pow(2), axis=1)/(2*var) + dim*torch.log(var)) + dim*math.log(2*math.pi))\n",
    "    if reduction==\"none\":\n",
    "        return - (torch.sum((y_true-mu).pow(2), axis=1)/(2*var) + dim*torch.log(var) + dim*math.log(2*math.pi))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SS_SVI(nn.Module):\n",
    "    def __init__(self, model, likelihood=\"GaussianNLL\", **kwargs):\n",
    "        super(SS_SVI, self).__init__()\n",
    "        self.model = model\n",
    "        self.one_hot = One_Hot(n_classes=self.model.y_dim)\n",
    "\n",
    "        if likelihood not in [\"BCE\", \"GaussianNLL\", \"MSE\"]:\n",
    "            raise ValueError(\"Unknown likelihood\")\n",
    "        else:\n",
    "            self.likelihood = likelihood\n",
    "        \n",
    "        self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\") if kwargs.get(\"set_device\") == None else kwargs.get(\"set_device\")\n",
    "        \n",
    "        self.encoder_independent_of_y = kwargs.get(\"encoder_independent_of_y\") if kwargs.get(\"encoder_independent_of_y\") != None else True \n",
    "\n",
    "    def reconstruction_loss(self, y_pred, y_true, reduction=\"mean\"):\n",
    "        if self.likelihood == \"BCE\":\n",
    "            reconstruction_loss = nn.BCELoss(reduction='sum')(y_pred, y_true)\n",
    "        elif self.likelihood == \"GaussianNLL\":\n",
    "            assert len(y_pred)==2\n",
    "            mu_out, sigma_out = y_pred\n",
    "            reconstruction_loss = gaussian_nll(y_true=y_true, mu=mu_out, sigma=sigma_out, reduction=reduction)\n",
    "            #self.loss_history[\"MSE\"].append(sample_mse(y_true=y_true, y_pred=mu_out).item())\n",
    "        elif self.loss_fn == \"MSE\":\n",
    "            reconstruction_loss = sample_mse(y_true=y_true, y_pred=y_pred)\n",
    "            #nn.MSELoss(reduction='sum')(y_pred, y_true)\n",
    "        return reconstruction_loss\n",
    "\n",
    "    def forward(self, X, y=None):\n",
    "        supervised = False if y is None else True\n",
    "\n",
    "        #\ty-> z \n",
    "        #\t|   |\n",
    "        #   ->x<-\n",
    "\n",
    "        if supervised:\n",
    "            p_y_pred = self.model.classify(X)\n",
    "        \n",
    "            y_oh = self.one_hot(y).to(self.device)\n",
    "            if self.encoder_independent_of_y:\n",
    "                z, z_mu, z_sigma = self.model.encoder(X) \n",
    "            else:\n",
    "                z, z_mu, z_sigma = self.model.encoder(torch.cat([X,y_oh], axis=1)) \n",
    "\n",
    "            X_hat = self.model.decoder(torch.cat([z, y_oh], axis=1))\n",
    "\n",
    "            # losses\n",
    "            # classification loss\n",
    "            loss_clf = nn.CrossEntropyLoss()(p_y_pred, y)\n",
    "\n",
    "            # log p_y of prior distribution\n",
    "            y_prior = 1/self.model.y_dim *torch.ones_like(y_oh)\n",
    "            log_py = torch.mean(torch.log(y_prior), axis=1).to(self.device)\n",
    "\n",
    "            # K-L divergence\n",
    "            kld = - 0.5 * torch.sum(1 + torch.log(z_sigma.pow(2)) - z_mu.pow(2) - z_sigma.pow(2), axis=1)\n",
    "            # with reduction ->  \n",
    "            #    kld = - 0.5 * torch.mean(torch.sum(1 + torch.log(z_sigma.pow(2)) - z_mu.pow(2) - z_sigma.pow(2), axis=1))\n",
    "            \n",
    "            # log p_x ... reconstruction error\n",
    "            log_px = self.reconstruction_loss(X_hat, X, reduction=\"none\")\n",
    "\n",
    "            # final loss of current flow\n",
    "            likelihood = log_px + log_py - kld \n",
    "\n",
    "            return torch.mean(likelihood), loss_clf # returns scalar losses\n",
    "\n",
    "        else:\n",
    "            X_expanded = torch.cat(self.model.y_dim*[X]).float()\n",
    "\n",
    "            # E[q(y|x)] = sum q(y|x) <- monte carlo improvement <- inaccurate decisions on start of training\n",
    "            y_oh = []\n",
    "            for i in range(self.model.y_dim):\n",
    "                y_oh.append(i*torch.ones(X.shape[0]))\n",
    "            y_oh_expanded = self.one_hot(torch.cat(y_oh,axis=0).long()).to(self.device)\n",
    "\n",
    "            z, z_mu, z_sigma = self.model.encoder(torch.cat([X_expanded,y_oh_expanded], axis=1)) \n",
    "\n",
    "            X_hat = self.model.decoder(torch.cat([z, y_oh_expanded.float()], axis=1))\n",
    "\n",
    "            y_pred = self.model.classify(X)\n",
    "            p_y_pred = F.softmax(y_pred, dim=1)\n",
    "\n",
    "            # losses\n",
    "            kld = - 0.5 * torch.sum(1 + torch.log(z_sigma.pow(2)) - z_mu.pow(2) - z_sigma.pow(2), axis=1)\n",
    "\n",
    "            log_px = self.reconstruction_loss(X_hat, X_expanded, reduction=\"none\")\n",
    "            \n",
    "            \"\"\" podle mě stačí jen log -> that should be same\n",
    "            y_prior = 1/self.model.y_dim *torch.ones_like(y_oh_expanded) #y_prior = 1/self.model.y_dim *torch.ones_like(p_y_pred)\n",
    "            log_py = - nn.CrossEntropyLoss(reduction=\"none\")(y_prior, torch.argmax(y_oh_expanded, axis=1)) #nn.CrossEntropyLoss(y_prior, p_y_pred, reduction=\"none\")\n",
    "            #https://github.com/wohlert/semi-supervised-pytorch/blob/master/semi-supervised/inference/distributions.py\n",
    "            #https://github.com/wohlert/semi-supervised-pytorch/blob/master/semi-supervised/inference/variational.py\n",
    "            \"\"\"\n",
    "            \n",
    "            log_py = torch.log(torch.tensor(1./self.model.y_dim))\n",
    "            \n",
    "            likelihood = log_px + log_py - kld \n",
    "\n",
    "            likelihood = torch.mul(p_y_pred, likelihood.view(self.model.y_dim, X.shape[0]).T - torch.log(p_y_pred+1e-8))\n",
    "\n",
    "            likelihood = torch.sum(likelihood, axis=1)\n",
    "\n",
    "            return torch.mean(likelihood)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>.container { width:80% !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:80% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generative_Model_Trainer(nn.Module):\n",
    "    def __init__(self, model, optimizer, scheduler=None, lr=1e-3, **kwargs):\n",
    "        super(Generative_Model_Trainer, self).__init__()\n",
    "        self.model = model\n",
    "        self.optimizer = optimizer(self.model.parameters(), lr=1e-3)\n",
    "        \n",
    "        self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\") if kwargs.get(\"set_device\")==None else kwargs.get(\"set_device\")\n",
    "        print(\"Using device {}\".format(self.device))\n",
    "        \n",
    "        self.elbo = SS_SVI(self.model, likelihood=\"GaussianNLL\", set_device=self.device)\n",
    "        self.scheduler = scheduler\n",
    "\n",
    "        #optional params\n",
    "        if kwargs.get(\"tensorboard\") == True:\n",
    "            self.tensorboard = True\n",
    "            if kwargs.get(\"model_name\")!= None:\n",
    "                self.tb = SummaryWriter(comment=kwargs.get(\"model_name\"))\n",
    "            else:\n",
    "                self.tb = SummaryWriter()\n",
    "        else: \n",
    "            self.tensorboard = False\n",
    "        \n",
    "        self.elbo.encoder_independent_of_y = kwargs.get(\"encoder_independent_of_y\") if kwargs.get(\"encoder_independent_of_y\") != None else True \n",
    "        \n",
    "        self.model = self.model.to(self.device)\n",
    "        self.elbo = self.elbo.to(self.device)\n",
    "        \n",
    "        self.verbose = kwargs.get(\"verbose\") if kwargs.get(\"verbose\") != None else False\n",
    "        if self.verbose:\n",
    "            print(self)\n",
    "        \n",
    "\n",
    "    def reset_losses(self):\n",
    "        self.loss_history = {\n",
    "                            \"train_total_loss\":0., \n",
    "                            \"train_classifier_loss\":0., \n",
    "                            \"train_supervised_loss\":0., \n",
    "                            \"train_unsupervised_loss\":0.,\n",
    "                            \"validation_total_loss\":0.,\n",
    "                            \"validation_classifier_loss\":0.,\n",
    "                            \"validation_supervised_loss\":0.,\n",
    "                            \"validation_unsupervised_loss\":0.,\n",
    "                            \"validation_accuracy\":0.\n",
    "                             }\n",
    "    \n",
    "    def tensorboard_push_losses(self, epoch, n_train_batches, n_valid_batches):\n",
    "        \"\"\"\n",
    "        function for saving losses to tensorboard\n",
    "        \"\"\"\n",
    "        self.tb.add_scalar(\"Loss/train_total_loss\", self.loss_history[\"train_total_loss\"]/n_train_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/train_classifier_loss\", self.loss_history[\"train_classifier_loss\"]/n_train_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/train_supervised_loss\", self.loss_history[\"train_supervised_loss\"]/n_train_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/train_unsupervised_loss\", self.loss_history[\"train_unsupervised_loss\"]/n_train_batches, epoch)\n",
    "        \n",
    "        self.tb.add_scalar(\"Loss/validation_total_loss\", self.loss_history[\"validation_total_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/validation_classifier_loss\", self.loss_history[\"validation_classifier_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/validation_supervised_loss\", self.loss_history[\"validation_supervised_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/validation_unsupervised_loss\", self.loss_history[\"validation_unsupervised_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Accuracy/validation\". self.loss_history[\"validation_accuracy\"]/n_valid_batches, epoch)\n",
    "\n",
    "\n",
    "    def forward(self, epochs, supervised_dataset, unsupervised_dataset, validation_dataset, batch_size):\n",
    "        if not isinstance(epochs, range):\n",
    "            epochs = range(epochs)\n",
    "        n_epochs = max(epochs)+1\n",
    "\n",
    "        unsupervised = torch.utils.data.DataLoader(\n",
    "                            dataset=torch.tensor(unsupervised_dataset).float(), \n",
    "                            batch_size=batch_size//2, \n",
    "                            shuffle=False, \n",
    "                            sampler=torch.utils.data.RandomSampler(\n",
    "                                unsupervised_dataset, \n",
    "                                replacement=False\n",
    "                                )\n",
    "                            )\n",
    "        supervised = torch.utils.data.DataLoader(\n",
    "                            dataset=supervised_dataset,\n",
    "                            batch_size=batch_size//2, \n",
    "                            shuffle=False, \n",
    "                            sampler=torch.utils.data.RandomSampler(\n",
    "                                supervised_dataset, \n",
    "                                replacement=True,\n",
    "                                num_samples=unsupervised_dataset.shape[0]\n",
    "                                )\n",
    "                            )\n",
    "\n",
    "        validation = torch.utils.data.DataLoader(dataset=validation_dataset, batch_size=batch_size)\n",
    "\n",
    "        for epoch in epochs:\n",
    "            self.model.train()\n",
    "            self.reset_losses()\n",
    "            # ============== training ==============\n",
    "            for i in range(len(unsupervised)):\n",
    "                X_unsup = next(iter(unsupervised))\n",
    "                X_sup, y_sup = next(iter(supervised))\n",
    "\n",
    "                X_unsup = X_unsup.to(self.device)\n",
    "                X_sup = X_sup.to(self.device)\n",
    "                y_sup = y_sup.to(self.device)\n",
    "\n",
    "                # ============== forward ===========\n",
    "                L, CLF = self.elbo(X_sup, y_sup)\n",
    "                L = -L\n",
    "                U = -self.elbo(X_unsup)\n",
    "\n",
    "                alpha = 0.1 * (batch_size//2)*2  # correction -> even numbers\n",
    "                J = L + U + alpha*CLF\n",
    "\n",
    "                # logging losses\n",
    "                self.loss_history[\"train_total_loss\"] += J.detach().item()\n",
    "                self.loss_history[\"train_supervised_loss\"] += L.detach().item()\n",
    "                self.loss_history[\"train_classifier_loss\"] += CLF.detach().item()\n",
    "                self.loss_history[\"train_unsupervised_loss\"] += U.detach().item()\n",
    "                \n",
    "                #print(f\"J: {np.exp(J.detach().item())}, L: {L.detach().item()}, U: {U.detach().item()}, CLF: {CLF.detach().item()}\")\n",
    "\n",
    "                # ============ backward ============\n",
    "                self.optimizer.zero_grad()\n",
    "                J.backward()\n",
    "                self.optimizer.step()\n",
    "\n",
    "            # ============= validation =============\n",
    "            self.model.eval()\n",
    "            with torch.no_grad():\n",
    "                acc = 0\n",
    "                for x, y in validation:\n",
    "                    x_sup = x[:batch_size//2]\n",
    "                    y_sup = y[:batch_size//2]\n",
    "                    x_unsup = x[batch_size//2:]\n",
    "                    \n",
    "                    x_sup = x_sup.to(self.device)\n",
    "                    y_sup = y_sup.to(self.device)\n",
    "                    x_unsup = x_unsup.to(self.device)\n",
    "\n",
    "                    # ============== forward ===========\n",
    "                    L, CLF = self.elbo(x_sup, y_sup)\n",
    "                    L = -L\n",
    "                    U = -self.elbo(x_unsup)\n",
    "\n",
    "                    alpha = 0.1 * (batch_size//2)*2  # correction -> even numbers\n",
    "                    J = L + U + alpha*CLF\n",
    "                    # logging losses\n",
    "                    self.loss_history[\"validation_total_loss\"] += J.detach().item()\n",
    "                    self.loss_history[\"validation_supervised_loss\"] += L.detach().item()\n",
    "                    self.loss_history[\"validation_classifier_loss\"] += CLF.detach().item()\n",
    "                    self.loss_history[\"validation_unsupervised_loss\"] += U.detach().item()\n",
    "                    \n",
    "                    # classification\n",
    "                    x = x.to(self.device)\n",
    "                    y = y.to(self.device)\n",
    "                    y_valid_pred = self.model.classify(x)\n",
    "                    self.loss_history[\"validation_accuracy\"] += accuracy_score(y.cpu().detach(), torch.argmax(y_valid_pred.cpu().detach(), axis=1))\n",
    "            if self.verbose:\n",
    "                print(\"Epoch [{}/{}], average_loss:{:.4f}, validation_loss:{:.4f}, val_accuracy:{:,.4f}\"\\\n",
    "                        .format(epoch+1, n_epochs,self.loss_history[\"train_total_loss\"]/len(unsupervised), self.loss_history[\"validation_total_loss\"]/len(validation), self.loss_history[\"validation_accuracy\"]/len(validation)))\n",
    "                    \n",
    "            if self.tensorboard:\n",
    "                self.tensorboard_push_losses(epoch=epoch, n_train_batches=len(unsupervised), n_valid_batches=len(validation))\n",
    "                \n",
    "            if self.scheduler!=None:\n",
    "                self.scheduler.step()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zadefinovat M2 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "class M2(nn.Module):\n",
    "    def __init__(self, latent_dim, n_classes):\n",
    "        super(M2, self).__init__()\n",
    "        self.y_dim = n_classes\n",
    "        self.encoder = nn.Sequential(\n",
    "                            nn.Linear(in_features=160+4, out_features=400),\n",
    "                            nn.ELU(),\n",
    "                            nn.Linear(in_features=400, out_features=200),\n",
    "                            nn.ELU(),\n",
    "                            layers.VariationalLayer(in_features=200, out_features=latent_dim, return_KL=False)\n",
    "                            )\n",
    "        self.decoder = nn.Sequential(\n",
    "                            nn.Linear(in_features=latent_dim+4, out_features=200),\n",
    "                            nn.ELU(),\n",
    "                            nn.Linear(in_features=200, out_features=400),\n",
    "                            nn.ELU(),\n",
    "                            layers.VariationalDecoderOutput(in_features=400, out_features=160)\n",
    "                            )\n",
    "        self.classify = nn.Sequential(\n",
    "                            nn.Linear(in_features=160, out_features=400),\n",
    "                            nn.ELU(),\n",
    "                            nn.Linear(in_features=400, out_features=200),\n",
    "                            nn.ELU(),\n",
    "                            nn.Linear(in_features=200, out_features=4)\n",
    "                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Připravit data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.vstack((np.load(\"data/sequenced_data_for_VAE_length-160_stride-10_pt1.npy\"),\n",
    "               np.load(\"data/sequenced_data_for_VAE_length-160_stride-10_pt2.npy\")))\n",
    "y = np.load(\"data/sequenced_data_for_VAE_length-160_stride-10_targets.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(83680, 160) (83680,)\n"
     ]
    }
   ],
   "source": [
    "print(X.shape, y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=666)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler = RobustScaler()\n",
    "X_train = scaler.fit_transform(X_train)\n",
    "X_test = scaler.transform(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_unsup, X_sup, y_unsup, y_sup = train_test_split(X_train, y_train, test_size=0.2, random_state=666)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "sup = d.H_alphaSequences(X_sup, y_sup)\n",
    "test = d.H_alphaSequences(X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sup_loader = torch.utils.data.DataLoader(dataset = sup, batch_size=512, shuffle=True)\n",
    "#unsup_loader = torch.utils.data.DataLoader(dataset = X_unsup, batch_size=512, shuffle=True)\n",
    "#test_loader = torch.utils.data.DataLoader(dataset = test, batch_size=512, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsupervised = torch.utils.data.DataLoader(\n",
    "                    dataset=X_unsup, \n",
    "                    batch_size=64, \n",
    "                    shuffle=False, \n",
    "                    sampler=torch.utils.data.RandomSampler(\n",
    "                        X_unsup, \n",
    "                        replacement=False\n",
    "                        )\n",
    "                    )\n",
    "supervised = torch.utils.data.DataLoader(\n",
    "                    dataset=sup,\n",
    "                    batch_size=64, \n",
    "                    shuffle=False, \n",
    "                    sampler=torch.utils.data.RandomSampler(\n",
    "                        sup, \n",
    "                        replacement=True,\n",
    "                        num_samples=X_unsup.shape[0]\n",
    "                        )\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.5261,  0.5227,  0.5327,  ...,  0.4995,  0.5320,  0.5514],\n",
       "        [-0.5435, -0.5480, -0.5401,  ..., -0.1841, -0.1531, -0.1185],\n",
       "        [-0.1430, -0.1262, -0.0768,  ..., -0.0498, -0.0645, -0.0490],\n",
       "        ...,\n",
       "        [ 1.2675,  1.2575,  1.2475,  ...,  1.1055,  1.1334,  1.2037],\n",
       "        [ 0.1804,  0.1748,  0.1696,  ...,  0.1585,  0.1507,  0.1609],\n",
       "        [-0.6210, -0.6033, -0.5980,  ..., -0.0828, -0.0999, -0.1116]],\n",
       "       dtype=torch.float64)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(unsupervised))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[ 0.0599,  0.0448,  0.0088,  ..., -0.0847, -0.0775, -0.0940],\n",
       "         [ 0.4667,  0.5947,  0.6855,  ...,  0.7528,  0.8496,  0.8665],\n",
       "         [ 0.2817,  0.2601,  0.2753,  ...,  0.3532,  0.3645,  0.3327],\n",
       "         ...,\n",
       "         [ 1.9015,  1.9384,  2.0870,  ...,  1.7680,  1.7477,  1.7215],\n",
       "         [-0.3499, -0.3711, -0.3812,  ..., -0.4057, -0.4134, -0.4139],\n",
       "         [-0.2339, -0.2697, -0.2638,  ..., -0.6638, -0.6342, -0.6662]]),\n",
       " tensor([1, 1, 1, 1, 0, 0, 2, 1, 1, 1, 1, 1, 3, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 0,\n",
       "         1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 2, 1, 1, 1,\n",
       "         1, 1, 1, 0, 1, 1, 1, 1, 3, 1, 0, 1, 1, 1, 1, 0])]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(supervised))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# kontrola dimenzí"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = M2(30,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "elbo = SS_SVI(model, likelihood=\"GaussianNLL\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J: 305.3947448730469, L: 140.54623413085938, U: 151.11927795410156, Clf: 1.072594165802002\n",
      "J: 304.8121643066406, L: 129.4605255126953, U: 159.42337036132812, Clf: 1.24439537525177\n",
      "J: 324.5542907714844, L: 157.43711853027344, U: 154.90997314453125, Clf: 0.9536867141723633\n",
      "J: 275.02392578125, L: 149.09475708007812, U: 113.04313659667969, Clf: 1.0067213773727417\n",
      "J: 253.57777404785156, L: 127.16284942626953, U: 112.63513946533203, Clf: 1.0765453577041626\n",
      "J: 258.19879150390625, L: 106.16460418701172, U: 139.84983825683594, Clf: 0.951903760433197\n",
      "J: 204.4355926513672, L: 101.17811584472656, U: 90.52375793457031, Clf: 0.9948216676712036\n",
      "J: 197.70603942871094, L: 91.07245635986328, U: 94.13410949707031, Clf: 0.9765209555625916\n",
      "J: 214.12820434570312, L: 129.99945068359375, U: 73.35533142089844, Clf: 0.8416734933853149\n",
      "J: 164.689697265625, L: 68.43266296386719, U: 85.19239807128906, Clf: 0.8644246459007263\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    X_unsup = next(iter(unsupervised))\n",
    "    X_sup, y_sup = next(iter(supervised))\n",
    "\n",
    "    # ============== forward ===========\n",
    "    L, CLF = elbo(X_sup, y_sup)\n",
    "    L = -L\n",
    "    U = -elbo(X_unsup.float())\n",
    "\n",
    "    alpha = 0.1 * 128\n",
    "    J = L + U + alpha*CLF\n",
    "    # ============= backward==============\n",
    "    optimizer.zero_grad()\n",
    "    J.backward()\n",
    "    optimizer.step()\n",
    "    print(f\"J: {J.item()}, L: {L.item()}, U: {U.item()}, Clf: {CLF.item()}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J: 41.21916580200195, L: 7.773998737335205, U: 22.048633575439453, Clf: 0.8903540968894958\n",
      "J: 140.17601013183594, L: 110.40641021728516, U: 18.319795608520508, Clf: 0.894515335559845\n",
      "J: 52.332054138183594, L: 32.473018646240234, U: 8.90212631225586, Clf: 0.8560085892677307\n",
      "J: 33.08345413208008, L: 0.8606748580932617, U: 20.63958740234375, Clf: 0.9049367904663086\n",
      "J: 48.522491455078125, L: 27.207210540771484, U: 8.677841186523438, Clf: 0.9873000383377075\n",
      "J: 29.506622314453125, L: 12.614885330200195, U: 5.646051406860352, Clf: 0.8785691261291504\n",
      "J: 16.454254150390625, L: 2.20969820022583, U: 2.203416347503662, Clf: 0.9407141208648682\n",
      "J: 13.802461624145508, L: -9.472404479980469, U: 12.043622016906738, Clf: 0.8774409294128418\n",
      "J: 4.527091979980469, L: -21.11874771118164, U: 16.887664794921875, Clf: 0.6842324137687683\n",
      "J: 25.184553146362305, L: 21.503582000732422, U: -8.526070594787598, Clf: 0.9536750912666321\n",
      "J: -7.714491844177246, L: -5.908435344696045, U: -13.083610534667969, Clf: 0.881058931350708\n",
      "J: 58.76394271850586, L: -1.9261574745178223, U: 50.45206832885742, Clf: 0.7998462915420532\n",
      "J: -14.550518035888672, L: -13.709857940673828, U: -12.300596237182617, Clf: 0.8953074812889099\n",
      "J: -11.644522666931152, L: -18.16075897216797, U: -4.8595099449157715, Clf: 0.8887302279472351\n",
      "J: 11.976177215576172, L: 12.227359771728516, U: -10.712787628173828, Clf: 0.8173128962516785\n",
      "J: -21.690879821777344, L: -5.7759928703308105, U: -25.16199493408203, Clf: 0.7224302887916565\n",
      "J: 2.247206687927246, L: -25.542587280273438, U: 18.52142333984375, Clf: 0.7240914106369019\n",
      "J: -0.6647529602050781, L: -7.228933334350586, U: -4.431354522705078, Clf: 0.8590261340141296\n",
      "J: -40.5006103515625, L: -44.0321159362793, U: -6.246212005615234, Clf: 0.7638840079307556\n",
      "J: -36.495174407958984, L: -23.929508209228516, U: -22.55841827392578, Clf: 0.7806838154792786\n",
      "J: -31.654541015625, L: -16.512325286865234, U: -23.792699813842773, Clf: 0.6758188009262085\n",
      "J: -10.591849327087402, L: -5.8349385261535645, U: -15.44448471069336, Clf: 0.8349666595458984\n",
      "J: -47.61458206176758, L: -30.26974105834961, U: -27.618547439575195, Clf: 0.8026334047317505\n",
      "J: -80.93330383300781, L: -50.21210479736328, U: -41.5400276184082, Clf: 0.8452212810516357\n",
      "J: -30.98052215576172, L: -10.974967002868652, U: -31.340808868408203, Clf: 0.885566771030426\n",
      "J: -14.890728950500488, L: -21.87282943725586, U: -2.9441027641296387, Clf: 0.7754846811294556\n",
      "J: -50.25548553466797, L: -29.144672393798828, U: -30.681276321411133, Clf: 0.7476926445960999\n",
      "J: 32.06899642944336, L: 20.06574821472168, U: 0.8650965690612793, Clf: 0.8701679706573486\n",
      "J: -40.01652145385742, L: -27.14531707763672, U: -20.70013999938965, Clf: 0.611635684967041\n",
      "J: -33.71699523925781, L: -26.083492279052734, U: -17.66634178161621, Clf: 0.783815324306488\n",
      "J: -31.710479736328125, L: -23.522815704345703, U: -15.700416564941406, Clf: 0.5869338512420654\n",
      "J: -69.70333099365234, L: -40.62921142578125, U: -41.743431091308594, Clf: 0.9897899031639099\n",
      "J: -39.009544372558594, L: -34.30625915527344, U: -17.453388214111328, Clf: 0.9961016774177551\n",
      "J: -25.158397674560547, L: -27.803203582763672, U: -5.4576520919799805, Clf: 0.6330046653747559\n",
      "J: -42.33679962158203, L: -25.093488693237305, U: -28.669443130493164, Clf: 0.8926665782928467\n",
      "J: -81.0386962890625, L: -47.170310974121094, U: -43.62593460083008, Clf: 0.7623089551925659\n",
      "J: -75.86442565917969, L: -54.75815963745117, U: -30.23941993713379, Clf: 0.7135278582572937\n",
      "J: -45.2384033203125, L: -24.446632385253906, U: -34.42467498779297, Clf: 1.0650707483291626\n",
      "J: -108.89185333251953, L: -55.58842849731445, U: -63.024681091308594, Clf: 0.7594736218452454\n",
      "J: -103.90495300292969, L: -57.58531951904297, U: -53.757930755615234, Clf: 0.5811166763305664\n",
      "J: -84.7451171875, L: -44.11774444580078, U: -49.77458190917969, Clf: 0.7146254777908325\n",
      "J: -111.96482849121094, L: -82.75782012939453, U: -40.49905014038086, Clf: 0.8821906447410583\n",
      "J: -122.64246368408203, L: -54.546504974365234, U: -78.87538146972656, Clf: 0.8421429395675659\n",
      "J: -131.33956909179688, L: -71.87358856201172, U: -68.4898681640625, Clf: 0.704992413520813\n",
      "J: -87.36941528320312, L: -23.667461395263672, U: -71.92108917236328, Clf: 0.642119824886322\n",
      "J: -81.49951171875, L: -54.077239990234375, U: -43.193233489990234, Clf: 1.2321069240570068\n",
      "J: -121.62738800048828, L: -54.58289337158203, U: -74.109375, Clf: 0.5519429445266724\n",
      "J: -139.57562255859375, L: -64.16610717773438, U: -83.51236724853516, Clf: 0.6330345273017883\n",
      "J: -142.1878204345703, L: -84.64818572998047, U: -65.07792663574219, Clf: 0.5889286994934082\n",
      "J: -133.8734130859375, L: -85.62664794921875, U: -60.01945495605469, Clf: 0.9197409749031067\n",
      "J: -98.72590637207031, L: -83.11165618896484, U: -26.76970863342285, Clf: 0.8715201616287231\n",
      "J: -80.92603302001953, L: -42.543678283691406, U: -52.99043273925781, Clf: 1.1412559747695923\n",
      "J: -142.23428344726562, L: -70.23193359375, U: -82.1836929321289, Clf: 0.7954162359237671\n",
      "J: -64.72835540771484, L: -43.739845275878906, U: -34.4726448059082, Clf: 1.0534485578536987\n",
      "J: -194.5736846923828, L: -111.63727569580078, U: -91.75804901123047, Clf: 0.6891907453536987\n",
      "J: -139.69583129882812, L: -71.60334777832031, U: -75.583740234375, Clf: 0.5852546095848083\n",
      "J: -116.60061645507812, L: -64.42929077148438, U: -60.04808044433594, Clf: 0.6153712272644043\n",
      "J: -124.30487060546875, L: -45.67039489746094, U: -90.83512878417969, Clf: 0.9531757235527039\n",
      "J: -150.24891662597656, L: -63.584815979003906, U: -98.81139373779297, Clf: 0.9490075707435608\n",
      "J: -104.7109375, L: -54.47902297973633, U: -63.26618957519531, Clf: 1.0183025598526\n",
      "J: -118.51876068115234, L: -51.87855911254883, U: -79.49453735351562, Clf: 1.0042449235916138\n",
      "J: -196.65042114257812, L: -86.58389282226562, U: -119.48416137695312, Clf: 0.7357522249221802\n",
      "J: -129.93133544921875, L: -97.93620300292969, U: -38.38920211791992, Clf: 0.4995373487472534\n",
      "J: -116.67993927001953, L: -62.6568603515625, U: -62.13630676269531, Clf: 0.6338458061218262\n",
      "J: -132.04039001464844, L: -71.59172821044922, U: -69.11066436767578, Clf: 0.676719069480896\n",
      "J: -94.59947204589844, L: -43.075748443603516, U: -65.70661163330078, Clf: 1.1080385446548462\n",
      "J: -170.76528930664062, L: -87.09695434570312, U: -91.21327209472656, Clf: 0.5894483327865601\n",
      "J: -162.67803955078125, L: -84.97600555419922, U: -90.4130859375, Clf: 0.9930517077445984\n",
      "J: -157.9476318359375, L: -62.87809753417969, U: -101.91097259521484, Clf: 0.5344868302345276\n",
      "J: -69.69927978515625, L: -23.611650466918945, U: -59.51056671142578, Clf: 1.048667073249817\n",
      "J: -91.94278717041016, L: -49.242374420166016, U: -56.54507827758789, Clf: 1.08161461353302\n",
      "J: -129.27706909179688, L: -46.79121017456055, U: -91.60749816894531, Clf: 0.7126286029815674\n",
      "J: -131.17791748046875, L: -74.81503295898438, U: -65.6458740234375, Clf: 0.7252340912818909\n",
      "J: -156.07566833496094, L: -93.00092315673828, U: -71.33080291748047, Clf: 0.6450045704841614\n",
      "J: -208.91952514648438, L: -117.321533203125, U: -101.59896850585938, Clf: 0.7813268303871155\n",
      "J: -172.8784637451172, L: -90.84275817871094, U: -90.18583679199219, Clf: 0.6367285251617432\n",
      "J: -103.42559814453125, L: -56.09288024902344, U: -57.49662399291992, Clf: 0.7940553426742554\n",
      "J: -202.1547088623047, L: -122.4793930053711, U: -87.30399322509766, Clf: 0.5959902405738831\n",
      "J: -165.8910369873047, L: -84.0583267211914, U: -90.37672424316406, Clf: 0.6675021052360535\n",
      "J: -125.16678619384766, L: -52.55272674560547, U: -83.684326171875, Clf: 0.8648653030395508\n",
      "J: -171.3170928955078, L: -96.95887756347656, U: -86.95347595214844, Clf: 0.9840041995048523\n",
      "J: -121.38831329345703, L: -65.51121520996094, U: -63.82140350341797, Clf: 0.6206480264663696\n",
      "J: -191.40638732910156, L: -112.63763427734375, U: -87.99717712402344, Clf: 0.7209703326225281\n",
      "J: -173.1829833984375, L: -99.03590393066406, U: -81.99498748779297, Clf: 0.6131176352500916\n",
      "J: -145.28408813476562, L: -36.3665771484375, U: -120.45468139648438, Clf: 0.9013420343399048\n",
      "J: -148.6740264892578, L: -81.67529296875, U: -78.82814025878906, Clf: 0.9241723418235779\n",
      "J: -187.94078063964844, L: -123.23890686035156, U: -74.19422912597656, Clf: 0.7415897250175476\n",
      "J: -157.35108947753906, L: -73.3741455078125, U: -93.47529602050781, Clf: 0.7420582175254822\n",
      "J: -134.37228393554688, L: -64.59823608398438, U: -83.1159439086914, Clf: 1.0423352718353271\n",
      "J: -176.19456481933594, L: -108.33670806884766, U: -76.90422821044922, Clf: 0.7067477703094482\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J: -182.5281219482422, L: -103.3153076171875, U: -87.80377197265625, Clf: 0.6711688041687012\n",
      "J: -121.03032684326172, L: -62.10957336425781, U: -69.95521545410156, Clf: 0.862067461013794\n",
      "J: -200.43798828125, L: -130.47694396972656, U: -78.63784790039062, Clf: 0.6778758764266968\n",
      "J: -167.4507598876953, L: -104.78825378417969, U: -70.606689453125, Clf: 0.6206389665603638\n",
      "J: -190.52044677734375, L: -118.4954833984375, U: -81.23854064941406, Clf: 0.7198104858398438\n",
      "J: -160.8734588623047, L: -71.12211608886719, U: -97.76535034179688, Clf: 0.6260940432548523\n",
      "J: -222.58456420898438, L: -123.96791076660156, U: -108.58290100097656, Clf: 0.7786133289337158\n",
      "J: -177.65866088867188, L: -86.23480224609375, U: -101.72545623779297, Clf: 0.8048129081726074\n",
      "J: -157.54981994628906, L: -70.93718719482422, U: -99.75155639648438, Clf: 1.0264791250228882\n",
      "J: -171.90257263183594, L: -116.13007354736328, U: -61.68026351928711, Clf: 0.4615439772605896\n"
     ]
    }
   ],
   "source": [
    "for i in range(100):\n",
    "    X_unsup = next(iter(unsupervised))\n",
    "    X_sup, y_sup = next(iter(supervised))\n",
    "\n",
    "    # ============== forward ===========\n",
    "    L, CLF = elbo(X_sup, y_sup)\n",
    "    L = -L\n",
    "    U = -elbo(X_unsup.float())\n",
    "\n",
    "    alpha = 0.1 * 128\n",
    "    J = L + U + alpha*CLF\n",
    "    # ============= backward==============\n",
    "    optimizer.zero_grad()\n",
    "    J.backward()\n",
    "    optimizer.step()\n",
    "    print(f\"J: {J.item()}, L: {L.item()}, U: {U.item()}, Clf: {CLF.item()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(244.8381, grad_fn=<NegBackward>)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.3806, grad_fn=<NllLossBackward>)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CLF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(469.2866, grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TEST\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = M2(30,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device cuda:0\n",
      "Generative_Model_Trainer(\n",
      "  (model): M2(\n",
      "    (encoder): Sequential(\n",
      "      (0): Linear(in_features=164, out_features=400, bias=True)\n",
      "      (1): ELU(alpha=1.0)\n",
      "      (2): Linear(in_features=400, out_features=200, bias=True)\n",
      "      (3): ELU(alpha=1.0)\n",
      "      (4): VariationalLayer(\n",
      "        (mu): Linear(in_features=200, out_features=30, bias=True)\n",
      "        (rho): Linear(in_features=200, out_features=30, bias=True)\n",
      "        (softplus): Softplus(beta=1, threshold=20)\n",
      "      )\n",
      "    )\n",
      "    (decoder): Sequential(\n",
      "      (0): Linear(in_features=34, out_features=200, bias=True)\n",
      "      (1): ELU(alpha=1.0)\n",
      "      (2): Linear(in_features=200, out_features=400, bias=True)\n",
      "      (3): ELU(alpha=1.0)\n",
      "      (4): VariationalDecoderOutput(\n",
      "        (mu): Linear(in_features=400, out_features=160, bias=True)\n",
      "        (rho): Linear(in_features=400, out_features=1, bias=True)\n",
      "        (softplus): Softplus(beta=1, threshold=20)\n",
      "      )\n",
      "    )\n",
      "    (classify): Sequential(\n",
      "      (0): Linear(in_features=160, out_features=400, bias=True)\n",
      "      (1): ELU(alpha=1.0)\n",
      "      (2): Linear(in_features=400, out_features=200, bias=True)\n",
      "      (3): ELU(alpha=1.0)\n",
      "      (4): Linear(in_features=200, out_features=4, bias=True)\n",
      "    )\n",
      "  )\n",
      "  (elbo): SS_SVI(\n",
      "    (model): M2(\n",
      "      (encoder): Sequential(\n",
      "        (0): Linear(in_features=164, out_features=400, bias=True)\n",
      "        (1): ELU(alpha=1.0)\n",
      "        (2): Linear(in_features=400, out_features=200, bias=True)\n",
      "        (3): ELU(alpha=1.0)\n",
      "        (4): VariationalLayer(\n",
      "          (mu): Linear(in_features=200, out_features=30, bias=True)\n",
      "          (rho): Linear(in_features=200, out_features=30, bias=True)\n",
      "          (softplus): Softplus(beta=1, threshold=20)\n",
      "        )\n",
      "      )\n",
      "      (decoder): Sequential(\n",
      "        (0): Linear(in_features=34, out_features=200, bias=True)\n",
      "        (1): ELU(alpha=1.0)\n",
      "        (2): Linear(in_features=200, out_features=400, bias=True)\n",
      "        (3): ELU(alpha=1.0)\n",
      "        (4): VariationalDecoderOutput(\n",
      "          (mu): Linear(in_features=400, out_features=160, bias=True)\n",
      "          (rho): Linear(in_features=400, out_features=1, bias=True)\n",
      "          (softplus): Softplus(beta=1, threshold=20)\n",
      "        )\n",
      "      )\n",
      "      (classify): Sequential(\n",
      "        (0): Linear(in_features=160, out_features=400, bias=True)\n",
      "        (1): ELU(alpha=1.0)\n",
      "        (2): Linear(in_features=400, out_features=200, bias=True)\n",
      "        (3): ELU(alpha=1.0)\n",
      "        (4): Linear(in_features=200, out_features=4, bias=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "gmt = Generative_Model_Trainer(model=model, optimizer=torch.optim.Adam, scheduler=None, lr=1e-3, tensorboard=False, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/20], average_loss:-247.5030, validation_loss:-340.2869, val_accuracy:0.7785\n",
      "Epoch [2/20], average_loss:-375.4871, validation_loss:-390.9652, val_accuracy:0.8009\n",
      "Epoch [3/20], average_loss:-399.0818, validation_loss:-409.1270, val_accuracy:0.8010\n",
      "Epoch [4/20], average_loss:-416.2138, validation_loss:-428.5248, val_accuracy:0.8179\n",
      "Epoch [5/20], average_loss:-435.0976, validation_loss:-404.9739, val_accuracy:0.8142\n",
      "Epoch [6/20], average_loss:-447.3335, validation_loss:-461.6093, val_accuracy:0.8072\n",
      "Epoch [7/20], average_loss:-462.6327, validation_loss:-467.0352, val_accuracy:0.8170\n",
      "Epoch [8/20], average_loss:-467.5099, validation_loss:-474.8092, val_accuracy:0.8212\n",
      "Epoch [9/20], average_loss:-473.7803, validation_loss:-479.2372, val_accuracy:0.8282\n",
      "Epoch [10/20], average_loss:-483.4562, validation_loss:-491.3083, val_accuracy:0.8117\n",
      "Epoch [11/20], average_loss:-492.1682, validation_loss:-495.8349, val_accuracy:0.8322\n",
      "Epoch [12/20], average_loss:-499.3201, validation_loss:-498.4437, val_accuracy:0.8393\n",
      "Epoch [13/20], average_loss:-504.5506, validation_loss:-506.8558, val_accuracy:0.8312\n",
      "Epoch [14/20], average_loss:-509.2443, validation_loss:-514.3692, val_accuracy:0.8537\n",
      "Epoch [15/20], average_loss:-514.3856, validation_loss:-517.9524, val_accuracy:0.8637\n",
      "Epoch [16/20], average_loss:-515.5543, validation_loss:-519.3082, val_accuracy:0.8681\n",
      "Epoch [17/20], average_loss:-521.2818, validation_loss:-521.0224, val_accuracy:0.8603\n",
      "Epoch [18/20], average_loss:-523.7026, validation_loss:-527.2988, val_accuracy:0.8726\n",
      "Epoch [19/20], average_loss:-515.7324, validation_loss:-522.6895, val_accuracy:0.8644\n",
      "Epoch [20/20], average_loss:-527.1089, validation_loss:-526.6531, val_accuracy:0.8681\n"
     ]
    }
   ],
   "source": [
    "gmt(epochs=range(20), supervised_dataset=sup, unsupervised_dataset = X_unsup, validation_dataset=test, batch_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train_total_loss': -441190.11865234375,\n",
       " 'train_classifier_loss': 243.34739600121975,\n",
       " 'train_supervised_loss': -222824.86700439453,\n",
       " 'train_unsupervised_loss': -221480.09805297852,\n",
       " 'validation_total_loss': -68991.56201171875,\n",
       " 'validation_classifier_loss': 47.51626372337341,\n",
       " 'validation_supervised_loss': -34807.221252441406,\n",
       " 'validation_unsupervised_loss': -34792.54878234863,\n",
       " 'validation_accuracy': 113.72395833333333}"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(22026.4648)"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = 1 \\\n",
    "    if temp==True else 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "class One_Hot(nn.Module):\n",
    "    def __init__(self, n_classes):\n",
    "        super(One_Hot, self).__init__()\n",
    "        self.n_classes = n_classes\n",
    "        self.register_buffer(\"class_matrix\", torch.diag(torch.ones(n_classes)))\n",
    "\n",
    "    def forward(self, p):\n",
    "        return self.class_matrix[p]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Gaussian_NLL(nn.Module):\n",
    "    \"\"\"\n",
    "    Negative log likelihood (loss function) of gaussian random variable\n",
    "\n",
    "    : param y_true:     target value\n",
    "    : param y_pred:     tuple(mu, sigma)\n",
    "    : param mu:         mean of distribution ... size = (batch_size, output_dim)\n",
    "    : param sigma:      standard deviation of distribution  ... size (batch_size, 1) <– same variance\n",
    "\n",
    "    returns mean loss per sample (not per point)\n",
    "    \"\"\"\n",
    "    def __init__(self, reduction=\"mean\"):\n",
    "        super(Gaussian_NLL, self).__init__()\n",
    "        self.reduction = reduction\n",
    "        \n",
    "    def forward(self, y_pred, y_true):\n",
    "        assert len(y_pred)==2\n",
    "        mu, sigma = y_pred\n",
    "        \n",
    "        dim = mu.shape[1]/2\n",
    "        var = sigma.pow(2).squeeze()\n",
    "        if self.reduction==\"mean\":\n",
    "            return - (torch.mean(torch.sum((y_true-mu).pow(2), axis=1)/(2*var) + dim*torch.log(var)) + dim*math.log(2*math.pi))\n",
    "        if self.reduction==\"none\":\n",
    "            return - (torch.sum((y_true-mu).pow(2), axis=1)/(2*var) + dim*torch.log(var) + dim*math.log(2*math.pi))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SemiSupervisedGenerativeModel(nn.Module):\n",
    "    def __init__(self, encoder, decoder, classifier, y_dim, likelihood=Gaussian_NLL, include_y=False):\n",
    "        super(SemiSupervisedGenerativeModel, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.classifier = classifier\n",
    "        self.likelihood = likelihood(reduction=\"none\")\n",
    "        self.y_dim = y_dim\n",
    "        self.one_hot = One_Hot(self.y_dim)\n",
    "        self.include_y = include_y\n",
    "        \n",
    "    def elbo(self, X, y=None):\n",
    "        supervised = False if y is None else True      \n",
    "        if supervised:\n",
    "            p_y_pred = self.classifier(X)\n",
    "            # y to one hot encoding\n",
    "            y_oh = self.one_hot(y)\n",
    "            # encoding to latent space\n",
    "            if self.include_y:\n",
    "                z, z_mu, z_sigma = self.encoder(torch.cat([X,y_oh], axis=1)) \n",
    "            else:\n",
    "                z, z_mu, z_sigma = self.encoder(X) \n",
    "            # decoding or reconstructing input\n",
    "            X_hat = self.decoder(torch.cat([z, y_oh], axis=1))\n",
    "            # compute ELBO\n",
    "            # classification loss\n",
    "            loss_clf = nn.CrossEntropyLoss()(p_y_pred, y)\n",
    "            # log p_y of prior distribution\n",
    "            log_py = torch.log(torch.tensor(1./self.y_dim))\n",
    "            # K-L divergence\n",
    "            kld = - 0.5 * torch.sum(1 + torch.log(z_sigma.pow(2)) - z_mu.pow(2) - z_sigma.pow(2), axis=1)\n",
    "            # log p_x ... reconstruction error\n",
    "            log_px = self.likelihood(X_hat, X)\n",
    "            # standard ELBO_xy\n",
    "            ELBO = log_px + log_py - kld \n",
    "            return torch.mean(ELBO), loss_clf # returns scalar losses        \n",
    "        else:\n",
    "            # E[q(y|x)] = sum q(y|x) <- monte carlo improvement <- inaccurate decisions on start of training\n",
    "            X_expanded = torch.cat(self.y_dim*[X]).float()        \n",
    "            y_oh = []\n",
    "            for i in range(self.y_dim):\n",
    "                y_oh.append(i*torch.ones(X.shape[0]))\n",
    "            y_oh_expanded = self.one_hot(torch.cat(y_oh,axis=0).long())\n",
    "            # encoding to latent space\n",
    "            if self.include_y:\n",
    "                z, z_mu, z_sigma = self.encoder(torch.cat([X_expanded,y_oh_expanded], axis=1)) \n",
    "            else:\n",
    "                z, z_mu, z_sigma = self.encoder(X_expanded) \n",
    "            # reconstructing input\n",
    "            X_hat = self.decoder(torch.cat([z, y_oh_expanded.float()], axis=1))\n",
    "            # predicting class of input (during trianing we can predict class after decoding thanks to MC treatment)\n",
    "            y_pred = self.classifier(X)\n",
    "            p_y_pred = F.softmax(y_pred, dim=1)\n",
    "            # compute ELBO\n",
    "            # log p_y of prior distribution\n",
    "            log_py = torch.log(torch.tensor(1./self.y_dim))\n",
    "            # K-L divergence\n",
    "            kld = - 0.5 * torch.sum(1 + torch.log(z_sigma.pow(2)) - z_mu.pow(2) - z_sigma.pow(2), axis=1)\n",
    "            # log p_x ... reconstruction error\n",
    "            log_px = self.likelihood(X_hat, X_expanded)\n",
    "            # standard ELBO_xy\n",
    "            ELBO = log_px + log_py - kld \n",
    "            # E_q(y|x)[-ELBO_xy - log(q(y|x))] = E_q(y|x)[-ELBO_xy] + H(q(y|x)) ... matrix version\n",
    "            ELBO = torch.mul(p_y_pred, ELBO.view(self.y_dim, X.shape[0]).T - torch.log(p_y_pred+1e-8))\n",
    "            # final ELBO_x\n",
    "            ELBO = torch.sum(ELBO, axis=1)\n",
    "            return torch.mean(ELBO)\n",
    "   \n",
    "    def encode(self, X, y):\n",
    "        if self.include_y:\n",
    "            y_oh = self.one_hot(y) if y.shape[1]!=self.y_dim else y\n",
    "            z, z_mu, z_sigma = self.encoder(torch.cat([X,y_oh], axis=1)) \n",
    "        else:\n",
    "            z, z_mu, z_sigma = self.encoder(X)\n",
    "        return z, z_mu, z_sigma\n",
    "    \n",
    "    def decode(self, Z, y):\n",
    "        y_oh = self.one_hot(y) if y.shape[1]!=self.y_dim else y\n",
    "        return self.decoder(torch.cat([Z, y_oh], axis=1))\n",
    "    \n",
    "    def classify(self, X):\n",
    "        return self.classifier(X)\n",
    "\n",
    "    def forward(self, x, y=None):\n",
    "        supervised = False if y is None else True\n",
    "        if not supervised:\n",
    "            y = self.classifier(x)\n",
    "        z, z_mu, z_sigma = self.encode(x,y)\n",
    "        return self.decode(z, y)\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generative_Model_Trainer(nn.Module):\n",
    "    def __init__(self, model, optimizer, scheduler=None, lr=1e-3, **kwargs):\n",
    "        super(Generative_Model_Trainer, self).__init__()\n",
    "        self.model = model\n",
    "        self.optimizer = optimizer(self.model.parameters(), lr=1e-3)\n",
    "        \n",
    "        self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\") if kwargs.get(\"set_device\")==None else kwargs.get(\"set_device\")\n",
    "        print(\"Using device {}\".format(self.device))\n",
    "        \n",
    "        self.scheduler = scheduler\n",
    "\n",
    "        #optional params\n",
    "        if kwargs.get(\"tensorboard\") == True:\n",
    "            self.tensorboard = True\n",
    "            if kwargs.get(\"model_name\")!= None:\n",
    "                self.tb = SummaryWriter(comment=kwargs.get(\"model_name\"))\n",
    "            else:\n",
    "                self.tb = SummaryWriter()\n",
    "        else: \n",
    "            self.tensorboard = False\n",
    "        \n",
    "        self.model = self.model.to(self.device)\n",
    "        \n",
    "        self.verbose = kwargs.get(\"verbose\") if kwargs.get(\"verbose\") != None else False\n",
    "        if self.verbose:\n",
    "            print(self)\n",
    "        \n",
    "\n",
    "    def reset_losses(self):\n",
    "        self.loss_history = {\n",
    "                            \"train_total_loss\":0., \n",
    "                            \"train_classifier_loss\":0., \n",
    "                            \"train_supervised_loss\":0., \n",
    "                            \"train_unsupervised_loss\":0.,\n",
    "                            \"validation_total_loss\":0.,\n",
    "                            \"validation_classifier_loss\":0.,\n",
    "                            \"validation_supervised_loss\":0.,\n",
    "                            \"validation_unsupervised_loss\":0.,\n",
    "                            \"validation_accuracy\":0.\n",
    "                             }\n",
    "    \n",
    "    def tensorboard_push_losses(self, epoch, n_train_batches, n_valid_batches):\n",
    "        \"\"\"\n",
    "        function for saving losses to tensorboard\n",
    "        \"\"\"\n",
    "        self.tb.add_scalar(\"Loss/train_total_loss\", self.loss_history[\"train_total_loss\"]/n_train_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/train_classifier_loss\", self.loss_history[\"train_classifier_loss\"]/n_train_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/train_supervised_loss\", self.loss_history[\"train_supervised_loss\"]/n_train_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/train_unsupervised_loss\", self.loss_history[\"train_unsupervised_loss\"]/n_train_batches, epoch)\n",
    "        \n",
    "        self.tb.add_scalar(\"Loss/validation_total_loss\", self.loss_history[\"validation_total_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/validation_classifier_loss\", self.loss_history[\"validation_classifier_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/validation_supervised_loss\", self.loss_history[\"validation_supervised_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Loss/validation_unsupervised_loss\", self.loss_history[\"validation_unsupervised_loss\"]/n_valid_batches, epoch)\n",
    "        self.tb.add_scalar(\"Accuracy/validation\". self.loss_history[\"validation_accuracy\"]/n_valid_batches, epoch)\n",
    "\n",
    "\n",
    "    def forward(self, epochs, supervised_dataset, unsupervised_dataset, validation_dataset, batch_size):\n",
    "        if not isinstance(epochs, range):\n",
    "            epochs = range(epochs)\n",
    "        n_epochs = max(epochs)+1\n",
    "\n",
    "        unsupervised = torch.utils.data.DataLoader(\n",
    "                            dataset=torch.tensor(unsupervised_dataset).float(), \n",
    "                            batch_size=batch_size//2, \n",
    "                            shuffle=False, \n",
    "                            sampler=torch.utils.data.RandomSampler(\n",
    "                                unsupervised_dataset, \n",
    "                                replacement=False\n",
    "                                )\n",
    "                            )\n",
    "        supervised = torch.utils.data.DataLoader(\n",
    "                            dataset=supervised_dataset,\n",
    "                            batch_size=batch_size//2, \n",
    "                            shuffle=False, \n",
    "                            sampler=torch.utils.data.RandomSampler(\n",
    "                                supervised_dataset, \n",
    "                                replacement=True,\n",
    "                                num_samples=unsupervised_dataset.shape[0]\n",
    "                                )\n",
    "                            )\n",
    "\n",
    "        validation = torch.utils.data.DataLoader(dataset=validation_dataset, batch_size=batch_size)\n",
    "\n",
    "        for epoch in epochs:\n",
    "            self.model.train()\n",
    "            self.reset_losses()\n",
    "            # ============== training ==============\n",
    "            for i in range(len(unsupervised)):\n",
    "                X_unsup = next(iter(unsupervised))\n",
    "                X_sup, y_sup = next(iter(supervised))\n",
    "\n",
    "                X_unsup = X_unsup.to(self.device)\n",
    "                X_sup = X_sup.to(self.device)\n",
    "                y_sup = y_sup.to(self.device)\n",
    "\n",
    "                # ============== forward ===========\n",
    "                L, CLF = self.model.elbo(X_sup, y_sup)\n",
    "                L = -L\n",
    "                U = -self.model.elbo(X_unsup)\n",
    "\n",
    "                alpha = 0.1 * (batch_size//2)*2  # correction -> even numbers\n",
    "                J = L + U + alpha*CLF\n",
    "\n",
    "                # logging losses\n",
    "                self.loss_history[\"train_total_loss\"] += J.detach().item()\n",
    "                self.loss_history[\"train_supervised_loss\"] += L.detach().item()\n",
    "                self.loss_history[\"train_classifier_loss\"] += CLF.detach().item()\n",
    "                self.loss_history[\"train_unsupervised_loss\"] += U.detach().item()\n",
    "                \n",
    "                #print(f\"J: {np.exp(J.detach().item())}, L: {L.detach().item()}, U: {U.detach().item()}, CLF: {CLF.detach().item()}\")\n",
    "\n",
    "                # ============ backward ============\n",
    "                self.optimizer.zero_grad()\n",
    "                J.backward()\n",
    "                self.optimizer.step()\n",
    "\n",
    "            # ============= validation =============\n",
    "            self.model.eval()\n",
    "            with torch.no_grad():\n",
    "                acc = 0\n",
    "                for x, y in validation:\n",
    "                    x_sup = x[:batch_size//2]\n",
    "                    y_sup = y[:batch_size//2]\n",
    "                    x_unsup = x[batch_size//2:]\n",
    "                    \n",
    "                    x_sup = x_sup.to(self.device)\n",
    "                    y_sup = y_sup.to(self.device)\n",
    "                    x_unsup = x_unsup.to(self.device)\n",
    "\n",
    "                    # ============== forward ===========\n",
    "                    L, CLF = self.model.elbo(x_sup, y_sup)\n",
    "                    L = -L\n",
    "                    U = -self.model.elbo(x_unsup)\n",
    "\n",
    "                    alpha = 0.1 * (batch_size//2)*2  # correction -> even numbers\n",
    "                    J = L + U + alpha*CLF\n",
    "                    # logging losses\n",
    "                    self.loss_history[\"validation_total_loss\"] += J.detach().item()\n",
    "                    self.loss_history[\"validation_supervised_loss\"] += L.detach().item()\n",
    "                    self.loss_history[\"validation_classifier_loss\"] += CLF.detach().item()\n",
    "                    self.loss_history[\"validation_unsupervised_loss\"] += U.detach().item()\n",
    "                    \n",
    "                    # classification\n",
    "                    x = x.to(self.device)\n",
    "                    y = y.to(self.device)\n",
    "                    y_valid_pred = self.model.classify(x)\n",
    "                    self.loss_history[\"validation_accuracy\"] += accuracy_score(y.cpu().detach(), torch.argmax(y_valid_pred.cpu().detach(), axis=1))\n",
    "            if self.verbose:\n",
    "                print(\"Epoch [{}/{}], average_loss:{:.4f}, validation_loss:{:.4f}, val_accuracy:{:,.4f}\"\\\n",
    "                        .format(epoch+1, n_epochs,self.loss_history[\"train_total_loss\"]/len(unsupervised), self.loss_history[\"validation_total_loss\"]/len(validation), self.loss_history[\"validation_accuracy\"]/len(validation)))\n",
    "                    \n",
    "            if self.tensorboard:\n",
    "                self.tensorboard_push_losses(epoch=epoch, n_train_batches=len(unsupervised), n_valid_batches=len(validation))\n",
    "                \n",
    "            if self.scheduler!=None:\n",
    "                self.scheduler.step()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SemiSupervisedGenerativeModel(\n",
    "                            encoder=nn.Sequential(\n",
    "                                nn.Linear(in_features=160+4, out_features=400),\n",
    "                                nn.ELU(),\n",
    "                                nn.Linear(in_features=400, out_features=200),\n",
    "                                nn.ELU(),\n",
    "                                layers.VariationalLayer(in_features=200, out_features=30, return_KL=False)), \n",
    "                            decoder=nn.Sequential(\n",
    "                                nn.Linear(in_features=30+4, out_features=200),\n",
    "                                nn.ELU(),\n",
    "                                nn.Linear(in_features=200, out_features=400),\n",
    "                                nn.ELU(),\n",
    "                                layers.VariationalDecoderOutput(in_features=400, out_features=160)\n",
    "                                ),\n",
    "                            classifier=nn.Sequential(\n",
    "                                nn.Linear(in_features=160, out_features=400),\n",
    "                                nn.ELU(),\n",
    "                                nn.Linear(in_features=400, out_features=200),\n",
    "                                nn.ELU(),\n",
    "                                nn.Linear(in_features=200, out_features=4)\n",
    "                                ),\n",
    "                            y_dim=4,\n",
    "                            include_y=True\n",
    "                            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device cuda:0\n",
      "Generative_Model_Trainer(\n",
      "  (model): SemiSupervisedGenerativeModel(\n",
      "    (encoder): Sequential(\n",
      "      (0): Linear(in_features=164, out_features=400, bias=True)\n",
      "      (1): ELU(alpha=1.0)\n",
      "      (2): Linear(in_features=400, out_features=200, bias=True)\n",
      "      (3): ELU(alpha=1.0)\n",
      "      (4): VariationalLayer(\n",
      "        (mu): Linear(in_features=200, out_features=30, bias=True)\n",
      "        (rho): Linear(in_features=200, out_features=30, bias=True)\n",
      "        (softplus): Softplus(beta=1, threshold=20)\n",
      "      )\n",
      "    )\n",
      "    (decoder): Sequential(\n",
      "      (0): Linear(in_features=34, out_features=200, bias=True)\n",
      "      (1): ELU(alpha=1.0)\n",
      "      (2): Linear(in_features=200, out_features=400, bias=True)\n",
      "      (3): ELU(alpha=1.0)\n",
      "      (4): VariationalDecoderOutput(\n",
      "        (mu): Linear(in_features=400, out_features=160, bias=True)\n",
      "        (rho): Linear(in_features=400, out_features=1, bias=True)\n",
      "        (softplus): Softplus(beta=1, threshold=20)\n",
      "      )\n",
      "    )\n",
      "    (classifier): Sequential(\n",
      "      (0): Linear(in_features=160, out_features=400, bias=True)\n",
      "      (1): ELU(alpha=1.0)\n",
      "      (2): Linear(in_features=400, out_features=200, bias=True)\n",
      "      (3): ELU(alpha=1.0)\n",
      "      (4): Linear(in_features=200, out_features=4, bias=True)\n",
      "    )\n",
      "    (likelihood): Gaussian_NLL()\n",
      "    (one_hot): One_Hot()\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "gmt = Generative_Model_Trainer(model=model, optimizer=torch.optim.Adam, scheduler=None, lr=1e-3, tensorboard=False, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10], average_loss:-248.8633, validation_loss:-356.8352, val_accuracy:0.7962\n",
      "Epoch [2/10], average_loss:-374.4327, validation_loss:-401.2551, val_accuracy:0.7962\n",
      "Epoch [3/10], average_loss:-407.1021, validation_loss:-411.1750, val_accuracy:0.7954\n",
      "Epoch [4/10], average_loss:-425.9852, validation_loss:-450.1399, val_accuracy:0.8075\n",
      "Epoch [5/10], average_loss:-446.1871, validation_loss:-439.4182, val_accuracy:0.7980\n",
      "Epoch [6/10], average_loss:-457.0127, validation_loss:-466.2881, val_accuracy:0.8066\n",
      "Epoch [7/10], average_loss:-463.7964, validation_loss:-455.0358, val_accuracy:0.8076\n",
      "Epoch [8/10], average_loss:-472.7484, validation_loss:-476.7056, val_accuracy:0.8187\n",
      "Epoch [9/10], average_loss:-479.4265, validation_loss:-478.5186, val_accuracy:0.8217\n",
      "Epoch [10/10], average_loss:-485.1885, validation_loss:-486.7399, val_accuracy:0.8066\n"
     ]
    }
   ],
   "source": [
    "gmt(epochs=range(10), supervised_dataset=sup, unsupervised_dataset = X_unsup, validation_dataset=test, batch_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'training': False,\n",
       " '_parameters': OrderedDict(),\n",
       " '_buffers': OrderedDict([('class_matrix', tensor([[1., 0., 0., 0.],\n",
       "                       [0., 1., 0., 0.],\n",
       "                       [0., 0., 1., 0.],\n",
       "                       [0., 0., 0., 1.]], device='cuda:0'))]),\n",
       " '_backward_hooks': OrderedDict(),\n",
       " '_forward_hooks': OrderedDict(),\n",
       " '_forward_pre_hooks': OrderedDict(),\n",
       " '_state_dict_hooks': OrderedDict(),\n",
       " '_load_state_dict_pre_hooks': OrderedDict(),\n",
       " '_modules': OrderedDict(),\n",
       " 'n_classes': 4}"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vars(gmt.model.one_hot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}