autoencoder_mlp.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade matplotlib"
      ],
      "metadata": {
        "id": "XS2yGufz1Lif"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ORwLkEDKHQUY"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import os\n",
        "from numpy.random import shuffle\n",
        "#%%\n",
        "from tqdm import tqdm\n",
        "\n",
        "rand_seed = 1234\n",
        "torch.manual_seed(rand_seed)\n",
        "#%%\n",
        "params = {'N_units': 30,   # numero di dimensioni del vettore latente\n",
        "          'L_rate_AE': 4e-5,\n",
        "          'B_size': 8,\n",
        "          'Epochs': 1000,\n",
        "          'W_N_f_p': 1e-1,\n",
        "          'W_N_p_f': 1e-1,\n",
        "          'N_evals': 10,\n",
        "          'enc_layers': [512, 256],\n",
        "          'dec_layers': [256],\n",
        "          'AE_activation': 'tanh',\n",
        "          'datastep': 10,\n",
        "          'bat_n_all': True,\n",
        "          'actual_epoch': 0,\n",
        "          'N_layers': [64, 128, 256, 128, 64], # Pi Layers\n",
        "          }"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "partial_shapes = np.load('./partial_shapes_matrix.npy')\n",
        "full_shapes = np.load('./total_shapes_matrix.npy')\n",
        "\n",
        "pix_f = torch.tensor(full_shapes)\n",
        "pix_p = torch.tensor(partial_shapes)"
      ],
      "metadata": {
        "id": "unKEhx_gIDeD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "save_every = 100\n",
        "test_every = 5\n",
        "#%%\n",
        "\n",
        "def distance_matrix(array1, array2):\n",
        "    \"\"\"\n",
        "    arguments:\n",
        "        array1: the array, size: (batch_size, num_point, num_feature)\n",
        "        array2: the samples, size: (batch_size, num_point, num_feature)\n",
        "    returns:\n",
        "        distances: each entry is the distance from a sample to array1\n",
        "            , it's size: (batch_size, num_point, num_point)\n",
        "    \"\"\"\n",
        "    batch_size, num_point, num_features = array1.shape\n",
        "    expanded_array1 = torch.tile(array1, dims=(1, num_point, 1))\n",
        "    expanded_array2 = torch.reshape(\n",
        "        torch.tile(torch.unsqueeze(array2, 2),\n",
        "                   (1, 1, num_point, 1)),\n",
        "        (batch_size, -1, num_features))\n",
        "\n",
        "    distances = torch.linalg.norm(expanded_array1-expanded_array2, dim=-1)\n",
        "    distances = torch.reshape(distances, (batch_size, num_point, num_point))\n",
        "    return distances\n",
        "\n",
        "\n",
        "def av_dist(array1, array2):\n",
        "    \"\"\"\n",
        "    arguments:\n",
        "        array1, array2: both size: (batch_size, num_points, num_feature)\n",
        "    returns:\n",
        "        distances: size: (1,)\n",
        "    \"\"\"\n",
        "    distances = distance_matrix(array1, array2)\n",
        "    distances, _ = torch.min(distances, dim=-1)\n",
        "    distances = torch.mean(distances, dim=-1)\n",
        "    return distances\n",
        "\n",
        "def av_dist_sum(array1, array2):\n",
        "    \"\"\"\n",
        "    arguments:\n",
        "        arrays: array1, array2\n",
        "    returns:\n",
        "        sum of av_dist(array1, array2) and av_dist(array2, array1)\n",
        "    \"\"\"\n",
        "    av_dist1 = av_dist(array1, array2)\n",
        "    av_dist2 = av_dist(array2, array1)\n",
        "    return av_dist1+av_dist2\n",
        "\n",
        "def chamfer_distance(array1, array2):\n",
        "    return torch.mean(av_dist_sum(array1, array2))\n",
        "\n",
        "def save_decoded_shape(shape, epoch):\n",
        "    with open('shape{}.npy'.format(epoch), 'wb') as f:\n",
        "      np.save(f, shape)"
      ],
      "metadata": {
        "id": "TPnujc7j8AMc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#%%\n",
        "model_path = 'models' \n",
        "if not os.path.exists(model_path):\n",
        "  os.mkdir(model_path)\n",
        "#%%\n",
        "# load datasets\n",
        "\n",
        "class Encoder(torch.nn.Module):\n",
        "\n",
        "    def __init__(self, n_points):\n",
        "\n",
        "        super(Encoder, self).__init__()\n",
        "        self.n_points = n_points\n",
        "        enc_layers = params['enc_layers'].copy()\n",
        "        enc_layers.insert(0, self.n_points * 3)\n",
        "        self.layers = torch.nn.ModuleList([])\n",
        "        self.flat = torch.nn.Flatten()\n",
        "        for n in range(1, len(enc_layers)):\n",
        "            self.layers.append(torch.nn.Linear(enc_layers[n - 1], enc_layers[n]))\n",
        "            self.layers.append(torch.nn.Tanh())\n",
        "        self.out = torch.nn.Linear(enc_layers[-1], params['N_units'])\n",
        "\n",
        "    def forward(self, x):\n",
        "        x =  self.flat(x)\n",
        "        for layer in self.layers:\n",
        "            x = layer(x)\n",
        "        x = self.out(x) # (N, n_points, 3)\n",
        "        return x\n",
        "\n",
        "class Decoder(torch.nn.Module):\n",
        "\n",
        "    def __init__(self, n_points, latent_dim):\n",
        "        super(Decoder, self).__init__()\n",
        "        self.latent_dim = latent_dim\n",
        "        self.n_points = n_points\n",
        "        dec_layers = params['dec_layers'].copy()\n",
        "        dec_layers.insert(0, self.latent_dim)\n",
        "        self.layers = torch.nn.ModuleList([])\n",
        "        for n in range(1, len(dec_layers)):\n",
        "            self.layers.append(torch.nn.Linear(dec_layers[n - 1], dec_layers[n]))\n",
        "            self.layers.append(torch.nn.Tanh())\n",
        "        self.out = torch.nn.Linear(dec_layers[-1], self.n_points * 3)\n",
        "\n",
        "    def forward(self, x):\n",
        "        for layer in self.layers:\n",
        "            x = layer(x)\n",
        "        x = self.out(x)\n",
        "        x = x.reshape(-1, self.n_points, 3) # (N, n_points, 3)\n",
        "        return x\n",
        "\n",
        "\n",
        "class N(torch.nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(N, self).__init__()\n",
        "        self.n_features = params['N_units']\n",
        "        self.bn1 = torch.nn.BatchNorm1d(self.n_features)\n",
        "        self.layers = torch.nn.ModuleList()\n",
        "        n1_layers = params['N_layers'].copy()\n",
        "        n1_layers.insert(0, self.n_features)\n",
        "        for n in range(1, len(n1_layers)):\n",
        "            self.layers.append(torch.nn.Linear(n1_layers[n - 1], n1_layers[n]))\n",
        "            self.layers.append(torch.nn.SELU())\n",
        "            if params['bat_n_all']:\n",
        "                self.layers.append(torch.nn.BatchNorm1d(n1_layers[n]))\n",
        "        self.out = torch.nn.Linear(n1_layers[-1], self.n_features)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # input dims: (N, 30)\n",
        "        x = self.bn1(x)\n",
        "        for layer in self.layers:\n",
        "            x = layer(x)\n",
        "        x = self.out(x)\n",
        "        return x\n",
        "\n",
        "class AE(torch.nn.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super(AE, self).__init__()\n",
        "    self.encoder = Encoder(n_points=pix_f.shape[1])\n",
        "    self.decoder = Decoder(n_points=pix_f.shape[1], latent_dim=params['N_units'])\n",
        "\n",
        "  def forward(self, x):\n",
        "    # full shape autoencoder\n",
        "    lat = self.encoder(x)\n",
        "    mesh_rec = self.decoder(lat)\n",
        "    # partial shape autoencoder\n",
        "    return mesh_rec, lat\n",
        "  \n",
        "  def encode(self, x):\n",
        "    return self.encoder(x)\n",
        "\n",
        "  def decode(self, lat):\n",
        "    return self.decoder(lat)\n",
        "\n",
        "\n",
        "class Model(torch.nn.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super(Model, self).__init__()\n",
        "    self.ae_f = AE()\n",
        "    self.ae_p = AE()\n",
        "    self.n_p_f = N()\n",
        "    self.n_f_p = N()\n",
        "\n",
        "  def forward(self, x_f, x_p):\n",
        "    # mlp network\n",
        "    mesh_rec_f, lat_f = self.ae_f(x_f)\n",
        "    mesh_rec_p, lat_p = self.ae_p(x_p)\n",
        "    lat_rec_f = self.map_p_to_f(lat_p)\n",
        "    lat_rec_p = self.map_f_to_p(lat_f)\n",
        "    return mesh_rec_f, mesh_rec_p, lat_f, lat_p, lat_rec_f, lat_rec_p\n",
        "\n",
        "  def map_p_to_f(self, lat_p):\n",
        "    return self.n_p_f(lat_p)\n",
        "\n",
        "  def map_f_to_p(self, lat_f):\n",
        "    return self.n_f_p(lat_f)"
      ],
      "metadata": {
        "id": "YCZF-DiK9n-Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#%%\n",
        "# Datasets\n",
        "N_train = 800\n",
        "\n",
        "train_f = pix_f[:N_train, :,:]\n",
        "test_f = pix_f[N_train:, :,:]\n",
        "train_dataset_f = torch.utils.data.TensorDataset(torch.tensor(train_f).to(torch.device(\"cuda:0\")))\n",
        "train_dataloader_f = torch.utils.data.DataLoader(train_dataset_f, batch_size=int(params['B_size']), num_workers=0)\n",
        "\n",
        "test_dataset_f = torch.utils.data.TensorDataset(torch.tensor(test_f).to(torch.device(\"cuda:0\")))\n",
        "test_dataloader_f = torch.utils.data.DataLoader(test_dataset_f, batch_size=int(params['B_size']), num_workers=0)\n",
        "\n",
        "\n",
        "train_p = pix_p[:N_train, :,:]\n",
        "test_p = pix_p[N_train:, :,:]\n",
        "train_dataset_p = torch.utils.data.TensorDataset(torch.tensor(train_p).to(torch.device(\"cuda:0\")))\n",
        "train_dataloader_p = torch.utils.data.DataLoader(train_dataset_p, batch_size=int(params['B_size']), num_workers=0)\n",
        "\n",
        "test_dataset_p = torch.utils.data.TensorDataset(torch.tensor(test_p).to(torch.device(\"cuda:0\")))\n",
        "test_dataloader_p = torch.utils.data.DataLoader(test_dataset_p, batch_size=int(params['B_size']),  num_workers=0)\n",
        "\n",
        "# Model\n",
        "model = Model().cuda()"
      ],
      "metadata": {
        "id": "DySd0NBQowtU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # Optimizer\n",
        "all_params = set(model.parameters())\n",
        "\n",
        "opt = torch.optim.Adam(params=[{\"params\": list(all_params)}], lr=params['L_rate_AE'])\n",
        "# #%%\n",
        "\n",
        "train_loss_f = []\n",
        "eval_loss_f = []\n",
        "train_loss_p = []\n",
        "eval_loss_p = []\n",
        "train_loss_n_f_p = []\n",
        "train_loss_n_p_f = []\n",
        "eval_loss_n_f_p = []\n",
        "eval_loss_n_p_f = []\n",
        "epochs = []\n",
        "\n",
        "losses_train = []\n",
        "losses_test = []\n",
        "\n",
        "for epoch in range(0, params['Epochs']):\n",
        "  epochs.append(epoch)\n",
        "  avg_loss_ae_f = torch.tensor(0.).cuda()\n",
        "  avg_loss_ae_p = torch.tensor(0.).cuda()\n",
        "  avg_loss_n_f_p = torch.tensor(0.).cuda()\n",
        "  avg_loss_n_p_f = torch.tensor(0.).cuda()\n",
        "  avg_loss = torch.tensor(0.).cuda()\n",
        "  for meshes_f, meshes_p in tqdm(zip(train_dataloader_f, train_dataloader_p)):\n",
        "    mesh_rec_f, mesh_rec_p, lat_f, lat_p, lat_rec_f, lat_rec_p = model(meshes_f[0], meshes_p[0])\n",
        "    loss_ae_f = torch.sum((meshes_f[0] - mesh_rec_f)**2, dim=-1).mean()\n",
        "    loss_ae_p = torch.sum((meshes_p[0] - mesh_rec_p)**2, dim=-1).mean()\n",
        "    loss_n_p_f = params['W_N_p_f']*torch.sum((lat_f - lat_rec_f)**2, dim=-1).mean()\n",
        "    loss_n_f_p = params['W_N_f_p']*torch.sum((lat_p - lat_rec_p)**2, dim=-1).mean()\n",
        "    loss = loss_ae_f + loss_ae_p + loss_n_p_f + loss_n_f_p\n",
        "    loss.backward()\n",
        "    opt.step()\n",
        "    opt.zero_grad()\n",
        "    avg_loss_ae_f += loss_ae_f\n",
        "    avg_loss_ae_p += loss_ae_p\n",
        "    avg_loss_n_p_f += loss_n_p_f\n",
        "    avg_loss += loss\n",
        "\n",
        "  avg_loss_ae_f /= len(train_dataloader_f)\n",
        "  avg_loss_ae_p /= len(train_dataloader_p)\n",
        "  avg_loss_n_p_f /= len(train_dataloader_p)\n",
        "  avg_loss /= len(train_dataloader_f)\n",
        "\n",
        "  losses_train.append(avg_loss)\n",
        "  print('Epoch {} of {}, Train Loss Full and Partial: {:.4f}'.format(epoch+1, params['Epochs'], loss_ae_f, loss_ae_p))\n",
        "  print(f'Train Avg: avg_loss_ae_f = {avg_loss_ae_f.item()}; avg_loss_ae_p = {avg_loss_ae_p.item()}; avg_loss_n_p_f = {avg_loss_n_p_f.item()}; avg_loss = {avg_loss.item()}')\n",
        "  #if (epoch+1) % 500 == 0:\n",
        "\n",
        "  if (epoch) % save_every == 0:\n",
        "    torch.save(model.state_dict(), model_path + '/ae_' + str(epoch) + '.pt')\n",
        "\n",
        "  if (epoch+1) % test_every == 0:\n",
        "    avg_loss_ae_f = torch.tensor(0.).cuda()\n",
        "    avg_loss_ae_p = torch.tensor(0.).cuda()\n",
        "    avg_loss_n_f_p = torch.tensor(0.).cuda()\n",
        "    avg_loss_n_p_f = torch.tensor(0.).cuda()\n",
        "    avg_loss = torch.tensor(0.).cuda()\n",
        "    with torch.no_grad():\n",
        "      for meshes in tqdm(zip(train_dataloader_f, train_dataloader_p)):\n",
        "        mesh_rec_f, mesh_rec_p, lat_f, lat_p, lat_rec_f , lat_rec_p = model(meshes_f[0], meshes_p[0])\n",
        "        # loss_ae = chamfer_distance(meshes[0].cuda(), g_mesh)\n",
        "        loss_ae_f = torch.sum((meshes_f[0] - mesh_rec_f)**2, dim = -1).mean()\n",
        "        loss_ae_p = torch.sum((meshes_p[0] - mesh_rec_p)**2, dim = -1).mean()\n",
        "        loss_n_p_f = params['W_N_p_f']*torch.nn.functional.mse_loss(lat_f, lat_rec_f)\n",
        "        loss_n_f_p = params['W_N_f_p']*torch.nn.functional.mse_loss(lat_p, lat_rec_p)\n",
        "        loss = loss_ae_f + loss_ae_p + loss_n_p_f + loss_n_f_p\n",
        "        avg_loss_n_p_f += loss_n_p_f\n",
        "        avg_loss_ae_f += loss_ae_f\n",
        "        avg_loss_ae_p += loss_ae_p\n",
        "        avg_loss += loss\n",
        "\n",
        "      avg_loss_ae_f /= len(test_dataloader_f)\n",
        "      avg_loss_ae_p /= len(test_dataloader_p)\n",
        "      avg_loss /= len(test_dataloader_p)\n",
        "      avg_loss_n_p_f /= len(test_dataloader_p)\n",
        "\n",
        "    losses_test.append(avg_loss)\n",
        "    print(f'Eval Avg: avg_loss_ae_f = {avg_loss_ae_f.item()}; avg_loss_ae_p = {avg_loss_ae_p.item()}; avg_loss_n_p_f = {avg_loss_n_p_f.item()}; avg_loss = {avg_loss.item()}')\n",
        "\n",
        "losses_train = [x.detach().cpu() for x in losses_train]\n",
        "losses_test = [x.detach().cpu() for x in losses_test]"
      ],
      "metadata": {
        "id": "YQEVvzMadQx8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.save(model.state_dict(), '/content/ae_model.pt')"
      ],
      "metadata": {
        "id": "B792hdgco2e4"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}