VAEs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "import torchvision.transforms as tranforms\n",
    "import torchvision.datasets as datasets\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VAE and Beta-VAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.device_count()\n",
    "torch.cuda.current_device()\n",
    "torch.cuda.get_device_name(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device=torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\n",
    "# device=\"cpu\"\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "tranform=tranforms.Compose([tranforms.Resize((64,64)),tranforms.ToTensor()])\n",
    "dataset=torchvision.datasets.ImageFolder(root=\"./img_align_celeba\",transform=tranform)\n",
    "train_data=DataLoader(dataset=dataset,batch_size=64,shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_two=np.load('./dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz')\n",
    "dataset_two=torch.tensor(dataset_two['imgs']).unsqueeze(dim=1)\n",
    "dataset_two=dataset_two.float()\n",
    "train_data_two=DataLoader(dataset=dataset_two,batch_size=256,shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_data=train_data\n",
    "# in_channels=3                   # Chossing CELEB \n",
    "train_data=train_data_two\n",
    "in_channels=1         # Chosing Dsprits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class VAE(nn.Module):\n",
    "\n",
    "\n",
    "    def __init__(self,in_channels,latent_dim):\n",
    "      \n",
    "        super(VAE, self).__init__()\n",
    "\n",
    "        self.latent_dim = latent_dim\n",
    "        self.in_channels=in_channels\n",
    "\n",
    "        # Build Encoder\n",
    "        units = []\n",
    "\n",
    "        units.append(nn.Sequential(\n",
    "                            nn.Conv2d(self.in_channels, out_channels=32,kernel_size= 3, stride= 2, padding  = 1),\n",
    "                            nn.BatchNorm2d(32),\n",
    "                            nn.LeakyReLU()\n",
    "                            ))\n",
    "        units.append(nn.Sequential(\n",
    "                            nn.Conv2d(32, out_channels=64,kernel_size= 3, stride= 2, padding  = 1),\n",
    "                            nn.BatchNorm2d(64),\n",
    "                            nn.LeakyReLU()\n",
    "                            ))\n",
    "        units.append(nn.Sequential(\n",
    "                            nn.Conv2d(64, out_channels=128,kernel_size= 3, stride= 2, padding  = 1),\n",
    "                            nn.BatchNorm2d(128),\n",
    "                            nn.LeakyReLU()\n",
    "                            ))\n",
    "        units.append(nn.Sequential(\n",
    "                            nn.Conv2d(128, out_channels=256,kernel_size= 3, stride= 2, padding  = 1),\n",
    "                            nn.BatchNorm2d(256),\n",
    "                            nn.LeakyReLU()\n",
    "                            ))\n",
    "        units.append(nn.Sequential(\n",
    "                            nn.Conv2d(256, out_channels=512,kernel_size= 3, stride= 2, padding  = 1),\n",
    "                            nn.BatchNorm2d(512),\n",
    "                            nn.LeakyReLU()\n",
    "                            ))\n",
    "\n",
    "        units.append(nn.Sequential(\n",
    "                nn.Conv2d(512, out_channels=1024,kernel_size= 3, stride= 2, padding  = 1),\n",
    "                nn.LeakyReLU())\n",
    "                )\n",
    "        units.append(\n",
    "            nn.Flatten(start_dim=1)\n",
    "            )\n",
    "        units.append(\n",
    "            nn.Linear(1024, latent_dim)\n",
    "            )\n",
    "        self.encoder=nn.Sequential(*units)\n",
    "        self.mean_module = nn.Linear(self.latent_dim, self.latent_dim)\n",
    "        self.variance_module = nn.Linear(self.latent_dim, self.latent_dim)\n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "        # Build Decoder\n",
    "        \n",
    "\n",
    "        self.decoder_input = nn.Linear(latent_dim, 1024)\n",
    "\n",
    "        units = []\n",
    "\n",
    "        units.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ConvTranspose2d(1024,\n",
    "                                       512,\n",
    "                                       kernel_size=3,\n",
    "                                       stride = 2,\n",
    "                                       padding=1,\n",
    "                                       output_padding=1),\n",
    "                    nn.BatchNorm2d(512),\n",
    "                    nn.LeakyReLU())\n",
    "                    )\n",
    "        units.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ConvTranspose2d(512,\n",
    "                                       256,\n",
    "                                       kernel_size=3,\n",
    "                                       stride = 2,\n",
    "                                       padding=1,\n",
    "                                       output_padding=1),\n",
    "                    nn.BatchNorm2d(256),\n",
    "                    nn.LeakyReLU())\n",
    "                    )\n",
    "        units.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ConvTranspose2d(256,\n",
    "                                       128,\n",
    "                                       kernel_size=3,\n",
    "                                       stride = 2,\n",
    "                                       padding=1,\n",
    "                                       output_padding=1),\n",
    "                    nn.BatchNorm2d(128),\n",
    "                    nn.LeakyReLU())\n",
    "                    )\n",
    "\n",
    "        units.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ConvTranspose2d(128,\n",
    "                                       64,\n",
    "                                       kernel_size=3,\n",
    "                                       stride = 2,\n",
    "                                       padding=1,\n",
    "                                       output_padding=1),\n",
    "                    nn.BatchNorm2d(64),\n",
    "                    nn.LeakyReLU())\n",
    "                    )\n",
    "        units.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ConvTranspose2d(64,\n",
    "                                       32,\n",
    "                                       kernel_size=3,\n",
    "                                       stride = 2,\n",
    "                                       padding=1,\n",
    "                                       output_padding=1),\n",
    "                    nn.BatchNorm2d(32),\n",
    "                    nn.LeakyReLU())\n",
    "                    )\n",
    "        \n",
    "        units.append(nn.Sequential(\n",
    "                            nn.ConvTranspose2d(32,\n",
    "                                               32,\n",
    "                                               kernel_size=3,\n",
    "                                               stride=2,\n",
    "                                               padding=1,\n",
    "                                               output_padding=1),\n",
    "                            nn.BatchNorm2d(32),\n",
    "                            nn.LeakyReLU())\n",
    "                        )\n",
    "\n",
    "        units.append(nn.Sequential(\n",
    "                        nn.Conv2d(32, out_channels= self.in_channels,\n",
    "                                      kernel_size= 3, padding= 1),\n",
    "                        nn.Sigmoid())\n",
    "                        )\n",
    "\n",
    "        self.decoder = nn.Sequential(*units)\n",
    "\n",
    "\n",
    "    def encode(self, input):\n",
    "        mu = self.mean_module(result)\n",
    "        log_var = self.variance_module(result)\n",
    "\n",
    "\n",
    "        result = self.encoder(input)\n",
    "        mu = result\n",
    "        log_var = result\n",
    "\n",
    "        return [mu, log_var]\n",
    "\n",
    "    def decode(self, z):\n",
    "\n",
    "        \n",
    "        result = self.decoder_input(z)\n",
    "        result = result.view(-1, 1024, 1, 1)\n",
    "        # result= result.to(\"cuda\")\n",
    "        result = self.decoder(result)\n",
    "        return result\n",
    "\n",
    "    def Reparm(self, mean,log_var):\n",
    "\n",
    "        m = torch.exp(0.5 * log_var)\n",
    "        epsilon = torch.randn_like(m)\n",
    "        return epsilon * m + mean\n",
    "\n",
    "    def forward(self, input):\n",
    "\n",
    "        mu, log_var = self.encode(input)\n",
    "        x_hat=self.decode(self.Reparm(mu, log_var))\n",
    "        return  [x_hat,mu, log_var]\n",
    "\n",
    "    def total_loss(self,\n",
    "                      reconstruction_error,\n",
    "                      mu,log_var,beta=1):\n",
    "\n",
    "        temp=torch.sum(1 + log_var - pow(mu,2) - log_var.exp(), dim = 1)\n",
    "        kld_loss = torch.mean(-0.5 * temp, dim = 0)\n",
    "\n",
    "        loss = reconstruction_error + beta*kld_loss\n",
    "        return loss\n",
    "\n",
    "    def sample_n(self,n):\n",
    "        z = torch.randn(n,\n",
    "                        self.latent_dim)\n",
    "\n",
    "        samples = self.decode(z)\n",
    "        return samples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 5 # epoch- means how much times the the \"entire\" dataset is used to memorise.\n",
    "\n",
    "batch_size = 256 # Since we plan to use Stochastic gradient descent \n",
    "\n",
    "lr = 0.001 # Step size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model=VAE(in_channels=in_channels,latent_dim=5)\n",
    "model=model.cuda()\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "criterion = nn.MSELoss(reduction='sum')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(model.parameters()).is_cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps=int(len(dataset)/batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "def fit_celeb(model, dataloader,beta):\n",
    "    for i in range(epochs):\n",
    "        print(\"\\n\\n\\nEpoch \",i)\n",
    "        running_loss = 0.0\n",
    "        for i, (data,_) in tqdm(enumerate(dataloader),total=num_steps):\n",
    "        \n",
    "            data = data.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            reconstruction, mu, logvar= model(data)\n",
    "            reconstruction= criterion(reconstruction, data)\n",
    "            loss = model.total_loss(reconstruction, mu, logvar,beta=1)\n",
    "            running_loss += loss.item()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        train_loss = running_loss/len(dataloader.dataset)\n",
    "    return train_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_dpsrite(model, dataloader,beta):\n",
    "    for i in range(epochs):\n",
    "        print(\"\\n\\n\\nEpoch \",i)\n",
    "        running_loss = 0.0\n",
    "        for i, data in tqdm(enumerate(dataloader),total=num_steps):\n",
    "        \n",
    "            data = data.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            \n",
    "            reconstruction, mu, logvar = model(data)\n",
    "            reconstruction= criterion(reconstruction, data)\n",
    "            loss = model.total_loss(reconstruction, mu, logvar, beta=beta)\n",
    "            running_loss += loss.item()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        train_loss = running_loss/len(dataloader.dataset)\n",
    "    return train_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit_dpsrite(model,train_data,beta=beta)\n",
    "fit_celeb(model,train_data,beta=beta)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64, 64])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_two[0]=torch.unsqueeze(dataset_two[0],dim=0)\n",
    "dataset_two[0].size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model,\"/home/liyana/disk/liyana/VAE/saved_models/model3_beta4_d1.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.load(\"/home/liyana/disk/liyana/VAE/saved_models/model1_beta1_d2.pth\")\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_test = torch.randn(1,model.latent_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "m=model.decode(z_test.to(device)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "m=m.detach().cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 64, 64])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "if (num_channels == 3):\n",
    "    plt.imshow(m.permute(1, 2, 0))\n",
    "else:\n",
    "    plt.imshow(m, cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = torch.randn(100,model.latent_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "m=model.decode(sample.to(device)).squeeze().detach().cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=plt.figure(1,figsize=(40,40))\n",
    "for idx in range(100):\n",
    "    # print(images.min())\n",
    "    ax=fig.add_subplot(10,10,idx+1,xticks=[],yticks=[])\n",
    "    #plt.imshow(m[idx].permute(1,2,0))\n",
    "    plt.imshow(m[idx], cmap=\"gray\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VanillaVAE_truth_celeb(\n",
       "  (encoder): Sequential(\n",
       "    (0): Sequential(\n",
       "      (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (4): Sequential(\n",
       "      (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (1): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (6): Flatten(start_dim=1, end_dim=-1)\n",
       "    (7): Linear(in_features=1024, out_features=5, bias=True)\n",
       "  )\n",
       "  (mean_module): Linear(in_features=5, out_features=5, bias=True)\n",
       "  (variance_module): Linear(in_features=5, out_features=5, bias=True)\n",
       "  (decoder_input): Linear(in_features=5, out_features=1024, bias=True)\n",
       "  (decoder): Sequential(\n",
       "    (0): Sequential(\n",
       "      (0): ConvTranspose2d(1024, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): ConvTranspose2d(512, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): ConvTranspose2d(256, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (4): Sequential(\n",
       "      (0): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): ConvTranspose2d(32, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): LeakyReLU(negative_slope=0.01)\n",
       "    )\n",
       "    (6): Sequential(\n",
       "      (0): Conv2d(32, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): Sigmoid()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#### DiSENTAGLEMENT PLOT ###############\n",
    "# Works for both ##\n",
    "model = torch.load(\"/home/liyana/disk/liyana/VAE/saved_models/model3_beta4_d2.pth\")\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "z=torch.randn(1,model.latent_dim) # dim(z)=1*latent_dim\n",
    "step=0.5\n",
    "N=10\n",
    "replacement_vec=torch.arange(start=-2.5,end=2.5,step=0.5)\n",
    "# m=model.decode(torch.randn(5,model.latent_dim)).detach().squeeze(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base_matrix=z.repeat(repeats=(N,1))\n",
    "index=0\n",
    "Base_matrix[:, index] = replacement_vec\n",
    "container=Base_matrix.unsqueeze(dim=0)\n",
    "# container.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index in range(1,model.latent_dim):\n",
    "    Base_matrix[:, index] = replacement_vec\n",
    "    container=torch.cat([container,Base_matrix.unsqueeze(dim=0)],dim=0)\n",
    "# container.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "m=model.decode(container.to(device)).detach().cpu()\n",
    "if model.in_channels==1:\n",
    "    m=m.squeeze(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=plt.figure(1,figsize=(40,40))\n",
    "for idx in range(model.latent_dim*N):\n",
    "    ax=fig.add_subplot(10,N,idx+1,xticks=[],yticks=[])\n",
    "    if model.in_channels==1:\n",
    "        plt.imshow(m[idx], cmap=\"gray\")\n",
    "    else:\n",
    "        plt.imshow(m[idx].permute(1,2,0), cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 (conda)",
   "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.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "4570e9e04e7f163b29fd78671957f1c8f837ebf8f072e0060b422ba09ee4ee51"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}