{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Create a Variational Autoencoder with pytorch (generate new samples)","metadata":{}},{"cell_type":"markdown","source":"*Autoencoders are a type of artificial neural network used in unsupervised machine learning and dimensionality reduction tasks. They are primarily designed for feature learning, data compression, and reconstruction. The core idea behind autoencoders is to encode input data into a compact representation (latent space) and then decode it back to the original data. This process allows the model to learn meaningful features and patterns in the input data. Autoencoders consist of two main components: an encoder and a decoder.*\n\nTypes of Autoencoders:\n\n- **Vanilla Autoencoder**: The simplest form of autoencoder, consisting of an encoder network to compress the input data into a lower-dimensional representation and a decoder network to reconstruct the original data from this representation.\n\n- **Denoising Autoencoder**: This type of autoencoder is trained to remove noise or corruption from input data. It helps in learning robust representations by encoding and decoding noisy versions of the input.\n\n- **Sparse Autoencoder**: Sparse autoencoders encourage the latent representations to be sparse, meaning that only a few units in the encoding layer are active at a time. This can help in feature selection and dimensionality reduction.\n\n- **Contractive Autoencoder**: These autoencoders include a regularization term in the loss function to make the model more robust to small perturbations in the input data, thus promoting stable representations.\n\n- **Variational Autoencoder (VAE)**: VAE is a specific type of autoencoder that introduces probabilistic elements into the latent space. It allows for the generation of new samples from the latent space and is particularly useful for generating data samples that are similar to the training data.\n\n- **Variational Autoencoders (VAE)**:\nVariational Autoencoders are a type of autoencoder that leverages probabilistic concepts to learn a probabilistic latent space. In VAE, instead of directly encoding data into a fixed latent representation, the encoder outputs the parameters of a probability distribution (usually a Gaussian distribution) in the latent space.\n\nThe key components of VAE include:\n\n- **Encoder**: This part of the network maps the input data to the mean and variance of a multivariate Gaussian distribution in the latent space. It effectively models the distribution of the latent variables.\n\n- **Reparameterization Trick**: To ensure that the network remains differentiable and can be trained using backpropagation, the sampling from the distribution is performed using a reparameterization trick, which involves sampling from a unit Gaussian and scaling and shifting the sample using the mean and variance from the encoder.\n\n- **Decoder**: The decoder takes samples from the latent space (which follow the learned distribution) and reconstructs the data.\n\nVAEs enable the generation of new samples from the probabilistic latent space by sampling points from the distribution and decoding them. This ability is what sets VAEs apart from other autoencoders. Since VAEs model the distribution of the latent space, they can generate new, meaningful data samples by sampling from the learned distribution. This is in contrast to traditional autoencoders, which produce fixed latent representations that cannot be readily used for generative tasks. VAEs are particularly useful for tasks like image generation and data generation in a controlled manner.","metadata":{}},{"cell_type":"markdown","source":"![](https://learnopencv.com/wp-content/uploads/2020/11/vae-diagram-1-1024x563.jpg)","metadata":{}},{"cell_type":"markdown","source":"## why variational autoencoders\n\n1. Probabilistic Latent Space: VAEs introduce a probabilistic approach to the latent space. Instead of producing a fixed, deterministic encoding like traditional autoencoders, VAEs model the latent space as a probability distribution. This allows VAEs to capture the inherent uncertainty and variations in the data, making them more robust and flexible in modeling complex, high-dimensional data.\n\n2. Generative Modeling: VAEs enable generative modeling, meaning they can generate new data samples that are similar to the training data. This is achieved by sampling from the learned latent space distribution and decoding the samples into new data points. VAEs have been used extensively for tasks like image generation, data synthesis, and content creation, making them essential for creative and generative applications.\n\n3. Data Generation with Controlled Properties: VAEs offer the ability to control the properties of generated data. By manipulating the latent variables, one can influence the characteristics of generated samples. This makes VAEs valuable in scenarios where you want to create data with specific attributes or styles, such as generating images with different poses or generating text with varying attributes.\n\n4. Semi-Supervised Learning: VAEs can be employed in semi-supervised learning settings. With labeled and unlabeled data, VAEs can use the probabilistic latent space to model the uncertainty in data labeling. This can improve the performance of models in scenarios where labeled data is limited.\n\n4. Anomaly Detection: VAEs can be used for anomaly detection by measuring how well the model can reconstruct a data point. Unusual or anomalous data points tend to have higher reconstruction errors, making VAEs valuable for detecting outliers in data.\n\n5. Interpolation and Smooth Transitions: VAEs allow for smooth interpolation between data points in the latent space. This means you can create smooth transitions between data samples, which is valuable for applications like morphing between images, style transfer, or generating new content that blends characteristics of existing data points.\n\n6. Dimensionality Reduction and Feature Learning: Like traditional autoencoders, VAEs can be used for dimensionality reduction and feature learning. However, the probabilistic nature of VAEs can capture more complex and nuanced features in the data.\n\n7. Regularization and Robustness: The probabilistic structure of VAEs acts as a form of regularization, making the model more robust and less prone to overfitting. This is especially beneficial in scenarios where data is noisy or limited.\n\n> ***without the need of expensive iterative inference schemes (such as MCMC) per datapoint. The learned approximate posterior inference model can also be used for a host of tasks such as recognition, denoising, representation and visualization purposes.***","metadata":{}},{"cell_type":"code","source":"!pip install torchsummary torchview","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:13:41.143980Z","iopub.execute_input":"2023-10-23T03:13:41.144903Z","iopub.status.idle":"2023-10-23T03:13:53.610056Z","shell.execute_reply.started":"2023-10-23T03:13:41.144867Z","shell.execute_reply":"2023-10-23T03:13:53.609120Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport csv\nimport cv2 as cv\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torchvision\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport seaborn as sns\nimport shutil\nimport time\n\nfrom torchvision import transforms\nfrom torchvision.transforms import v2\nfrom sklearn.metrics import accuracy_score\nfrom matplotlib import pyplot as plt\nfrom tqdm.notebook import tqdm\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torch import nn\nfrom torchsummary import summary\nfrom torchview import draw_graph\nfrom warnings import filterwarnings\n\nfilterwarnings(action='ignore')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-10-23T03:13:53.612354Z","iopub.execute_input":"2023-10-23T03:13:53.612680Z","iopub.status.idle":"2023-10-23T03:13:58.269726Z","shell.execute_reply.started":"2023-10-23T03:13:53.612651Z","shell.execute_reply":"2023-10-23T03:13:58.268819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Utils","metadata":{}},{"cell_type":"markdown","source":"this notebook use MNIST dataset to do the implementation of the variational autoencoder and theses are the pre-processing functions to convert csv image data into real images and save in the working directory to use in the dataloaders while training and validation steps later.","metadata":{}},{"cell_type":"code","source":"def read_data(filename):\n    with open(filename, \"r\") as file:\n        reader = csv.reader(file)\n        reader.__next__()\n\n        for row in reader:\n            yield np.array(row, dtype=float)\n\n\ndef get_label_and_img(array):\n    return array[1:], array[0]\n\n\ndef convert_to_img(array, filename):\n    img = array.reshape((28, 28))\n    cv.imwrite(filename, img)\n\n\ndef do_preprocess(filename, dst_dir):\n    Y = []\n\n    for i, arr in enumerate(read_data(filename)):\n        img, label = get_label_and_img(arr)\n        file = f'img_{i}.jpg'\n\n        Y.append([file, label])\n        dst = os.path.join(dst_dir, file)\n        convert_to_img(img, dst)\n\n    df = pd.DataFrame(Y, columns=['file', 'label'])\n    df.to_csv('label.csv', index=False)","metadata":{"_kg_hide-input":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-10-23T03:13:58.271010Z","iopub.execute_input":"2023-10-23T03:13:58.271572Z","iopub.status.idle":"2023-10-23T03:13:58.280671Z","shell.execute_reply.started":"2023-10-23T03:13:58.271537Z","shell.execute_reply":"2023-10-23T03:13:58.279642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create seperate folders to keep training and testing data\n\nif not os.path.isdir('./train'):\n    os.mkdir('./train')\n    \nif not os.path.isdir('./test'):\n    os.mkdir('./test')\n\ndo_preprocess('../input/digit-recognizer/train.csv', './train')","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:13:58.281800Z","iopub.execute_input":"2023-10-23T03:13:58.282063Z","iopub.status.idle":"2023-10-23T03:14:11.031962Z","shell.execute_reply.started":"2023-10-23T03:13:58.282040Z","shell.execute_reply":"2023-10-23T03:14:11.030799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# select and move randomly selected sample of training images to test folder\n\nimg_files = os.listdir('./train/')\ntest_sample = np.random.choice(img_files, size=256, replace=False)\n\nfor file in tqdm(test_sample):\n    curr = os.path.join('./train', file)\n    new = os.path.join('./test', file)\n\n    shutil.move(curr, new)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.035305Z","iopub.execute_input":"2023-10-23T03:14:11.035695Z","iopub.status.idle":"2023-10-23T03:14:11.114753Z","shell.execute_reply.started":"2023-10-23T03:14:11.035660Z","shell.execute_reply":"2023-10-23T03:14:11.113887Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"labels = pd.read_csv('./label.csv')\ntest_set = labels[labels['file'].isin(test_sample)]\nlabels.drop(test_set.index, axis=0).to_csv('./label.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.115872Z","iopub.execute_input":"2023-10-23T03:14:11.116140Z","iopub.status.idle":"2023-10-23T03:14:11.279161Z","shell.execute_reply.started":"2023-10-23T03:14:11.116117Z","shell.execute_reply":"2023-10-23T03:14:11.278372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def move_to(obj, device):\n    \"\"\"this is utility function to move tensor objects like tensors, models\n    to specified device\"\"\"\n    \n    if isinstance(obj, list):\n        return [move_to(x, device) for x in obj]\n    \n    elif isinstance(obj, tuple):\n        return tuple(move_to(obj, device))\n    \n    elif isinstance(obj, set):\n        return set(move_to(obj, device))\n    \n    elif isinstance(obj, dict):\n        to_return = dict()\n        \n        for key, value in obj.items():\n            to_return[move_to(key, device)] = move_to(value, device)\n        \n        return to_return\n    \n    elif hasattr(obj, 'to'):\n        return obj.to(device)\n    \n    else:\n        return obj","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.280273Z","iopub.execute_input":"2023-10-23T03:14:11.280562Z","iopub.status.idle":"2023-10-23T03:14:11.288002Z","shell.execute_reply.started":"2023-10-23T03:14:11.280538Z","shell.execute_reply":"2023-10-23T03:14:11.287140Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_epoch(model, optimizer, data_loader,\n              loss_func, device, results,\n              score_funcs, prefix=\"\", desc=None, keep=False, epoch=None):\n    \"\"\"this will run a single epoch on training data and calculate loss\n    then use that to perform backpropergation and update weights accordingly\"\"\"\n    running_loss = []\n    y_true = []\n    y_pred = []\n    start = time.time()\n\n    for inputs, labels in tqdm(data_loader, desc=desc, leave=keep):\n        # Move the batch to the device we are using.\n        inputs = move_to(inputs, device)\n        labels = move_to(labels, device)\n\n        y_hat, mu, sigma = model(inputs)  # this just computed f_Θ(x(i))\n        \n        # Compute loss.\n        reconstructed_loss = loss_func(y_hat, labels)\n        kdl_ele = mu.pow(2).add_(sigma.exp()).mul_(-1).add_(1).add_(sigma)\n        kl_div = torch.sum(kdl_ele).mul_(-0.5)\n        loss = reconstructed_loss + kl_div\n\n        if model.training:\n            loss.backward()\n            optimizer.step()\n            optimizer.zero_grad()\n\n        # Now we are just grabbing some information we would like to have\n        running_loss.append(loss.item())\n\n        if len(score_funcs) > 0 and isinstance(labels, torch.Tensor):\n            # moving labels & predictions back to CPU for computing / storing predictions\n            labels = labels.detach().cpu().numpy()\n            y_hat = y_hat.detach().cpu().numpy()\n            # add to predictions so far\n            y_true.extend(labels.tolist())\n            y_pred.extend(y_hat.tolist())\n\n    # end training epoch\n    end = time.time()\n\n    y_pred = np.asarray(y_pred)\n    if len(y_pred.shape) == 2 and y_pred.shape[1] > 1:  # We have a classification problem, convert to labels\n        y_pred = np.argmax(y_pred, axis=1)\n    # Else, we assume we are working on a regression problem\n\n    results[prefix + \" loss\"].append(np.mean(running_loss))\n    result_str = [f\"{prefix} loss: {np.mean(running_loss)}\"]\n\n    for name, score_func in score_funcs.items():\n        try:\n            score = score_func(y_true, y_pred)\n\n            results[prefix + \" \" + name].append(score)\n            result_str.append(f\"{prefix} {name}: {score}\")\n\n            if prefix == 'val' or prefix == 'test':\n                checkpointer(score, (epoch + 1), model, optimizer)\n\n        except:\n            results[prefix + \" \" + name].append(float(\"NaN\"))\n            result_str.append(f\"{prefix} {name}: {float('NaN')}\")\n\n    print(\" \".join(result_str))\n\n    return end - start  # time spent on epoch","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:02:17.397996Z","iopub.execute_input":"2023-10-23T04:02:17.398371Z","iopub.status.idle":"2023-10-23T04:02:17.413461Z","shell.execute_reply.started":"2023-10-23T04:02:17.398340Z","shell.execute_reply":"2023-10-23T04:02:17.412352Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_loop(model, loss_func,\n               train_loader, test_loader=None, val_loader=None,\n               score_funcs=None, epochs=50, device=\"cpu\",\n               optimizer=None, lr_schedule=None, checkpoint_file=None, keep=True):\n    \"\"\"full training loop to train a neural network from scratch. this will run through a\n    specific number of epochs and calculate scores by using given score functions\"\"\"\n    if score_funcs is None:\n        score_funcs = {}\n\n    to_track = [\"epoch\", \"total time\", \"train loss\", \"lr\"]\n\n    if val_loader is not None:\n        to_track.append(\"val loss\")\n\n    if test_loader is not None:\n        to_track.append(\"test loss\")\n\n    for eval_score in score_funcs:\n        to_track.append(\"train \" + eval_score)\n\n        if val_loader is not None:\n            to_track.append(\"val \" + eval_score)\n\n        if test_loader is not None:\n            to_track.append(\"test \" + eval_score)\n\n    total_train_time = 0  # How long have we spent in the training loop?\n    results = {}\n\n    # Initialize every item with an empty list\n    for item in to_track:\n        results[item] = []\n\n    if optimizer is None:\n        optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\n        del_opt = True\n\n    else:\n        del_opt = False\n\n    # Place the model on the correct computed resource (CPU or GPU)\n    model.to(device)\n\n    for epoch in tqdm(range(epochs), desc=\"Epoch\", leave=keep):\n        model = model.train()  # Put our model in training mode\n\n        run_epoch(model, optimizer, train_loader, loss_func, device, results, score_funcs,\n                  prefix=\"train\", desc=\"Training\")\n\n        results[\"total time\"].append(total_train_time)\n        results[\"epoch\"].append(epoch)\n        results[\"lr\"].append(optimizer.param_groups[0]['lr'])\n\n        if val_loader is not None:\n            model = model.eval()\n\n            with torch.no_grad():\n                run_epoch(model, optimizer, val_loader, loss_func, device, results, score_funcs, prefix=\"val\",\n                          desc=\"Validation\", epoch=epoch)\n\n        if test_loader is not None:\n            model = model.eval()\n\n            with torch.no_grad():\n                run_epoch(model, optimizer, test_loader, loss_func, device, results, score_funcs, prefix=\"test\",\n                          desc=\"Testing\", epoch=epoch)\n\n        if lr_schedule is not None:\n            if isinstance(lr_schedule, torch.optim.lr_scheduler.ReduceLROnPlateau):\n                lr_schedule.step(results[\"val loss\"][-1])\n\n            else:\n                lr_schedule.step()\n\n        if checkpoint_file is not None:\n            torch.save({\n                'epoch': epoch,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'results': results\n            }, checkpoint_file)\n\n    if del_opt:\n        del optimizer\n\n    return pd.DataFrame.from_dict(results)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:02:20.728356Z","iopub.execute_input":"2023-10-23T04:02:20.729456Z","iopub.status.idle":"2023-10-23T04:02:20.744578Z","shell.execute_reply.started":"2023-10-23T04:02:20.729415Z","shell.execute_reply":"2023-10-23T04:02:20.743572Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Load data ","metadata":{}},{"cell_type":"markdown","source":"loading image data from training directory to pass for the training of the neural netowork. in here we use few image data augmentaions to improve model performance by generalizing it to different image inputs. to autoencoders we need to pass original training image as the ground truth to calculate the loss, so we return input image as lebel in the data loader.","metadata":{}},{"cell_type":"code","source":"class AutoEncodeDataset(Dataset):\n    \"\"\"loading input as output for x,x tuples in training data\"\"\"\n\n    def __init__(self, annotation_file, img_dir, transform=None):\n        self.img_labels = pd.read_csv(annotation_file)\n        self.img_dir = img_dir\n        self.transform = transform\n        self.resize = v2.Compose([v2.Resize(size=27),\n                                  transforms.ToTensor()])\n\n    def __len__(self):\n        return len(self.img_labels)\n\n    def __getitem__(self, idx):\n        img_path = os.path.join(self.img_dir, self.img_labels.loc[idx, 'file'])\n        image = torchvision.io.read_image(img_path).to(torch.float32) / 255\n\n        if self.transform is not None:\n            image = transforms.ToPILImage()(image)\n            image = self.transform(image)\n            \n        label = transforms.ToPILImage()(image)\n        label = self.resize(label)\n\n        return image, label","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.323338Z","iopub.execute_input":"2023-10-23T03:14:11.323616Z","iopub.status.idle":"2023-10-23T03:14:11.343610Z","shell.execute_reply.started":"2023-10-23T03:14:11.323593Z","shell.execute_reply":"2023-10-23T03:14:11.342635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# data augmentation steps to use inside dataloders\n\ntrain_transform = v2.Compose([v2.RandomAffine(degrees=5, translate=(0.05, 0.05), scale=(0.98, 1.02)),\n                              v2.RandomPerspective(distortion_scale=0.5, p=0.5),\n                              transforms.ToTensor()])","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.344834Z","iopub.execute_input":"2023-10-23T03:14:11.345131Z","iopub.status.idle":"2023-10-23T03:14:11.355488Z","shell.execute_reply.started":"2023-10-23T03:14:11.345101Z","shell.execute_reply":"2023-10-23T03:14:11.354612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"training_data = AutoEncodeDataset('./label.csv', './train/', transform=train_transform)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.356433Z","iopub.execute_input":"2023-10-23T03:14:11.356702Z","iopub.status.idle":"2023-10-23T03:14:11.394269Z","shell.execute_reply.started":"2023-10-23T03:14:11.356673Z","shell.execute_reply":"2023-10-23T03:14:11.393277Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_size = int(len(training_data) * 0.8)\ntest_size = len(training_data) - train_size\n\ntrain_dataset, test_dataset = random_split(training_data, (train_size, test_size))\n\ntrain_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)\ntest_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.395733Z","iopub.execute_input":"2023-10-23T03:14:11.396126Z","iopub.status.idle":"2023-10-23T03:14:11.426515Z","shell.execute_reply.started":"2023-10-23T03:14:11.396092Z","shell.execute_reply":"2023-10-23T03:14:11.425743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model building","metadata":{}},{"cell_type":"code","source":"# use cuda if it's available\n\nif torch.cuda.is_available():\n    device = torch.device('cuda')\n    \nelse:\n    device = torch.device('cpu')","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.431545Z","iopub.execute_input":"2023-10-23T03:14:11.431856Z","iopub.status.idle":"2023-10-23T03:14:11.458744Z","shell.execute_reply.started":"2023-10-23T03:14:11.431829Z","shell.execute_reply":"2023-10-23T03:14:11.457756Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:11.459980Z","iopub.execute_input":"2023-10-23T03:14:11.460337Z","iopub.status.idle":"2023-10-23T03:14:11.476551Z","shell.execute_reply.started":"2023-10-23T03:14:11.460304Z","shell.execute_reply":"2023-10-23T03:14:11.475643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Variational Autoencoder model","metadata":{}},{"cell_type":"code","source":"class Encode(nn.Module):\n\n    def __init__(self, channels, embeddings, kernel):\n        super(Encode, self).__init__()\n\n        self.relu = nn.ReLU()\n        self.filters = [32, 64, 128, 256, 512]\n\n        self.conv1 = nn.Conv2d(channels, self.filters[0], kernel_size=kernel)\n        self.bn1 = nn.BatchNorm2d(self.filters[0])\n        self.conv2 = nn.Conv2d(self.filters[0], self.filters[1], kernel_size=kernel, stride=2)\n        self.bn2 = nn.BatchNorm2d(self.filters[1])\n        self.conv3 = nn.Conv2d(self.filters[1], self.filters[2], kernel_size=kernel)\n        self.bn3 = nn.BatchNorm2d(self.filters[2])\n        self.conv4 = nn.Conv2d(self.filters[2], self.filters[3], kernel_size=kernel, stride=2)\n        self.bn4 = nn.BatchNorm2d(self.filters[3])\n        self.conv5 = nn.Conv2d(self.filters[3], self.filters[4], kernel_size=kernel)\n        self.bn5 = nn.BatchNorm2d(self.filters[4])\n\n        self.mu = nn.Linear(self.filters[4] * 4, embeddings)\n        self.sigma = nn.Linear(self.filters[4] * 4, embeddings)\n\n    def forward(self, x):\n        x = self.bn1(self.relu(self.conv1(x)))\n        x = self.bn2(self.relu(self.conv2(x)))\n        x = self.bn3(self.relu(self.conv3(x)))\n        x = self.bn4(self.relu(self.conv4(x)))\n        x = self.bn5(self.relu(self.conv5(x)))\n        x = torch.flatten(x, start_dim=1)\n\n        return self.mu(x), self.sigma(x)\n    \n    \nclass Decode(nn.Module):\n\n    def __init__(self, channels, embeddings, kernel):\n        super(Decode, self).__init__()\n\n        self.relu = nn.ReLU()\n        self.filters = [512, 256, 128, 64, 32]\n\n        self.decode_in = nn.Linear(embeddings, self.filters[0] * 4)\n        self.deconv1 = nn.ConvTranspose2d(self.filters[0], self.filters[1], kernel_size=kernel, stride=2, dilation=1, output_padding=1, padding=1)\n        self.bn1 = nn.BatchNorm2d(self.filters[1])\n        self.deconv2 = nn.ConvTranspose2d(self.filters[1], self.filters[2], kernel_size=kernel, stride=2, dilation=1, output_padding=0, padding=1)\n        self.bn2 = nn.BatchNorm2d(self.filters[2])\n        self.deconv3 = nn.ConvTranspose2d(self.filters[2], self.filters[3], kernel_size=kernel, stride=2, dilation=1, output_padding=1, padding=1)\n        self.bn3 = nn.BatchNorm2d(self.filters[3])\n        self.deconv4 = nn.ConvTranspose2d(self.filters[3], self.filters[4], kernel_size=kernel, stride=2, dilation=1, output_padding=0, padding=1)\n        self.bn4 = nn.BatchNorm2d(self.filters[4])\n        self.decode_out = nn.Conv2d(self.filters[4], 1, kernel_size=kernel, padding=1)\n\n    def forward(self, x):\n        x = self.decode_in(x)\n        x = x.view(-1, 512, 2, 2)\n        x = self.bn1(self.relu(self.deconv1(x)))\n        x = self.bn2(self.relu(self.deconv2(x)))\n        x = self.bn3(self.relu(self.deconv3(x)))\n        x = self.bn4(self.relu(self.deconv4(x)))\n        x = self.decode_out(x)\n\n        return x\n\n    \nclass ConvVariationalAutoEncoder(nn.Module):\n    \"\"\"variational autoencoder model\"\"\"\n\n    def __init__(self, channels, embeddings, kernel):\n        super(ConvVariationalAutoEncoder, self).__init__()\n\n        self.encoder = Encode(channels, embeddings, kernel)\n        self.decoder = Decode(channels, embeddings, kernel)\n\n    def forward(self, x):\n        mu, sigma = self.encoder(x)\n        \n        std = torch.exp(0.5 * sigma)\n        epsilon = torch.randn_like(std)\n        z = epsilon.mul(std).add_(mu)\n\n        return self.decoder(z), mu, sigma","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:01:20.106406Z","iopub.execute_input":"2023-10-23T04:01:20.106761Z","iopub.status.idle":"2023-10-23T04:01:20.131962Z","shell.execute_reply.started":"2023-10-23T04:01:20.106735Z","shell.execute_reply":"2023-10-23T04:01:20.130886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = ConvVariationalAutoEncoder(1, 20, 3)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:01:24.556799Z","iopub.execute_input":"2023-10-23T04:01:24.557637Z","iopub.status.idle":"2023-10-23T04:01:24.593056Z","shell.execute_reply.started":"2023-10-23T04:01:24.557604Z","shell.execute_reply":"2023-10-23T04:01:24.592257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Model summary","metadata":{}},{"cell_type":"code","source":"summary(model.cuda(), input_size=(1, 28, 28), batch_size=-1)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:01:27.555165Z","iopub.execute_input":"2023-10-23T04:01:27.555986Z","iopub.status.idle":"2023-10-23T04:01:27.576800Z","shell.execute_reply.started":"2023-10-23T04:01:27.555951Z","shell.execute_reply":"2023-10-23T04:01:27.575907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### model architecture","metadata":{}},{"cell_type":"code","source":"model_graph = draw_graph(model.cuda(), input_size=(1, 1, 28, 28), expand_nested=True)\nmodel_graph.visual_graph","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:19.202233Z","iopub.execute_input":"2023-10-23T03:14:19.202627Z","iopub.status.idle":"2023-10-23T03:14:19.429296Z","shell.execute_reply.started":"2023-10-23T03:14:19.202571Z","shell.execute_reply":"2023-10-23T03:14:19.428347Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"we use binary cross entropy as loss function here but also can use mean squared error function as well.","metadata":{}},{"cell_type":"code","source":"loss_func = nn.MSELoss(reduction='sum')","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:01:41.555186Z","iopub.execute_input":"2023-10-23T04:01:41.555800Z","iopub.status.idle":"2023-10-23T04:01:41.560036Z","shell.execute_reply.started":"2023-10-23T04:01:41.555768Z","shell.execute_reply":"2023-10-23T04:01:41.559059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"images, labels = next(iter(train_dataloader))\nfig, axes = plt.subplots(nrows=4, ncols=4, figsize=[13, 14], dpi=100)\naxes = axes.ravel()\n\nfor i, (img, label) in enumerate(zip(images, labels)):\n    axes[i].imshow(img.numpy()[0], cmap=plt.cm.gray)\n#     axes[i].set_title(label.numpy())\n    \n    if i >= 15:\n        break\n\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:19.436667Z","iopub.execute_input":"2023-10-23T03:14:19.436926Z","iopub.status.idle":"2023-10-23T03:14:21.631426Z","shell.execute_reply.started":"2023-10-23T03:14:19.436903Z","shell.execute_reply":"2023-10-23T03:14:21.630414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"labels[0].shape","metadata":{"execution":{"iopub.status.busy":"2023-10-23T03:14:21.632565Z","iopub.execute_input":"2023-10-23T03:14:21.633308Z","iopub.status.idle":"2023-10-23T03:14:21.639004Z","shell.execute_reply.started":"2023-10-23T03:14:21.633280Z","shell.execute_reply":"2023-10-23T03:14:21.638049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Training model","metadata":{}},{"cell_type":"code","source":"optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:01:49.997408Z","iopub.execute_input":"2023-10-23T04:01:49.998242Z","iopub.status.idle":"2023-10-23T04:01:50.003431Z","shell.execute_reply.started":"2023-10-23T04:01:49.998208Z","shell.execute_reply":"2023-10-23T04:01:50.002341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"results = train_loop(model, \n                     loss_func,\n                     train_loader=train_dataloader,\n                     val_loader=test_dataloader,\n                     optimizer=optimizer,\n                     lr_schedule=scheduler,\n                     device=device,\n                     epochs=100\n                    )","metadata":{"execution":{"iopub.status.busy":"2023-10-23T04:02:34.416900Z","iopub.execute_input":"2023-10-23T04:02:34.417613Z","iopub.status.idle":"2023-10-23T05:34:37.336000Z","shell.execute_reply.started":"2023-10-23T04:02:34.417579Z","shell.execute_reply":"2023-10-23T05:34:37.335043Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nplt.figure(figsize=[12, 6], dpi=200)\nsns.lineplot(x='epoch', y='train loss', data=results, label='train')\nsns.lineplot(x='epoch', y='val loss', data=results, label='validation')\nplt.ylabel('binary cross entropy loss')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-10-23T05:43:15.961003Z","iopub.execute_input":"2023-10-23T05:43:15.961493Z","iopub.status.idle":"2023-10-23T05:43:16.357003Z","shell.execute_reply.started":"2023-10-23T05:43:15.961458Z","shell.execute_reply":"2023-10-23T05:43:16.356027Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Inference","metadata":{}},{"cell_type":"code","source":"model = model.to('cpu')\n\n\ndef get_digits(dataset):\n    idx = torch.randint(len(dataset), (1,))\n    \n    return dataset[idx][0].unsqueeze(0)\n\n\ndef get_encodings(image):\n    with torch.no_grad():\n        mu, sigma = model.encoder(image)\n\n    return mu, sigma\n\ndef get_reconstructed(image):\n    with torch.no_grad():\n        recon, _, _ = model(image)\n        \n    return recon\n\n\ndef inference(*, dataset, n_examples=1):\n    out = []\n    image = get_digits(dataset)\n    mu, sigma = get_encodings(image)\n    recon_img = get_reconstructed(image)\n\n    for example in range(n_examples):\n        epsilon = torch.rand_like(sigma)\n        z = mu + sigma * epsilon\n\n        out.append(model.decoder(z))\n        \n    return image, out, recon_img","metadata":{"execution":{"iopub.status.busy":"2023-10-23T05:44:01.271180Z","iopub.execute_input":"2023-10-23T05:44:01.271559Z","iopub.status.idle":"2023-10-23T05:44:01.286542Z","shell.execute_reply.started":"2023-10-23T05:44:01.271531Z","shell.execute_reply":"2023-10-23T05:44:01.285760Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_digit, outs, recon_img = inference(dataset=test_dataset, n_examples=10)\n\nfig, axes = plt.subplots(nrows=2, ncols=6, layout='constrained', figsize=(14, 4))\ngridspec = axes[0, 0].get_subplotspec().get_gridspec()\n\nfor a in axes[:, 0]:\n    a.remove()\n    \nfor i, a in enumerate(axes[:, 1:].flat):\n    a.axis('off')\n    a.imshow(outs[i].detach().numpy()[0][0])\n    \nsubfig = fig.add_subfigure(gridspec[:, 0])\naxsLeft = subfig.subplots(1, 2, sharey=True)\n\naxsLeft[0].imshow(input_digit[0][0])\naxsLeft[0].set_title('original', fontsize=12)\naxsLeft[1].imshow(recon_img[0][0])\naxsLeft[1].set_title('reconstructed', fontsize=12)\n\naxsLeft[0].axis('off')\naxsLeft[1].axis('off')\n\nsubfig.suptitle('encoder-decoder', fontsize='x-large')\nfig.suptitle('generated samples', fontsize='xx-large')","metadata":{"execution":{"iopub.status.busy":"2023-10-23T06:05:38.693297Z","iopub.execute_input":"2023-10-23T06:05:38.694162Z","iopub.status.idle":"2023-10-23T06:05:39.413451Z","shell.execute_reply.started":"2023-10-23T06:05:38.694124Z","shell.execute_reply":"2023-10-23T06:05:39.412516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.rmtree(\"./train/\")\nshutil.rmtree(\"./test/\")","metadata":{"execution":{"iopub.status.busy":"2023-10-23T06:06:07.352962Z","iopub.execute_input":"2023-10-23T06:06:07.353712Z","iopub.status.idle":"2023-10-23T06:06:08.595953Z","shell.execute_reply.started":"2023-10-23T06:06:07.353681Z","shell.execute_reply":"2023-10-23T06:06:08.595088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = 'conv-var-autoencoder.pth'\ntorch.save(model.state_dict(), PATH)","metadata":{"execution":{"iopub.status.busy":"2023-10-23T06:06:25.202162Z","iopub.execute_input":"2023-10-23T06:06:25.202916Z","iopub.status.idle":"2023-10-23T06:06:25.239086Z","shell.execute_reply.started":"2023-10-23T06:06:25.202884Z","shell.execute_reply":"2023-10-23T06:06:25.238284Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## References\n\n- D. P. Kingma and M. Welling, “Auto-Encoding variational Bayes,” arXiv (Cornell University), Dec. 2013, doi: 10.48550/arxiv.1312.6114. \n- Aladdin Persson, “Variational AutoEncoder Paper Walkthrough,” YouTube. Sep. 12, 2022. [Online]. Available: https://www.youtube.com/watch?v=5bA6gwo36Cw \n- H. Bandyopadhyay, “Autoencoders in Deep Learning: Tutorial & Use Cases [2023],” V7, Jul. 11, 2023. https://www.v7labs.com/blog/autoencoders guide#:~:text=Generation%20of%20image%20and%20time%20series%20data&text=The%20parameterized%20distribution%20at%20the,to%20generation%20of%20image%20data.","metadata":{}}]}