{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Load libs"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom math import log2, floor\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data as data_utils\nimport torchvision\nfrom torchvision import transforms\nfrom torchvision.models import densenet121\nfrom torch.distributions.multivariate_normal import MultivariateNormal","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define dataset and model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_transform():\n    \"\"\"\n    Make data transform.\n    \"\"\"\n    return transforms.Compose([\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomVerticalFlip(),\n        transforms.RandomAffine(40),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ])\n\nclass RecursionDataset(data_utils.Dataset):\n    def __init__(self,\n                 train=True,\n                 img_dir=\"/kaggle/input/recursion-simple\",\n                 transform=make_transform()):\n\n        self.mode = \"train\" if train else \"test\"\n        self.img_dir = img_dir\n        self.transform = transform\n\n        df = pd.read_csv(os.path.join(img_dir, f\"new_{self.mode}.csv\"))\n        self.records = df.to_records(index=False)\n        \n    def __getitem__(self, index):\n        img_path = os.path.join(\n            self.img_dir, self.mode, self.mode, self.records[index].filename)\n        img = Image.open(img_path)\n        img = self.transform(img)\n        sirna = int(self.records[index].sirna) if self.mode == \"train\" else -1\n        sirna = torch.tensor(sirna)\n        return (img, sirna)\n\n    def __len__(self):\n        return len(self.records)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DCGAN(nn.Module):\n    \"\"\"Deep Convolutional Generative Adversarial Network\"\"\"\n\n    def __init__(self,\n                 num_latents=100,\n                 num_features=64,\n                 image_channels=3,\n                 image_size=64,\n                 gan_type=\"gan\",\n                 fully_convolutional=True,\n                 activation=None,\n                 use_batchnorm=True,\n                 use_spectralnorm=False,):\n        \"\"\"\n        Initializes DCGAN.\n\n        Args:\n            num_latents: Number of latent factors.\n            num_features: Number of features in the convolutions.\n            image_channels: Number of channels in the input image.\n            image_size: Size (i.e. height or width) of image.\n            gan_type: Type of GAN (e.g. \"gan\" or \"wgan-gp\").\n        \"\"\"\n        super().__init__()\n\n        self.num_latents = num_latents\n        self.num_features = num_features\n        self.image_channels = image_channels\n        self.image_size = image_size\n        self.gan_type = gan_type\n        self.fully_convolutional = fully_convolutional\n        self.activation = activation\n        self.use_batchnorm = use_batchnorm\n        self.use_spectralnorm = use_spectralnorm\n\n        D_params = {\n            \"num_latents\": 1,  # XXX\n            \"num_features\": num_features,\n            \"image_channels\": image_channels,\n            \"image_size\": image_size,\n            \"gan_type\": gan_type,\n            \"fully_convolutional\": fully_convolutional,\n            \"activation\": activation,\n            \"use_batchnorm\": use_batchnorm,\n            \"use_spectralnorm\": use_spectralnorm,\n        }\n        G_params = {\n            \"num_latents\": num_latents,\n            \"num_features\": num_features,\n            \"image_channels\": image_channels,\n            \"image_size\": image_size,\n            \"gan_type\": gan_type,\n            \"fully_convolutional\": fully_convolutional,\n            \"activation\": activation,\n            \"use_batchnorm\": use_batchnorm,\n            \"use_spectralnorm\": use_spectralnorm,\n        }\n\n        self.D = DCGAN_Discriminator(**D_params)\n        self.G = DCGAN_Generator(**G_params)\n\n\nclass DCGAN_DiscriminatorBlock(nn.Module):\n    \"\"\"\n    A discriminator convolutional block.\n    Default stride and padding half the size of features,\n    e.g. if input is [in_channels, 64, 64], output will be [out_channels, 32, 32].\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1,\n                 use_batchnorm=True, use_spectralnorm=False, activation=None):\n        super().__init__()\n\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,\n                              stride=stride, padding=padding, bias=False)\n        if use_spectralnorm:\n            self.conv = nn.utils.spectral_norm(self.conv)\n        self.batchnorm = nn.BatchNorm2d(out_channels) if use_batchnorm else None\n        self.activation = nn.LeakyReLU(0.2, inplace=True) if activation is None else activation()\n\n    def forward(self, x):\n        x = self.conv(x)\n        if self.batchnorm:\n            x = self.batchnorm(x)\n        x = self.activation(x)\n        return x\n\n\nclass DCGAN_GeneratorBlock(nn.Module):\n    \"\"\"\n    A generator convolutional block.\n    Default stride and padding double the size of features,\n    e.g. if input is [in_channels, 32, 32], output will be [out_channels, 64, 64].\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1,\n                 use_batchnorm=True, use_spectralnorm=False, activation=None):\n        super().__init__()\n\n        self.convT = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size,\n                                        stride=stride, padding=padding, bias=False)\n        if use_spectralnorm:\n            self.convT = nn.utils.spectral_norm(self.convT)\n        self.batchnorm = nn.BatchNorm2d(out_channels) if use_batchnorm else None\n        self.activation = nn.LeakyReLU(0.2, inplace=True) if activation is None else activation() # XXX: ReLU?\n\n    def forward(self, x):\n        x = self.convT(x)\n        if self.batchnorm:\n            x = self.batchnorm(x)\n        x = self.activation(x)\n        return x\n\n\nclass DCGAN_Discriminator(nn.Module):\n    \"\"\"The discriminator of a DCGAN\"\"\"\n\n    def __init__(self,\n                 num_latents=1,\n                 num_features=64,\n                 image_channels=3,\n                 image_size=64,\n                 max_features=512,\n                 gan_type=\"gan\",\n                 fully_convolutional=True,\n                 activation=None,\n                 use_batchnorm=True,\n                 use_spectralnorm=False,\n                 D_block=DCGAN_DiscriminatorBlock):\n        super().__init__()\n\n        using_grad_penalty = gan_type in (\"gan-gp\", \"wgan-gp\")\n        output_sigmoid = gan_type in (\"gan\", \"gan-gp\")\n\n        block_config = {\n            \"activation\": activation,\n            \"use_batchnorm\": use_batchnorm and not using_grad_penalty,\n            \"use_spectralnorm\": use_spectralnorm,\n        }\n\n        # Calculate intermediate image sizes\n        image_sizes = [image_size]\n        while image_sizes[-1] > 5:\n            image_sizes.append(image_sizes[-1] // 2)\n        latent_kernel = image_sizes[-1]  # should be either 3, 4, or 5\n        num_layers = len(image_sizes) - 1\n\n        # Calculate feature sizes\n        features = [min(num_features * 2**i, max_features) for i in range(num_layers)]\n\n        # Input layer\n        self.input_layer = D_block(image_channels, features[0], **block_config)\n\n        # Intermediate layers\n        self.main_layers = nn.Sequential(*[\n            D_block(in_features, out_features, **block_config)\n            for in_features, out_features in zip(features, features[1:])\n        ])\n\n        # Output layer (feature_size = 3, 4, or 5 -> 1)\n        if fully_convolutional:\n            self.output_layer = nn.Sequential(\n                nn.Conv2d(features[-1], num_latents, latent_kernel, bias=False),\n                View(-1),\n            )\n        else:\n            self.output_layer = nn.Sequential(\n                View(-1),\n                nn.Linear(features[-1] * latent_kernel**2, num_latents, bias=False)\n            )\n\n        # Add sigmoid activation if using regular GAN loss\n        self.output_activation = nn.Sigmoid() if output_sigmoid else None\n\n    def forward(self, x):\n        x = self.input_layer(x)\n        x = self.main_layers(x)\n        x = self.output_layer(x)\n        if self.output_activation:\n            x = self.output_activation(x)\n        # Remove H and W dimensions, infer channels dim (remove if 1)\n        x = x.view(x.size(0), -1).squeeze(1)\n        return x\n\n\nclass DCGAN_Generator(nn.Module):\n    \"\"\"The generator of a DCGAN\"\"\"\n\n    def __init__(self,\n                 num_latents=100,\n                 num_features=64,\n                 image_channels=3,\n                 image_size=64,\n                 max_features=512,\n                 gan_type=\"gan\",\n                 fully_convolutional=True,\n                 activation=None,\n                 use_batchnorm=True,\n                 use_spectralnorm=False,\n                 G_block=DCGAN_GeneratorBlock):\n        super().__init__()\n\n        block_config = {\n            \"activation\": activation,\n            \"use_batchnorm\": use_batchnorm,\n            \"use_spectralnorm\": use_spectralnorm\n        }\n\n        # Calculate intermediate image sizes\n        image_sizes = [image_size]\n        while image_sizes[-1] > 5:\n            image_sizes.append(image_sizes[-1] // 2)\n        latent_kernel = image_sizes[-1]  # should be either 3, 4, or 5\n        num_layers = len(image_sizes) - 1\n\n        # Calculate feature sizes\n        features = [min(num_features * 2**i, max_features) for i in range(num_layers)]\n\n        # Reverse order of image sizes and features for generator\n        image_sizes = image_sizes[::-1]\n        features = features[::-1]\n\n        # Input layer\n        if fully_convolutional:\n            self.input_layer = G_block(num_latents, features[0], kernel_size=latent_kernel,\n                                       stride=1, padding=0, **block_config)\n        else:\n            self.input_layer = nn.Sequential(\n                View(-1),\n                nn.Linear(num_latents, features[0] * image_sizes[0]**2, bias=False),\n                View(features[0], image_sizes[0], image_sizes[0])\n            )\n\n        # Intermediate layers\n        self.main_layers = nn.Sequential(*[\n            G_block(in_features, out_features, kernel_size=4+(expected_size%2), **block_config)\n            for in_features, out_features, expected_size in zip(features, features[1:], image_sizes[1:])\n        ])\n\n        # Output layer\n        self.output_layer = nn.ConvTranspose2d(features[-1], image_channels, kernel_size=4+(image_size%2),\n                                               stride=2, padding=1, bias=False)\n        self.output_activation = nn.Tanh()\n\n    def forward(self, x):\n        # Add H and W dimensions, infer channels dim (add if none)\n        x = x.view(x.size(0), -1, 1, 1)\n        x = self.input_layer(x)\n        x = self.main_layers(x)\n        x = self.output_layer(x)\n        x = self.output_activation(x)\n        return x\n\n\nclass View(nn.Module):\n    def __init__(self, *shape, including_batch=False):\n        super().__init__()\n        self.shape = shape\n        self.including_batch = including_batch\n    \n    def forward(self, x):\n        if self.including_batch:\n            return x.view(*self.shape)\n        else:\n            return x.view(x.size(0), *self.shape)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class RAE(nn.Module):\n    \"\"\"\n    Regularized Auto-encoders.\n    \"\"\"\n    def __init__(self, encoder=None, decoder=None, noise_std=0.01, *args, **kwargs):\n        super().__init__()\n        self.encoder = DCGAN_Discriminator(gan_type=\"wgan-gp\", *args, **kwargs) if encoder is None else encoder\n        self.decoder = DCGAN_Generator(*args, **kwargs) if decoder is None else decoder\n        self.noise_std = noise_std\n    \n    def sample_latent(self, x):\n        device = x.device\n        with torch.no_grad():\n            z = self.encoder(x)\n        # Just in case\n        z = z.view(z.size(0), -1)\n        batch_size, num_latents = z.size()\n        # The covariance of z.t()\n        mu = z.mean(dim=0, keepdim=True)\n        z -= mu # normalize\n        z_cov = z.t() @ z / (batch_size-1)  # (num_latents, num_latents)\n        z_cov += 1e-5 * torch.eye(num_latents).to(z)  # to avoid singular matrix\n        sampler = MultivariateNormal(mu.cpu().squeeze(), z_cov.cpu())\n        samples = torch.stack([sampler.sample() for _ in range(batch_size)])\n        return samples.to(device)\n\n    def reconstruct(self, z):\n        with torch.no_grad():\n            x_re = self.decoder(z)\n        return x_re\n    \n    def sample(self, x):\n        z = self.sample_latent(x)\n        x_re = self.reconstruct(z)\n        return x_re, z\n    \n    def forward(self, x, noise_std=None):\n        noise_std = self.noise_std if noise_std is None else noise_std\n        z = self.encoder(x)\n        z = z + torch.randn_like(z) * noise_std\n        x_re = self.decoder(z)\n        return x_re, z","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_transform():\n    \"\"\"\n    Make data transform.\n    \"\"\"\n    # Original image size is 400x400\n    return transforms.Compose([\n        #transforms.RandomCrop(300),\n        transforms.Resize(100),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomVerticalFlip(),\n        #transforms.RandomGrayscale(),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ])\n\nclass RecursionDataset(data_utils.Dataset):\n    def __init__(self,\n                 train=True,\n                 img_dir=\"/kaggle/input/recursion-simple\",\n                 transform=make_transform()):\n\n        self.mode = \"train\" if train else \"test\"\n        self.img_dir = img_dir\n        self.transform = transform\n\n        df = pd.read_csv(os.path.join(img_dir, f\"new_{self.mode}.csv\"))\n        self.records = df.to_records(index=False)\n        \n    def __getitem__(self, index):\n        img_path = os.path.join(\n            self.img_dir, self.mode, self.mode, self.records[index].filename)\n        img = Image.open(img_path)\n        img = self.transform(img)\n        sirna = int(self.records[index].sirna) if self.mode == \"train\" else -1\n        sirna = torch.tensor(sirna)\n        return (img, sirna)\n\n    def __len__(self):\n        return len(self.records)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### MNIST ###\nimage_size = 100  # original 400x400\nimage_channels = 3\n\n### Hyperparameters ###\nnum_classes = 1108\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nbatch_size = 128\nnum_workers = 32\nnum_latents = num_classes*2\nnum_features = 64\nnum_epochs = 30\nreport_interval = 200\nuse_spectralnorm = False\nuse_gradnorm = True\nz_sparsity = True\noptimizer = torch.optim.Adam\nlr = 1e-4","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = RecursionDataset(train=True)\ntest_dataset  = RecursionDataset(train=False)\ntrainloader = data_utils.DataLoader(train_dataset, batch_size=batch_size,\n                                     shuffle=True, num_workers=num_workers)\ntestloader = data_utils.DataLoader(test_dataset, batch_size=batch_size,\n                                    shuffle=False, num_workers=num_workers)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Dataset, model, optimizer\nrae = RAE(image_channels=image_channels, image_size=image_size,\n          num_latents=num_latents, num_features=num_features)\noptim = optimizer(rae.parameters(), lr=lr)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Losses\ndef reconstruction_loss(x, y):\n    return 0.5 * (x - y).pow(2).flatten(start_dim=1).sum(dim=1).mean()\n    #return F.mse_loss(x, y)\n\ndef rae_loss(z):\n    return 0.5 * z.pow(2).flatten(start_dim=1).sum(dim=1).mean()\n    #return F.mse_loss(z, torch.zeros_like(z))\n\ndef grad_norm(x_re, z):\n    grad = torch.autograd.grad(x_re, z, torch.ones_like(x_re), retain_graph=True)\n    grad_norm = grad[0].flatten(start_dim=1).norm(dim=1)\n    return grad_norm\n\ndef gradnorm_penalty(x_re, z):\n    return grad_norm(x_re, z).mean()\n\ndef num_parameters(model):\n    return sum(p.numel() for p in model.parameters())\n\nprint(\"Number of parameters =\", num_parameters(rae))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"losses = []\nreconstructions = []","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n### Train ###\nrae.train()\nrae = rae.to(device)\nfixed_x, _ = next(iter(testloader))\nfixed_x = fixed_x.to(device)\n\nfor epoch in range(num_epochs):\n    for iters, data in enumerate(trainloader):\n        optim.zero_grad()\n        x, y = data\n        x, y = x.to(device), y.to(device)\n        x_re, z = rae(x)\n\n        loss = reconstruction_loss(x_re, x)\n        if use_gradnorm:\n            loss += 0.02 * gradnorm_penalty(x_re, z)\n        if z_sparsity:\n            loss += z.abs().sum(dim=-1).mean()\n        loss += rae_loss(z)\n        loss.backward()\n        optim.step()\n\n        if (iters+1) % report_interval == 0:\n            progress = f\"[{epoch+1}/{num_epochs}][{iters+1}/{len(trainloader)}]\"\n            loss_stats = f\"Loss = {loss.item():.4f}\"\n            print(progress, loss_stats)\n            with torch.no_grad():\n                fixed_x_re, _ = rae(fixed_x)\n                reconstruction = torch.cat([fixed_x, fixed_x_re], dim=0).detach().cpu()\n            reconstructions.append(reconstruction)\n            losses.append(loss.item())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Using log scale to examine convergence\n_ = plt.plot(np.log(np.array(losses)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LatentClassifier(nn.Module):\n    def __init__(self, in_features=num_latents, out_features=num_classes, hidden_features=256, n_hidden=0):\n        super().__init__()\n        hidden_layers = []\n        for _ in range(n_hidden):\n            hidden_layers += [nn.Linear(hidden_features, hidden_features), nn.LeakyReLU(0.2)]\n        self.main = nn.Sequential(\n            nn.Linear(in_features, hidden_features),\n            *hidden_layers,\n            nn.Linear(hidden_features, out_features)\n        )\n\n    def forward(self, x):\n        return self.main(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainloader = data_utils.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)\ntestloader = data_utils.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n\nn_hidden = 8\nhidden_features = 512\nclassifier = LatentClassifier(in_features=num_latents, out_features=num_classes,\n                              hidden_features=num_classes, n_hidden=n_hidden)\nclassifier_optim = torch.optim.Adam(classifier.parameters())\nclassifier_loss = nn.CrossEntropyLoss()\nnum_parameters(classifier)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def topk_accuracy(pred_y, y, k=3):\n    topk_accuracies = [0] * (k+1)\n    with torch.no_grad():\n        topk = torch.topk(pred_y, k=k)\n        for kk in range(1, k+1):\n            hits = torch.zeros(pred_y.size(0)).byte()\n            for k_i in range(kk):\n                hits |= topk.indices[:, k_i].int() == y.int()\n            topk_accuracies[kk] = hits.float().mean().item()\n        return topk_accuracies","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n### Train ###\nclassifier.train()\nclassifier = classifier.to(device)\nclassifier_losses = []\nfor epoch in range(num_epochs):\n    for iters, data in enumerate(trainloader):\n        classifier_optim.zero_grad()\n        x, y = data\n        x, y = x.to(device), y.to(device)\n        with torch.no_grad():\n            z = rae.encoder(x)\n        y_pred = classifier(z)\n\n        loss = classifier_loss(y_pred, y)\n        loss.backward()\n        classifier_optim.step()\n\n        if (iters+1) % report_interval == 0:\n            progress = f\"[{epoch+1}/{num_epochs}][{iters+1}/{len(trainloader)}]\"\n            loss_stats = f\"Loss = {loss.item():.4f}\"\n            hits = (y_pred.argmax(dim=1) == y).sum()\n            accuracy_stats = f\"Accuracy = {hits.float() / y.size(0) * 100 :.2f}%\"\n            print(f\"{progress} {loss_stats}, {accuracy_stats}\")\n            classifier_losses.append(loss.item())\n\n            accuracies = topk_accuracy(y_pred.cpu(), y.cpu(), k=10)\n            [print(f\"- {k+1}: {100*acc:.2f}%\") for k, acc in enumerate(accuracies[1:])]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Using log scale to examine convergence\n_ = plt.plot(np.log(np.array(classifier_losses)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}