2-1-EC2VAE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "You can skip this if you only want to test the SketchNet\n",
    "In this file we train the EC2VAE \"Deep Music Analogy Via Latent Representation Disentanglement\", published in ISMIR 2019\n",
    "The core model code and training code are from their releasing codes.\n",
    "'''\n",
    "import torch\n",
    "import os\n",
    "import numpy as np\n",
    "from torch import optim\n",
    "from torch.distributions import kl_divergence, Normal\n",
    "from torch.nn import functional as F\n",
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "from ec2vae.model_dis import RECVAE\n",
    "from loader.dataloader import DataLoader\n",
    "\n",
    "class MinExponentialLR(ExponentialLR):\n",
    "    def __init__(self, optimizer, gamma, minimum, last_epoch=-1):\n",
    "        self.min = minimum\n",
    "        super(MinExponentialLR, self).__init__(optimizer, gamma, last_epoch=-1)\n",
    "\n",
    "    def get_lr(self):\n",
    "        return [\n",
    "            max(base_lr * self.gamma**self.last_epoch, self.min)\n",
    "            for base_lr in self.base_lrs\n",
    "        ]\n",
    "###############################\n",
    "# initial parameters\n",
    "batch_size = 128\n",
    "n_epochs = 100\n",
    "save_path = \"model_backup\" # model_save_path\n",
    "save_period = 2 # save every 2 epoches\n",
    "data_path = [\"data/irish_train.npy\",\n",
    "            \"data/irish_validate.npy\",\n",
    "            \"data/irish_test.npy\"]\n",
    "lr = 1e-4\n",
    "decay = 0.9999\n",
    "if_parallel = False\n",
    "hidden_dims = 2048\n",
    "z1_dims = 128\n",
    "z2_dims = 128\n",
    "vae_beta = 0.1\n",
    "input_dims = 130\n",
    "rhythm_dims = 3\n",
    "seq_len = 6 * 4\n",
    "##############################\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input data\n",
    "train_x = np.load(data_path[0],allow_pickle=True)\n",
    "validate_x = np.load(data_path[1],allow_pickle=True)\n",
    "test_x = np.load(data_path[2],allow_pickle=True)\n",
    "dl = DataLoader(train = train_x, validate = validate_x, test = test_x)\n",
    "# each measure is 24-token, we split it here\n",
    "dl.process_split(split_size = seq_len)\n",
    "print(len(dl.train_set),len(dl.validate_set),len(dl.test_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import model\n",
    "model = RECVAE(input_dims, hidden_dims, rhythm_dims, z1_dims,z2_dims,seq_len, 3000)\n",
    "optimizer = optim.Adam(model.parameters(), lr = lr)\n",
    "if decay > 0:\n",
    "    scheduler = MinExponentialLR(optimizer, gamma = decay, minimum = 1e-5)\n",
    "if torch.cuda.is_available():\n",
    "    print('Using: ', torch.cuda.get_device_name(torch.cuda.current_device()))\n",
    "    model.cuda()\n",
    "else:\n",
    "    print('Using: CPU')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loss function\n",
    "def std_normal(shape):\n",
    "    N = Normal(torch.zeros(shape), torch.ones(shape))\n",
    "    if torch.cuda.is_available():\n",
    "        N.loc = N.loc.cuda()\n",
    "        N.scale = N.scale.cuda()\n",
    "    return N\n",
    "\n",
    "def loss_function(recon, recon_rhythm, target, target_rhythm, dis1, dis2, beta):\n",
    "    CE = F.nll_loss(recon.view(-1, recon.size(-1)), target, reduction = \"mean\")\n",
    "    rhy_CE = F.nll_loss(recon_rhythm.view(-1, recon_rhythm.size(-1)), target_rhythm, reduction = \"mean\")\n",
    "    normal1 = std_normal(dis1.mean.size())\n",
    "    normal2=  std_normal(dis2.mean.size())\n",
    "    KLD1 = kl_divergence(dis1, normal1).mean()\n",
    "    KLD2 = kl_divergence(dis2, normal2).mean()\n",
    "    return CE, rhy_CE, CE + rhy_CE + beta * (KLD1 + KLD2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start training\n",
    "model.train()\n",
    "step = 0\n",
    "for epoch in range(n_epochs):\n",
    "    print(\"epoch: %d\\n__________________________________________\" % (epoch), flush = True)\n",
    "    train_batches, validate_batches = dl.start_new_epoch(batch_size = batch_size)\n",
    "    for i in range(len(train_batches)):\n",
    "        # validate display\n",
    "        j = i % len(validate_batches)\n",
    "        raw_x = dl.convert_onehot(train_batches[i])\n",
    "        raw_vx = dl.convert_onehot(validate_batches[j])\n",
    "        \n",
    "        x = torch.from_numpy(raw_x).float()\n",
    "        target_rhythm = np.expand_dims(raw_x[:,:,:-2].sum(-1),-1)\n",
    "        target_rhythm = np.concatenate((target_rhythm, raw_x[:,:,-2:]), -1)\n",
    "        target_rhythm = torch.from_numpy(target_rhythm).float()\n",
    "        target_rhythm = target_rhythm.view(-1, target_rhythm.size(-1)).max(-1)[1]\n",
    "        target = x.view(-1, x.size(-1)).max(-1)[1]\n",
    "        \n",
    "        \n",
    "        vx = torch.from_numpy(raw_vx).float()\n",
    "        target_rhythm_vx = np.expand_dims(raw_vx[:,:,:-2].sum(-1),-1)\n",
    "        target_rhythm_vx = np.concatenate((target_rhythm_vx, raw_vx[:,:,-2:]), -1)\n",
    "        target_rhythm_vx = torch.from_numpy(target_rhythm_vx).float()\n",
    "        target_rhythm_vx = target_rhythm_vx.view(-1, target_rhythm_vx.size(-1)).max(-1)[1]\n",
    "        target_vx = vx.view(-1, vx.size(-1)).max(-1)[1]\n",
    "        \n",
    "        \n",
    "        if torch.cuda.is_available():\n",
    "            x = x.cuda()\n",
    "            target = target.cuda()\n",
    "            target_rhythm = target_rhythm.cuda()\n",
    "            vx = vx.cuda()\n",
    "            target_vx = target_vx.cuda()\n",
    "            target_rhythm_vx = target_rhythm_vx.cuda()\n",
    "            \n",
    "        optimizer.zero_grad()\n",
    "        recon_x, recon_rhythm, dis1, dis2 = model(x)\n",
    "        dis1 = Normal(dis1.mean, dis1.stddev)\n",
    "        dis2 = Normal(dis2.mean, dis2.stddev)\n",
    "        recon_loss, rhythm_loss,total_loss = loss_function(recon_x, recon_rhythm,target,target_rhythm, dis1, dis2, vae_beta)\n",
    "        total_loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n",
    "        optimizer.step()\n",
    "        \n",
    "        v_total_loss = 0.0\n",
    "        v_recon_loss = 0.0\n",
    "        v_rhythm_loss = 0.0\n",
    "        with torch.no_grad():\n",
    "            vrecon_x, vrecon_rhythm, vdis1, vdis2 = model(vx)\n",
    "            vdis1 = Normal(vdis1.mean, vdis1.stddev)\n",
    "            vdis2 = Normal(vdis2.mean, vdis2.stddev)\n",
    "            v_recon_loss, v_rhythm_loss, v_total_loss = loss_function(vrecon_x, vrecon_rhythm,target_vx,target_rhythm_vx, vdis1, vdis2,vae_beta)\n",
    "        step += 1\n",
    "        if decay > 0:\n",
    "            scheduler.step()\n",
    "        print(\"batch %d (total_loss, recon_loss, rhythm_loss) = (%.5f, %.5f, %.5f) | validate = (%.5f, %.5f, %.5f)\"  \n",
    "              % (i,total_loss.item(), recon_loss.item(), rhythm_loss.item(),v_total_loss.item(),v_recon_loss.item(),v_rhythm_loss.item()),flush = True)\n",
    "    if (epoch + 1) % save_period == 0:\n",
    "        filename = \"ec2vae-\" + 'loss_' + str(v_total_loss.item()) + \"_epoch_\" + str(epoch+1) + \".pt\"\n",
    "        torch.save(model.cpu().state_dict(),os.path.join(save_path,filename))\n",
    "        model.cuda()\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}