PI_ACVI_MNIST.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mQ1cgc7Y_0bH"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YIJdtVyG1xZU"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gGThtfYc_lKU"
      },
      "outputs": [],
      "source": [
        "import numpy\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from collections import OrderedDict\n",
        "import os\n",
        "import sys\n",
        "import warnings\n",
        "from torch.utils.data import DataLoader\n",
        "import argparse\n",
        "import time\n",
        "import copy\n",
        "import math\n",
        "import torchvision.utils as vision_utils\n",
        "import json\n",
        "import numpy as np\n",
        "from torch.distributions import bernoulli\n",
        "from scipy import linalg\n",
        "import torchvision.datasets as _datasets\n",
        "import torchvision.transforms as _transforms\n",
        "import matplotlib.pyplot as plt\n",
        "import shutil\n",
        "import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KltKnDnTw8sM"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m7rmPgOn_4Yk"
      },
      "source": [
        "# Models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "suSPfCc-_0vc"
      },
      "outputs": [],
      "source": [
        "_NOISE_DIM = 128\n",
        "_H_FILTERS = 64\n",
        "\n",
        "\n",
        "class DiscriminatorCNN28(nn.Module):\n",
        "\n",
        "    def __init__(self, img_channels=1, h_filters=_H_FILTERS,\n",
        "                 spectral_norm=False, img_size=None, n_outputs=1):\n",
        "        if any(not isinstance(_arg, int) for _arg in [img_channels, h_filters, n_outputs]):\n",
        "            raise TypeError(\"Unsupported operand type. Expected integer.\")\n",
        "        if not isinstance(spectral_norm, bool):\n",
        "            raise TypeError(f\"Unsupported operand type: {type(spectral_norm)}. \"\n",
        "                            \"Expected bool.\")\n",
        "        if min([img_channels, h_filters, n_outputs]) <= 0:\n",
        "            raise ValueError(\"Expected nonzero positive input arguments for: the \"\n",
        "                             \"number of output channels, the dimension of the noise \"\n",
        "                             \"vector, as well as the depth of the convolution kernels.\")\n",
        "        super(DiscriminatorCNN28, self).__init__()\n",
        "        # _conv = nn.utils.spectral_norm(nn.Conv2d) if spectral_norm else nn.Conv2d\n",
        "        _apply_sn = lambda x: nn.utils.spectral_norm(x) if spectral_norm else x\n",
        "        self.img_channels = img_channels\n",
        "        self.img_size = img_size\n",
        "        self.n_outputs = n_outputs\n",
        "        self.main = nn.Sequential(\n",
        "            _apply_sn(nn.Conv2d(img_channels, h_filters, 4, 2, 1, bias=False)),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            _apply_sn(nn.Conv2d(h_filters, h_filters * 2, 4, 2, 1, bias=False)),\n",
        "            nn.BatchNorm2d(h_filters * 2),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            _apply_sn(nn.Conv2d(h_filters * 2, h_filters * 4, 4, 2, 1, bias=False)),\n",
        "            nn.BatchNorm2d(h_filters * 4),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            _apply_sn(nn.Conv2d(h_filters * 4, self.n_outputs, 3, 1, 0, bias=False))        \n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        if self.img_channels is not None and self.img_size is not None:\n",
        "            if numpy.prod(list(x.size())) % (self.img_size ** 2 * self.img_channels) != 0:\n",
        "                raise ValueError(f\"Size mismatch. Input size: {numpy.prod(list(x.size()))}. \"\n",
        "                                 f\"Expected input divisible by: {self.noise_dim}\")\n",
        "            x = x.view(-1, self.img_channels, self.img_size, self.img_size)\n",
        "        x = self.main(x)\n",
        "        return x.view(-1, self.n_outputs)\n",
        "\n",
        "    def load(self, model):\n",
        "      self.load_state_dict(model.state_dict())\n",
        "\n",
        "\n",
        "class GeneratorCNN28(nn.Module):\n",
        "\n",
        "    def __init__(self, img_channels=1, noise_dim=_NOISE_DIM, h_filters=_H_FILTERS, out_tanh=False):\n",
        "        if any(not isinstance(_arg, int) for _arg in [img_channels, noise_dim, h_filters]):\n",
        "            raise TypeError(\"Unsupported operand type. Expected integer.\")\n",
        "        if min([img_channels, noise_dim, h_filters]) <= 0:\n",
        "            raise ValueError(\"Expected strictly positive input arguments for the \"\n",
        "                             \"number of output channels, the dimension of the noise \"\n",
        "                             \"vector, as well as the depth of the convolution kernels.\")\n",
        "        super(GeneratorCNN28, self).__init__()\n",
        "        self.noise_dim = noise_dim\n",
        "        self.main = nn.Sequential(\n",
        "            nn.ConvTranspose2d(noise_dim, h_filters * 8, 3, 1, 0, bias=False),\n",
        "            nn.BatchNorm2d(_H_FILTERS * 8),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.ConvTranspose2d(h_filters * 8, h_filters * 4, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(_H_FILTERS * 4),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.ConvTranspose2d(h_filters * 4, h_filters * 2, 4, 2, 0, bias=False),\n",
        "            nn.BatchNorm2d(_H_FILTERS * 2),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.ConvTranspose2d(h_filters * 2, img_channels, 4, 2, 1, bias=False),\n",
        "            nn.Tanh() if out_tanh else nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        if numpy.prod(list(x.size())) % self.noise_dim != 0:\n",
        "            raise ValueError(f\"Size mismatch. Input size: {numpy.prod(list(x.size()))}. \"\n",
        "                             f\"Expected input divisible by: {self.noise_dim}\")\n",
        "        x = x.view(-1, self.noise_dim, 1, 1)\n",
        "        x = self.main(x)\n",
        "        return x\n",
        "\n",
        "    def load(self, model):\n",
        "      self.load_state_dict(model.state_dict())\n",
        "\n",
        "\n",
        "class MLP_mnist(nn.Module):\n",
        "  def __init__(self, input_dims, n_hiddens, n_class):\n",
        "    super(MLP_mnist, self).__init__()\n",
        "    assert isinstance(input_dims, int), 'Expected int for input_dims'\n",
        "    self.input_dims = input_dims\n",
        "    current_dims = input_dims\n",
        "    layers = OrderedDict()\n",
        "\n",
        "    if isinstance(n_hiddens, int):\n",
        "      n_hiddens = [n_hiddens]\n",
        "    else:\n",
        "      n_hiddens = list(n_hiddens)\n",
        "    for i, n_hidden in enumerate(n_hiddens):\n",
        "      layers['fc{}'.format(i+1)] = nn.Linear(current_dims, n_hidden)\n",
        "      layers['relu{}'.format(i+1)] = nn.ReLU()\n",
        "      layers['drop{}'.format(i+1)] = nn.Dropout(0.2)\n",
        "      current_dims = n_hidden\n",
        "    layers['out'] = nn.Linear(current_dims, n_class)\n",
        "    self.layers = layers\n",
        "    self.model= nn.Sequential(layers)\n",
        "    #print(self.model)\n",
        "\n",
        "  def forward(self, input):\n",
        "    input = input.view(input.size(0), -1)\n",
        "    assert input.size(1) == self.input_dims\n",
        "    return self.model.forward(input)\n",
        "\n",
        "  def get_logits_and_fc2_outputs(self, x):\n",
        "    x = x.view(x.size(0), -1)\n",
        "    assert x.size(1) == self.input_dims\n",
        "    fc2_out = None\n",
        "    for l in self.model:\n",
        "      x = l(x)\n",
        "      if l == self.layers[\"fc2\"]:\n",
        "        fc2_out = x\n",
        "    return x, fc2_out\n",
        "\n",
        "\n",
        "def pretrained_mnist_model(input_dims=784, n_hiddens=[256, 256], n_class=10, \n",
        "                           pretrained=None):\n",
        "    model = MLP_mnist(input_dims, n_hiddens, n_class)\n",
        "    if pretrained is not None:\n",
        "        if os.path.exists(pretrained):\n",
        "            print('Loading trained model from %s' % pretrained)\n",
        "            state_dict = torch.load(pretrained,\n",
        "                    map_location='cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "            if 'parallel' in pretrained:\n",
        "                new_state_dict = OrderedDict()\n",
        "                for k, v in state_dict.items():\n",
        "                    name = k[7:]  # remove `module.`\n",
        "                    new_state_dict[name] = v\n",
        "                state_dict = new_state_dict\n",
        "        else:\n",
        "            raise FileNotFoundError(f\"Could not find pretrained model: {pretrained}.\")\n",
        "        model.load_state_dict(state_dict)\n",
        "    if torch.cuda.is_available():\n",
        "        model = model.cuda()\n",
        "    return model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pDYOE4e8DRAY"
      },
      "source": [
        "# Dataloader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ylrrJcSnDQOP"
      },
      "outputs": [],
      "source": [
        "class Binarize(object):\n",
        "  def __init__(self, threshold=0.3):\n",
        "    self.threshold = threshold\n",
        "      \n",
        "  def __call__(self, t):\n",
        "    t = (t > self.threshold).float()\n",
        "    return t\n",
        "  \n",
        "  def __repr__(self):\n",
        "    return self.__class__.__name__ + '(th={0})'.format(self.threshold)\n",
        "\n",
        "\n",
        "class Smooth(object):\n",
        "  def __init__(self, smooth=0.1):\n",
        "    self.smooth = smooth\n",
        "      \n",
        "  def __call__(self, t):\n",
        "    t[t == 1.] = 1 - self.smooth\n",
        "    t[t == 0.] = 0 + self.smooth\n",
        "    return t\n",
        "  \n",
        "  def __repr__(self):\n",
        "    return self.__class__.__name__ + '(smooth={0})'.format(self.smooth)\n",
        "\n",
        "\n",
        "def load_mnist(_data_root='datasets', binarized=False, bin_th=0.3, smooth=None):\n",
        "    trans = [_transforms.ToTensor()]\n",
        "    if binarized:\n",
        "      binarizor = Binarize(bin_th)\n",
        "      trans.append(binarizor)\n",
        "    if smooth is not None:\n",
        "      smoother = Smooth(smooth)\n",
        "      trans.append(smoother)\n",
        "    trans = _transforms.Compose(trans)\n",
        "    _data = _datasets.MNIST(_data_root, train=True, download=True,\n",
        "                            transform=trans)\n",
        "    return _data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mowQ0xtCtDJS"
      },
      "source": [
        "# Generate constraints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GEayFgAWtItV"
      },
      "outputs": [],
      "source": [
        "num=100\n",
        "\n",
        "random.seed(0)\n",
        "torch.manual_seed(0)\n",
        "def generate_constraints(num, M, b_lb=3, b_ub=10, device='cuda'):\n",
        "  device=torch.device(device)\n",
        "  l,r=[],[]\n",
        "  for i in range(num):\n",
        "    r.append(b_lb+(b_ub-b_lb)*random.random())\n",
        "    li=copy.deepcopy(M)\n",
        "    li.to(device)\n",
        "    for p in li.parameters():\n",
        "      p=torch.randn(*p.shape,device=device)\n",
        "    l.append(li)\n",
        "  return [l,r]\n",
        "\n",
        "G = GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "D = DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "C_g=generate_constraints(num, G)\n",
        "C_d=generate_constraints(num, D)\n",
        "\n",
        "\n",
        "def proj(A,generator,eps=1e-5, maxit=10000):\n",
        "  l,r=None,None\n",
        "  if generator:\n",
        "    l,r=C_g\n",
        "  else:\n",
        "    l,r=C_d\n",
        "  c_num=len(r)\n",
        "  it=0\n",
        "  while it<maxit:\n",
        "    idx=-1\n",
        "    max_norm=0\n",
        "    d=eps\n",
        "    for i in range(c_num):\n",
        "      temp=-r[i]\n",
        "      norm=0\n",
        "      for p,pl in zip(A.parameters(), l[i].parameters()):\n",
        "        with torch.no_grad():\n",
        "          temp+=torch.sum(p*pl)\n",
        "          norm+=torch.sum(pl**2)\n",
        "      norm=math.sqrt(norm)\n",
        "      temp/=norm\n",
        "      if temp>d:\n",
        "        d=temp\n",
        "        idx=i\n",
        "        max_norm=norm\n",
        "    if idx==-1:\n",
        "      break\n",
        "    for p,pl in zip(A.parameters(), l[idx].parameters()):\n",
        "      with torch.no_grad():\n",
        "        p-=d/max_norm*pl\n",
        "    it+=1\n",
        "    if it==maxit:\n",
        "      print('*************Projection Failed!****************')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# PI-ACVI training function"
      ],
      "metadata": {
        "id": "JFoqmqiC1SFZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_disciminator_loss_x(Dx, Dy, Dlmd, x_real, x_gen, lbl_real, lbl_fake, beta, use_acvi=True):\n",
        "  \"\"\"\"\"\"\n",
        "  D_x = Dx(x_real)\n",
        "  D_G_z = Dx(x_gen)\n",
        "  lossD_real = torch.binary_cross_entropy_with_logits(D_x, lbl_real).mean()\n",
        "  lossD_fake = torch.binary_cross_entropy_with_logits(D_G_z, lbl_fake).mean()\n",
        "  lossD = lossD_real + lossD_fake\n",
        "  if use_acvi:\n",
        "    for px,py,plmd in zip(Dx.parameters(),Dy.parameters(),Dlmd.parameters()):\n",
        "      lossD+=beta/2*torch.sum((px-(py-plmd/beta))**2)\n",
        "  return lossD\n",
        "\n",
        "\n",
        "def get_generator_loss_x(Gx, Gy, Glmd, Dx, z, lbl_real, beta, use_acvi=True):\n",
        "  \"\"\"\"\"\"\n",
        "  D_G_z = Dx(Gx(z))\n",
        "  lossG = torch.binary_cross_entropy_with_logits(D_G_z, lbl_real).mean()\n",
        "  if use_acvi:\n",
        "    for px,py,plmd in zip(Gx.parameters(),Gy.parameters(),Glmd.parameters()):\n",
        "      lossG += beta/2*torch.sum((px-(py-plmd/beta))**2)\n",
        "  return lossG\n",
        "\n",
        "\n",
        "def get_sampler(dataset, batch_size, shuffle=True, drop_last=True):\n",
        "  dataloader = DataLoader(dataset, batch_size, shuffle=shuffle, drop_last=drop_last)\n",
        "  dataloader_iterator = iter(dataloader)\n",
        "  def sampler():\n",
        "    nonlocal dataloader_iterator\n",
        "    try:\n",
        "        data = next(dataloader_iterator) \n",
        "    except StopIteration:\n",
        "        dataloader_iterator = iter(dataloader)\n",
        "        data = next(dataloader_iterator) \n",
        "    return data\n",
        "  return sampler\n",
        "\n",
        "\n",
        "def train(Gx, Gy, Glmd, Dx, Dy, Dlmd, dataset, iterations, batch_size=32, lrDx=0.01, lrGx=0.01,\n",
        "          beta1=0.99, eval_every=100, device=torch.device('cpu'), use_acvi=True,\n",
        "          plot_func=lambda a,b,c,d: None, extragrad=False, \n",
        "          out_dir=None, lx=1, lx_warmup=100, train_time=1500, beta=0.5, ceps=-0.001):\n",
        "  \n",
        "  XY = [(x,y) for x,y in dataset]\n",
        "  X = torch.stack([x[0] for x in XY]).to('cuda')\n",
        "  Y = torch.tensor([x[1] for x in XY]).long().to('cuda')\n",
        "  dataset = torch.utils.data.TensorDataset(X, Y)\n",
        "    \n",
        "  sampler = get_sampler(dataset, batch_size, shuffle=True, drop_last=True)\n",
        "\n",
        "  if extragrad:\n",
        "    D_extrax = copy.deepcopy(Dx)\n",
        "    G_extrax = copy.deepcopy(Gx)\n",
        "  else:\n",
        "    D_extrax = Dx\n",
        "    G_extrax = Gx\n",
        "\n",
        "  # Optimizers\n",
        "  optimizerDx = torch.optim.Adam(Dx.parameters(), lr=lrDx, betas=(beta1, 0.999))\n",
        "  optimizerGx = torch.optim.Adam(Gx.parameters(), lr=lrGx, betas=(beta1, 0.999))\n",
        "\n",
        "  optimizerD_extrax = torch.optim.Adam(D_extrax.parameters(), lr=lrDx, betas=(beta1, 0.999))\n",
        "  optimizerG_extrax = torch.optim.Adam(G_extrax.parameters(), lr=lrGx, betas=(beta1, 0.999))\n",
        "\n",
        "  # LBLs\n",
        "  lbl_real = torch.ones(batch_size, 1, device=device)\n",
        "  lbl_fake = torch.zeros(batch_size, 1, device=device)\n",
        "\n",
        "  fixed_noise = torch.randn(100, Gx.noise_dim, device=device)\n",
        "\n",
        "  Gx.to(device)\n",
        "  Dx.to(device)\n",
        "  Gy.to(device)\n",
        "  Dy.to(device)\n",
        "  Glmd.to(device)\n",
        "  Dlmd.to(device)\n",
        "\n",
        "  G_extrax.to(device)\n",
        "  D_extrax.to(device)\n",
        "\n",
        "  start_time = time.perf_counter()\n",
        "\n",
        "  for itr in range(iterations):\n",
        "\n",
        "    # UPDATE THE X\n",
        "    for _ in range(lx_warmup if itr == 0 else lx):\n",
        "\n",
        "      err_x = 0\n",
        "\n",
        "      if extragrad:\n",
        "        # STEP 1: get G_{t+1} (G_extra)\n",
        "        optimizerG_extrax.zero_grad()\n",
        "        z = torch.randn(batch_size, G_extrax.noise_dim, device=device)\n",
        "        lossGx = get_generator_loss_x(G_extrax, Gy, Glmd, Dx, z, lbl_real, beta, use_acvi=use_acvi)\n",
        "        lossGx.backward()\n",
        "        optimizerG_extrax.step()\n",
        "        if not use_acvi:\n",
        "          with torch.no_grad():\n",
        "            proj(G_extrax,generator=True)\n",
        "\n",
        "        # STEP 2: get D_{t+1} (D_extra)\n",
        "        optimizerD_extrax.zero_grad()\n",
        "        x_real, _ = sampler()\n",
        "        x_real = x_real.to(device)\n",
        "        z = torch.randn(batch_size, Gx.noise_dim, device=device)\n",
        "        with torch.no_grad():\n",
        "          x_gen = Gx(z)\n",
        "        lossDx = get_disciminator_loss_x(D_extrax, Dy, Dlmd, x_real, x_gen, lbl_real, lbl_fake, beta, use_acvi=use_acvi)\n",
        "        lossDx.backward()\n",
        "        optimizerD_extrax.step()\n",
        "        if not use_acvi:\n",
        "          with torch.no_grad():\n",
        "            proj(D_extrax,generator=False)\n",
        "\n",
        "      # STEP 3: D optimization step using G_extra\n",
        "      x_real, _ = sampler()\n",
        "      x_real = x_real.to(device)\n",
        "      z = torch.randn(batch_size, Gx.noise_dim, device=device)\n",
        "      with torch.no_grad():\n",
        "        x_gen = G_extrax(z) # using G_{t+1}\n",
        "      optimizerDx.zero_grad()\n",
        "      lossDx = get_disciminator_loss_x(Dx, Dy, Dlmd, x_real, x_gen, lbl_real, lbl_fake, beta, use_acvi=use_acvi)\n",
        "      lossDx.backward()\n",
        "      #with torch.no_grad():\n",
        "      #  for px, py, plmd in zip(Dx.parameters(), Dy.parameters(), Dlmd.parameters()):\n",
        "      #    err_x = max((px + 1/beta * px.grad - py + 1/beta * plmd).abs().max(),  err_x)\n",
        "      optimizerDx.step()\n",
        "      optimizerDx.zero_grad()\n",
        "      if not use_acvi:\n",
        "          with torch.no_grad():\n",
        "            proj(Dx,generator=False)\n",
        "\n",
        "      # STEP 4: G optimization step using D_extra\n",
        "      z = torch.randn(batch_size, Gx.noise_dim, device=device)\n",
        "      optimizerGx.zero_grad()\n",
        "      lossGx = get_generator_loss_x(Gx, Gy, Glmd, D_extrax, z, lbl_real, beta, use_acvi=use_acvi) # we use the unrolled D\n",
        "      lossGx.backward()\n",
        "      #with torch.no_grad():\n",
        "      #  for px, py, plmd in zip(Gx.parameters(), Gy.parameters(), Glmd.parameters()):\n",
        "      #    err_x = max((px + 1/beta * px.grad - py + 1/beta * plmd).abs().max(), err_x)\n",
        "      optimizerGx.step()\n",
        "      if not use_acvi:\n",
        "          with torch.no_grad():\n",
        "            proj(Gx,generator=True)\n",
        "\n",
        "      if extragrad:\n",
        "        G_extrax.load_state_dict(Gx.state_dict())\n",
        "        D_extrax.load_state_dict(Dx.state_dict())\n",
        "\n",
        "    time_tick=time.perf_counter() - start_time\n",
        "\n",
        "    # Just plotting things\n",
        "    if itr % eval_every == 0 or itr == iterations-1 or time_tick>=train_time:\n",
        "      with torch.no_grad():\n",
        "        probasx = torch.sigmoid(Dx(Gx(fixed_noise)))\n",
        "        mean_probax = probasx.mean().cpu().item()\n",
        "        std_probax = probasx.std().cpu().item()\n",
        "        samplesx = Gx(fixed_noise)\n",
        "      print(f\"Iter {itr}: Mean proba from Dx(Gx(z)): {mean_probax:.4f} +/- {std_probax:.4f}\")\n",
        "      plot_func(samplesx.detach().cpu(), time_tick=time_tick, D=Dx, G=Gx, iteration=itr, G_avg=None, G_ema=None)\n",
        "\n",
        "    if time_tick >= train_time:\n",
        "      break\n",
        "\n",
        "    # UPDATE THE Y\n",
        "    if use_acvi:\n",
        "      with torch.no_grad():\n",
        "        for px, py, plmd in zip(Gx.parameters(),Gy.parameters(),Glmd.parameters()):\n",
        "          py.data = px + plmd/beta\n",
        "        proj(Gy,generator=True)\n",
        "        for px, py, plmd in zip(Dx.parameters(),Dy.parameters(),Dlmd.parameters()):\n",
        "          py.data = px + plmd/beta\n",
        "        proj(Dy,generator=False)\n",
        "\n",
        "    # UPDATE THE LAMBDAS\n",
        "    if use_acvi:\n",
        "      with torch.no_grad():\n",
        "        for px,py,plmd in zip(Gx.parameters(),Gy.parameters(),Glmd.parameters()):\n",
        "          plmd += beta*(px-py)\n",
        "        for px,py,plmd in zip(Dx.parameters(),Dy.parameters(),Dlmd.parameters()):\n",
        "          plmd += beta*(px-py)\n",
        "\n",
        "    if time_tick >= train_time:\n",
        "      break"
      ],
      "metadata": {
        "id": "gy1pAgL21SOY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JEMA5J179bCN"
      },
      "source": [
        "# Display & Eval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9tBw8aDhlAMx"
      },
      "outputs": [],
      "source": [
        "def compute_mu_sigma_pretrained_model(dataset, pretrained_clf):\n",
        "  dataloader = DataLoader(dataset, batch_size=512, num_workers=2, drop_last=True)\n",
        "  cuda = next(pretrained_clf.parameters()).is_cuda\n",
        "  all_fc2_out = []\n",
        "  pretrained_clf.eval()\n",
        "  for batch, _ in dataloader:\n",
        "    with torch.no_grad():\n",
        "      if cuda:\n",
        "        batch = batch.cuda()\n",
        "      _, fc2_out = pretrained_clf.get_logits_and_fc2_outputs(batch)\n",
        "    all_fc2_out.append(fc2_out.cpu())\n",
        "  all_fc2_out = torch.cat(all_fc2_out, dim=0).numpy()\n",
        "  mu_real = np.mean(all_fc2_out, axis=0)\n",
        "  sigma_real = np.cov(all_fc2_out, rowvar=False)\n",
        "  return mu_real, sigma_real\n",
        "\n",
        "\n",
        "def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):\n",
        "    \"\"\"Numpy implementation of the Frechet Distance.\"\"\"\n",
        "    mu1 = np.atleast_1d(mu1)\n",
        "    mu2 = np.atleast_1d(mu2)\n",
        "\n",
        "    sigma1 = np.atleast_2d(sigma1)\n",
        "    sigma2 = np.atleast_2d(sigma2)\n",
        "\n",
        "    assert mu1.shape == mu2.shape, \"Training and test mean vectors have different lengths\"\n",
        "    assert sigma1.shape == sigma2.shape, \"Training and test covariances have different dimensions\"\n",
        "\n",
        "    diff = mu1 - mu2\n",
        "\n",
        "    # product might be almost singular\n",
        "    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)\n",
        "    if not np.isfinite(covmean).all():\n",
        "        msg = \"fid calculation produces singular product; adding %s to diagonal of cov estimates\" % eps\n",
        "        print(msg)\n",
        "        # warnings.warn(msg)\n",
        "        offset = np.eye(sigma1.shape[0]) * eps\n",
        "        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n",
        "\n",
        "    # numerical error might give slight imaginary component\n",
        "    if np.iscomplexobj(covmean):\n",
        "        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):\n",
        "            m = np.max(np.abs(covmean.imag))\n",
        "            raise ValueError(\"Imaginary component {}\".format(m))\n",
        "        covmean = covmean.real\n",
        "    tr_covmean = np.trace(covmean)\n",
        "\n",
        "    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean\n",
        "\n",
        "\n",
        "def _calculate_metrics(pretrained_clf, G, dataset_length, mu_real, sigma_real, \n",
        "                       n_classes=10, batch_size=1024):\n",
        "    cuda = next(pretrained_clf.parameters()).is_cuda\n",
        "    if cuda:\n",
        "      device = torch.device('cuda')\n",
        "    else:\n",
        "      device = torch.device('cpu')\n",
        "    # Using pretrained clf to get predictions over fake data\n",
        "    inception_predictions, all_fc2_out, class_probas = [], [], []\n",
        "    dataloader = DataLoader(list(range(dataset_length)), batch_size, num_workers=2, drop_last=True)\n",
        "    pretrained_clf.eval()\n",
        "    for batch in dataloader:\n",
        "      with torch.no_grad():\n",
        "        noise = torch.randn(batch_size, G.noise_dim, device=device)\n",
        "        probas, fc2_out = pretrained_clf.get_logits_and_fc2_outputs(G(noise).view(batch_size, -1))\n",
        "      all_fc2_out.append(fc2_out.cpu())\n",
        "      class_probas.append(probas.cpu())\n",
        "    all_fc2_out = torch.cat(all_fc2_out, dim=0).numpy()\n",
        "    class_probas = torch.cat(class_probas, dim=0)\n",
        "    inception_predictions = torch.softmax(class_probas, dim=1).numpy()\n",
        "    class_probas = class_probas.numpy()\n",
        "    pred_prob = np.maximum(class_probas, 1e-20 * np.ones_like(class_probas))\n",
        "\n",
        "    y_vec = 1e-20 * np.ones((len(pred_prob), n_classes), dtype=np.float)  # pred label distr\n",
        "    gnd_vec = 0.1 * np.ones((1, n_classes), dtype=np.float)  # gnd label distr, uniform over classes\n",
        "\n",
        "    for i, label in enumerate(pred_prob):\n",
        "        y_vec[i, np.argmax(pred_prob[i])] = 1.0\n",
        "    y_vec = np.sum(y_vec, axis=0, keepdims=True)\n",
        "    y_vec = y_vec / np.sum(y_vec)\n",
        "\n",
        "    label_entropy = np.sum(-y_vec * np.log(y_vec)).tolist()\n",
        "    label_tv = np.true_divide(np.sum(np.abs(y_vec - gnd_vec)), 2).tolist()\n",
        "    label_l2 = np.sum((y_vec - gnd_vec) ** 2).tolist()\n",
        "\n",
        "    # --- is ----\n",
        "    inception_scores = []\n",
        "    for i in range(n_classes):\n",
        "        part = inception_predictions[(i * inception_predictions.shape[0]\n",
        "                                      // n_classes):((i + 1) * inception_predictions.shape[0]\n",
        "                                                     // n_classes), :]\n",
        "        kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))\n",
        "        kl = np.mean(np.sum(kl, 1))\n",
        "        inception_scores.append(np.exp(kl))\n",
        "\n",
        "    mu = np.mean(all_fc2_out, axis=0)\n",
        "    sigma = np.cov(all_fc2_out, rowvar=False)\n",
        "    _fid = calculate_frechet_distance(mu, sigma, mu_real, sigma_real)\n",
        "\n",
        "    return (label_entropy, label_tv, label_l2,\n",
        "            float(np.mean(inception_scores)),\n",
        "            float(np.std(inception_scores)),\n",
        "            _fid)\n",
        "\n",
        "\n",
        "def get_metrics(pretrained_clf, dataset_length, mu_real, sigma_real, G):\n",
        "    \"\"\"Calculates entropy, TV, L2, and inception scores.\"\"\"\n",
        "    e, tv, l2, is_m, is_std, fid = _calculate_metrics(pretrained_clf,\n",
        "                                                      G,\n",
        "                                                      dataset_length,\n",
        "                                                      mu_real,\n",
        "                                                      sigma_real)\n",
        "    m_result = {\n",
        "        'entropy': e,\n",
        "        'TV': tv,\n",
        "        'L2': l2,\n",
        "        'inception_mean': is_m,\n",
        "        'inception_std': is_std,\n",
        "        'fid': fid\n",
        "    }\n",
        "    return m_result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9T2JfCAd5SXA"
      },
      "outputs": [],
      "source": [
        "def save_models(G, D, opt_G, opt_D, out_dir, suffix):\n",
        "  torch.save(G.state_dict(), os.path.join(out_dir, f\"gen_{suffix}.pth\"))\n",
        "  torch.save(D.state_dict(), os.path.join(out_dir, f\"disc_{suffix}.pth\"))\n",
        "  torch.save(opt_G.state_dict(), os.path.join(out_dir, f\"gen_optim_{suffix}.pth\"))\n",
        "  torch.save(opt_D.state_dict(), os.path.join(out_dir, f\"disc_optim_{suffix}.pth\"))\n",
        "\n",
        "\n",
        "def get_plot_func(out_dir, img_size, num_samples_eval=10000, save_curves=None):\n",
        "  dataset = load_mnist(_data_root='datasets', binarized=False)\n",
        "  #shutil.rmtree(out_dir, ignore_errors=True)\n",
        "  #if not os.path.exists(out_dir):\n",
        "  #  os.makedirs(out_dir)\n",
        "  pretrained_clf = pretrained_mnist_model(pretrained='/content/drive/My Drive/mnist_exp/ACVI/mnist/mnist.pth')\n",
        "  mu_real, sigma_real = compute_mu_sigma_pretrained_model(dataset, pretrained_clf)\n",
        "  inception_means, inception_stds, inception_means_ema, inception_means_avg, fids, fids_ema, fids_avg = [], [], [], [], [], [], []\n",
        "  iterations, times = [], []\n",
        "  def plot_func(samples, iteration, time_tick, G=None, D=None, G_avg=None, G_ema=None):\n",
        "    fig = plt.figure(figsize=(12,5), dpi=100)\n",
        "    plt.subplot(1,2,1)\n",
        "    samples = samples.view(100, *img_size)\n",
        "    file_name = os.path.join(out_dir, '%08d.png' % iteration)\n",
        "    vision_utils.save_image(samples, file_name, nrow=10)\n",
        "    grid_img = vision_utils.make_grid(samples, nrow=10, normalize=True, padding=0)\n",
        "    plt.imshow(grid_img.permute(1, 2, 0), interpolation='nearest')\n",
        "    plt.subplot(1,2,2)\n",
        "    metrics = get_metrics(pretrained_clf, num_samples_eval, mu_real, sigma_real, G)\n",
        "    fids.append(metrics['fid'])\n",
        "    inception_means.append(metrics['inception_mean'])\n",
        "    inception_stds.append(metrics['inception_std'])\n",
        "    if G_avg is not None:\n",
        "      metrics = get_metrics(pretrained_clf, num_samples_eval, mu_real, sigma_real, G_avg)\n",
        "      fids_avg.append(metrics['fid'])\n",
        "      inception_means_avg.append(metrics['inception_mean'])\n",
        "    if G_ema is not None:\n",
        "      metrics = get_metrics(pretrained_clf, num_samples_eval, mu_real, sigma_real, G_ema)\n",
        "      fids_ema.append(metrics['fid'])\n",
        "      inception_means_ema.append(metrics['inception_mean'])\n",
        "    iterations.append(iteration)\n",
        "    times.append(time_tick)\n",
        "    #  is\n",
        "    is_low  = [m - s for m, s in zip(inception_means, inception_stds)]\n",
        "    is_high = [m + s for m, s in zip(inception_means, inception_stds)]\n",
        "    plt.plot(times, inception_means, label=\"is\", color='r')\n",
        "    plt.fill_between(times, is_low, is_high, facecolor='r', alpha=.3)\n",
        "    plt.yticks(np.arange(0, 10+1, 0.5))\n",
        "    # fid\n",
        "    plt.plot(times, fids, label=\"fid\", color='b')\n",
        "    plt.xlabel('Time (sec)')\n",
        "    plt.ylabel('Metric')\n",
        "    plt.grid()\n",
        "    ax = fig.gca()\n",
        "    ax.set_ylim(-0.1, 10)\n",
        "    plt.legend(fancybox=True, framealpha=.5)\n",
        "    curves_img_file_name = os.path.join(out_dir, 'curves.png')\n",
        "    fig.savefig(curves_img_file_name)\n",
        "    plt.show()\n",
        "    curves_file_name = os.path.join(out_dir, 'curves.json')\n",
        "    curves = {\n",
        "        'inception_means': list(inception_means),\n",
        "        'inception_stds': list(inception_stds),\n",
        "        'inception_means_ema': list(inception_means_ema),\n",
        "        'inception_means_avg': list(inception_means_avg),\n",
        "        'fids_ema': list(fids_ema),\n",
        "        'fids_avg': list(fids_avg),\n",
        "        'fids': list(fids),\n",
        "        'iterations':iterations,\n",
        "        'times': times\n",
        "    }\n",
        "    with open(curves_file_name, 'w') as fs:\n",
        "      json.dump(curves, fs)\n",
        "  return plot_func"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wh8ag1jR9t2K"
      },
      "source": [
        "# Experiments"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## GDA"
      ],
      "metadata": {
        "id": "1aXoXq7icSK0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "args = dict(iterations = 10000,\n",
        "            batch_size = 50,\n",
        "            lrDx = 0.001,\n",
        "            lrGx = 0.001,\n",
        "            beta1 = 0.05,\n",
        "            extragrad = False,\n",
        "            eval_every = 100,\n",
        "            device = 'cuda',\n",
        "            lx_warmup=1,\n",
        "            lx=1,\n",
        "            beta=0.5,\n",
        "            use_acvi=False)\n",
        "\n",
        "for k in range(0,4):\n",
        "  torch.manual_seed(k)\n",
        "  torch.cuda.manual_seed(k)\n",
        "  np.random.seed(k)\n",
        "  exp_key = f\"iter{args['iterations']}_bs{args['batch_size']}_lrD{args['lrDx']}\" + \\\n",
        "            f\"_lrG{args['lrGx']}_beta1{args['beta1']}_beta{args['beta']}\" + \\\n",
        "            f\"_extragrad{args['extragrad']}_ee{args['eval_every']}_lx{args['lx']}_lxw{args['lx_warmup']}_useacvi{args['use_acvi']}\"\n",
        "  out_dir = f\"/content/drive/My Drive/mnist_exp/ACVI/aistats/EG/{exp_key}/{k}/\"\n",
        "\n",
        "  shutil.rmtree(out_dir, ignore_errors=True)\n",
        "  if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "  with open(os.path.join(out_dir, 'args.json'), 'w') as fs:\n",
        "    json.dump(args, fs)\n",
        "\n",
        "  dataset = load_mnist(_data_root='datasets', binarized=False)\n",
        "\n",
        "  plot_func = get_plot_func(out_dir=out_dir, \n",
        "                img_size=dataset[0][0].size(),\n",
        "                num_samples_eval=10000)\n",
        "\n",
        "  Gx = GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  Dx = DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "\n",
        "  Gy=GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  with torch.no_grad():\n",
        "    for px, py in zip(Gx.parameters(), Gy.parameters()):\n",
        "      py.data = copy.deepcopy(px.data)\n",
        "  Dy=DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "  with torch.no_grad():\n",
        "    for px, py in zip(Dx.parameters(), Dy.parameters()):\n",
        "      py.data = copy.deepcopy(px.data)\n",
        "\n",
        "  Glmd=GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  with torch.no_grad():\n",
        "    for p in Glmd.parameters():\n",
        "      p.data.fill_(0.0)\n",
        "  Dlmd=DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "  with torch.no_grad():\n",
        "    for p in Dlmd.parameters():\n",
        "      p.data.fill_(0.0)\n",
        "\n",
        "  train(Gx, Gy, Glmd, Dx, Dy, Dlmd, dataset, \n",
        "        iterations=args['iterations'], \n",
        "        batch_size=args['batch_size'], \n",
        "        lrDx=args['lrDx'], \n",
        "        lrGx=args['lrGx'],\n",
        "        beta1=args['beta1'], \n",
        "        eval_every=args['eval_every'], \n",
        "        device=torch.device(args['device']), \n",
        "        plot_func=plot_func, \n",
        "        extragrad=args['extragrad'], \n",
        "        out_dir=out_dir, \n",
        "        lx=args['lx'], \n",
        "        lx_warmup=args['lx_warmup'], \n",
        "        beta=args['beta'],\n",
        "        use_acvi=args['use_acvi'])"
      ],
      "metadata": {
        "id": "u0FNpyECcBhr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## EG"
      ],
      "metadata": {
        "id": "l18pEvu01dhQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "args = dict(iterations = 5000,\n",
        "            batch_size = 50,\n",
        "            lrDx = 0.001,\n",
        "            lrGx = 0.001,\n",
        "            beta1 = 0.05,\n",
        "            extragrad = True,\n",
        "            eval_every = 100,\n",
        "            device = 'cuda',\n",
        "            lx_warmup=1,\n",
        "            lx=1,\n",
        "            beta=0.5,\n",
        "            use_acvi=False)\n",
        "\n",
        "for k in range(0,4):\n",
        "  torch.manual_seed(k)\n",
        "  torch.cuda.manual_seed(k)\n",
        "  np.random.seed(k)\n",
        "  exp_key = f\"iter{args['iterations']}_bs{args['batch_size']}_lrD{args['lrDx']}\" + \\\n",
        "            f\"_lrG{args['lrGx']}_beta1{args['beta1']}_beta{args['beta']}\" + \\\n",
        "            f\"_extragrad{args['extragrad']}_ee{args['eval_every']}_lx{args['lx']}_lxw{args['lx_warmup']}_useacvi{args['use_acvi']}\"\n",
        "  out_dir = f\"/content/drive/My Drive/mnist_exp/ACVI/aistats/EG/{exp_key}/{k}/\"\n",
        "\n",
        "  shutil.rmtree(out_dir, ignore_errors=True)\n",
        "  if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "  with open(os.path.join(out_dir, 'args.json'), 'w') as fs:\n",
        "    json.dump(args, fs)\n",
        "\n",
        "  dataset = load_mnist(_data_root='datasets', binarized=False)\n",
        "\n",
        "  plot_func = get_plot_func(out_dir=out_dir, \n",
        "                img_size=dataset[0][0].size(),\n",
        "                num_samples_eval=10000)\n",
        "\n",
        "  Gx = GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  Dx = DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "\n",
        "  Gy=GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  with torch.no_grad():\n",
        "    for px, py in zip(Gx.parameters(), Gy.parameters()):\n",
        "      py.data = copy.deepcopy(px.data)\n",
        "  Dy=DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "  with torch.no_grad():\n",
        "    for px, py in zip(Dx.parameters(), Dy.parameters()):\n",
        "      py.data = copy.deepcopy(px.data)\n",
        "\n",
        "  Glmd=GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  with torch.no_grad():\n",
        "    for p in Glmd.parameters():\n",
        "      p.data.fill_(0.0)\n",
        "  Dlmd=DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "  with torch.no_grad():\n",
        "    for p in Dlmd.parameters():\n",
        "      p.data.fill_(0.0)\n",
        "\n",
        "  train(Gx, Gy, Glmd, Dx, Dy, Dlmd, dataset, \n",
        "        iterations=args['iterations'], \n",
        "        batch_size=args['batch_size'], \n",
        "        lrDx=args['lrDx'], \n",
        "        lrGx=args['lrGx'],\n",
        "        beta1=args['beta1'], \n",
        "        eval_every=args['eval_every'], \n",
        "        device=torch.device(args['device']), \n",
        "        plot_func=plot_func, \n",
        "        extragrad=args['extragrad'], \n",
        "        out_dir=out_dir, \n",
        "        lx=args['lx'], \n",
        "        lx_warmup=args['lx_warmup'], \n",
        "        beta=args['beta'],\n",
        "        use_acvi=args['use_acvi'])"
      ],
      "metadata": {
        "id": "S0c_NjB91cXx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## PI-ACVI"
      ],
      "metadata": {
        "id": "r1s99t3_4b2n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "args = dict(iterations = 5000,\n",
        "            batch_size = 50,\n",
        "            lrDx = 0.001,\n",
        "            lrGx = 0.001,\n",
        "            beta1 = 0.05,\n",
        "            extragrad = False,\n",
        "            eval_every = 100,\n",
        "            device = 'cuda',\n",
        "            lx_warmup=500,\n",
        "            lx=20,\n",
        "            beta=0.5,\n",
        "            use_acvi=True)\n",
        "\n",
        "for k in range(0,4):\n",
        "  torch.manual_seed(k)\n",
        "  torch.cuda.manual_seed(k)\n",
        "  np.random.seed(k)\n",
        "  exp_key = f\"iter{args['iterations']}_bs{args['batch_size']}_lrD{args['lrDx']}\" + \\\n",
        "            f\"_lrG{args['lrGx']}_beta1{args['beta1']}_beta{args['beta']}\" + \\\n",
        "            f\"_extragrad{args['extragrad']}_ee{args['eval_every']}_lx{args['lx']}_lxw{args['lx_warmup']}_useacvi{args['use_acvi']}\"\n",
        "  out_dir = f\"/content/drive/My Drive/mnist_exp/ACVI/aistats/EG/{exp_key}/{k}/\"\n",
        "\n",
        "  shutil.rmtree(out_dir, ignore_errors=True)\n",
        "  if not os.path.exists(out_dir):\n",
        "    os.makedirs(out_dir)\n",
        "\n",
        "  with open(os.path.join(out_dir, 'args.json'), 'w') as fs:\n",
        "    json.dump(args, fs)\n",
        "\n",
        "  dataset = load_mnist(_data_root='datasets', binarized=False)\n",
        "\n",
        "  plot_func = get_plot_func(out_dir=out_dir, \n",
        "                img_size=dataset[0][0].size(),\n",
        "                num_samples_eval=10000)\n",
        "\n",
        "  Gx = GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  Dx = DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "\n",
        "  Gy=GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  with torch.no_grad():\n",
        "    for px, py in zip(Gx.parameters(), Gy.parameters()):\n",
        "      py.data = copy.deepcopy(px.data)\n",
        "  Dy=DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "  with torch.no_grad():\n",
        "    for px, py in zip(Dx.parameters(), Dy.parameters()):\n",
        "      py.data = copy.deepcopy(px.data)\n",
        "\n",
        "  Glmd=GeneratorCNN28(noise_dim=_NOISE_DIM, out_tanh=True)\n",
        "  with torch.no_grad():\n",
        "    for p in Glmd.parameters():\n",
        "      p.data.fill_(0.0)\n",
        "  Dlmd=DiscriminatorCNN28(spectral_norm=False, img_size=28)\n",
        "  with torch.no_grad():\n",
        "    for p in Dlmd.parameters():\n",
        "      p.data.fill_(0.0)\n",
        "\n",
        "  train(Gx, Gy, Glmd, Dx, Dy, Dlmd, dataset, \n",
        "        iterations=args['iterations'], \n",
        "        batch_size=args['batch_size'], \n",
        "        lrDx=args['lrDx'], \n",
        "        lrGx=args['lrGx'],\n",
        "        beta1=args['beta1'], \n",
        "        eval_every=args['eval_every'], \n",
        "        device=torch.device(args['device']), \n",
        "        plot_func=plot_func, \n",
        "        extragrad=args['extragrad'], \n",
        "        out_dir=out_dir, \n",
        "        lx=args['lx'], \n",
        "        lx_warmup=args['lx_warmup'], \n",
        "        beta=args['beta'],\n",
        "        use_acvi=args['use_acvi'])"
      ],
      "metadata": {
        "id": "95mNeZP76l8g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "xJlgnZgI1XoJ"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "standard",
    "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": 0
}