2-3-2-SketchVAE-train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "In this file \n",
    "We train our proposed SketchVAE\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 SketchVAE.sketchvae import SketchVAE\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\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",
    "s_dir = \"\" # folder_address\n",
    "batch_size = 64\n",
    "n_epochs = 100\n",
    "data_path = [\"data/irish_train_chord_rhythm.npy\",\n",
    "             \"data/irish_validate_chord_rhythm.npy\",\n",
    "             \"data/irish_test_chord_rhythm.npy\"]\n",
    "save_path = \"model_backup\" # save_model_address \n",
    "lr = 1e-4\n",
    "decay = 0.9999\n",
    "hidden_dims = 1024\n",
    "zp_dims = 128\n",
    "zr_dims = 128\n",
    "vae_beta = 0.1\n",
    "input_dims = 130\n",
    "pitch_dims = 129\n",
    "rhythm_dims = 3\n",
    "seq_len = 4 * 6\n",
    "beat_num = 4\n",
    "tick_num = 6\n",
    "# set here to config your save_period (2 i.e. save the model every 2 epochs)\n",
    "save_period = 2\n",
    "##############################\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "def processed_data_tensor(data):\n",
    "    print(\"processed data:\")\n",
    "    gd = np.array([d[0] for d in data])\n",
    "    px = np.array([d[1] for d in data])\n",
    "    rx = np.array([d[2] for d in data])\n",
    "    len_x = np.array([d[3] for d in data])\n",
    "    nrx = []\n",
    "    for i,r in enumerate(rx):\n",
    "        temp = np.zeros((seq_len, rhythm_dims))\n",
    "        lins = np.arange(0, len(r))\n",
    "        temp[lins, r - 1] = 1\n",
    "        nrx.append(temp)\n",
    "    nrx = np.array(nrx)\n",
    "    gd = torch.from_numpy(gd).long()\n",
    "    px = torch.from_numpy(px).long()\n",
    "    rx = torch.from_numpy(rx).float()\n",
    "    len_x = torch.from_numpy(len_x).long()\n",
    "    nrx = torch.from_numpy(nrx).float()\n",
    "    print(\"processed finish!\")\n",
    "    return TensorDataset(px, rx, len_x, nrx, gd)\n",
    "train_set = np.load(os.path.join(s_dir,data_path[0]), allow_pickle = True)\n",
    "validate_set = np.load(os.path.join(s_dir,data_path[1]),allow_pickle = True)\n",
    "train_set = DataLoader(\n",
    "    dataset = processed_data_tensor(train_set),\n",
    "    batch_size = batch_size, \n",
    "    shuffle = True, \n",
    "    num_workers = 8, \n",
    "    pin_memory = True, \n",
    "    drop_last = True\n",
    ")\n",
    "validate_set = DataLoader(\n",
    "    dataset = processed_data_tensor(validate_set),\n",
    "    batch_size = batch_size, \n",
    "    shuffle = False, \n",
    "    num_workers = 8, \n",
    "    pin_memory = True, \n",
    "    drop_last = True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import model\n",
    "model = SketchVAE(input_dims, pitch_dims, rhythm_dims, hidden_dims, zp_dims, zr_dims, seq_len, beat_num, tick_num, 4000)\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": [
    "# process validete data from the dataloder\n",
    "validate_data = []\n",
    "for i,d in enumerate(validate_set):\n",
    "    validate_data.append(d)\n",
    "print(len(validate_data))"
   ]
  },
  {
   "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, target, p_dis, r_dis, beta):\n",
    "    CE = F.cross_entropy(recon.view(-1, recon.size(-1)), target, reduction = \"mean\")\n",
    "    normal1 = std_normal(p_dis.mean.size())\n",
    "    normal2=  std_normal(r_dis.mean.size())\n",
    "    KLD1 = kl_divergence(p_dis, normal1).mean()\n",
    "    KLD2 = kl_divergence(r_dis, normal2).mean()\n",
    "    max_indices = recon.view(-1, recon.size(-1)).max(-1)[-1]\n",
    "    correct = max_indices == target\n",
    "    acc = torch.sum(correct.float()) / target.size(0)\n",
    "    return acc, CE + beta * (KLD1 + KLD2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# start training\n",
    "logs = []\n",
    "device = torch.device(torch.cuda.current_device())\n",
    "iteration = 0\n",
    "step = 0\n",
    "for epoch in range(n_epochs):\n",
    "    print(\"epoch: %d\\n__________________________________________\" % (epoch), flush = True)\n",
    "    mean_loss = 0.0\n",
    "    mean_acc = 0.0\n",
    "    v_mean_loss = 0.0\n",
    "    v_mean_acc = 0.0\n",
    "    total = 0\n",
    "    for i, d in enumerate(train_set):\n",
    "        # validate display\n",
    "        model.train()\n",
    "        j = i % len(validate_data)\n",
    "        px, rx, len_x, nrx, gd = d\n",
    "        v_px, v_rx, v_len_x, v_nrx, v_gd = validate_data[j]\n",
    "        \n",
    "        px = px.to(device = device,non_blocking = True)\n",
    "        len_x = len_x.to(device = device,non_blocking = True)\n",
    "        nrx = nrx.to(device = device,non_blocking = True)\n",
    "        gd = gd.to(device = device,non_blocking = True)\n",
    "\n",
    "        v_px = v_px.to(device = device,non_blocking = True)\n",
    "        v_len_x = v_len_x.to(device = device,non_blocking = True)\n",
    "        v_nrx = v_nrx.to(device = device,non_blocking = True)\n",
    "        v_gd = v_gd.to(device = device,non_blocking = True)\n",
    "            \n",
    "        optimizer.zero_grad()\n",
    "        recon, p_dis, r_dis, iteration = model(px, nrx, len_x, gd)\n",
    "        \n",
    "        acc, loss = loss_function(recon, gd.view(-1), p_dis, r_dis, vae_beta)\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n",
    "        optimizer.step()\n",
    "        mean_loss += loss.item()\n",
    "        mean_acc += acc.item()\n",
    "        \n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            v_recon, v_p_dis, v_r_dis, _ = model(v_px, v_nrx, v_len_x, v_gd)\n",
    "            v_acc, v_loss = loss_function(v_recon, v_gd.view(-1), v_p_dis, v_r_dis, vae_beta)\n",
    "            v_mean_loss += v_loss.item()\n",
    "            v_mean_acc += v_acc.item()\n",
    "        step += 1\n",
    "        total += 1\n",
    "        if decay > 0:\n",
    "            scheduler.step()\n",
    "        print(\"batch %d loss: %.5f acc: %.5f | val loss %.5f acc: %.5f iteration: %d\"  \n",
    "              % (i,loss.item(), acc.item(), v_loss.item(),v_acc.item(),iteration),flush = True)\n",
    "    mean_loss /= total\n",
    "    mean_acc /= total\n",
    "    v_mean_loss /= total\n",
    "    v_mean_acc /= total\n",
    "    print(\"epoch %d loss: %.5f acc: %.5f | val loss %.5f acc: %.5f iteration: %d\"  \n",
    "              % (epoch, mean_loss, mean_acc, v_mean_loss, v_mean_acc, iteration),flush = True)\n",
    "    logs.append([mean_loss,mean_acc,v_mean_loss,v_mean_acc,iteration])\n",
    "    if (epoch + 1) % save_period == 0:\n",
    "        filename = \"sketchvae-\" + 'loss_' + str(v_mean_loss) + \"_acc_\" + str(v_mean_acc) + \"_epoch_\" +  str(epoch+1) + \"_it_\" + str(iteration) + \".pt\"\n",
    "        torch.save(model.cpu().state_dict(),os.path.join(s_dir, save_path, filename))\n",
    "        model.cuda()\n",
    "    np.save(os.path.join(s_dir,\"sketchvae-log.npy\"), logs)\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
}