3-2-SketchVAE_InpaintRNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "In this file we train SketchVAE + InpaintRNN\n",
    "InpaintRNN is from \"Learning to Traverse Latent Spaces for Musical Score Inpainting\", published in ISMIR 2019\n",
    "The core model code is from their releasing code.\n",
    "'''\n",
    "import torch\n",
    "import os\n",
    "import numpy as np\n",
    "from torch import optim\n",
    "from torch.nn import functional as F\n",
    "from SketchVAE.sketchvae import SketchVAE\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "from torch.distributions import Normal\n",
    "from InpaintRNN.inpaintrnn_dis import InpaintingNet\n",
    "from utils.helpers import *\n",
    "import time\n",
    "###############################\n",
    "# initial parameters\n",
    "s_dir = \"\"\n",
    "input_dims = 256\n",
    "pf_hidden_dims = 512\n",
    "g_h_dims = 1024\n",
    "pf_num = 2\n",
    "inpaint_len = 4\n",
    "seq_len = 16\n",
    "batch_size = 32\n",
    "n_epochs = 100\n",
    "\n",
    "'''\n",
    "---------!!!NOTICE!!!---------\n",
    "Here we cut the irish_train/validate/test.npy into 16-measure as the group\n",
    "And we save them as the array \"irish-dis-measure-vae-xxx-whole.npy\"\n",
    "The size is (data length, measure length, four attributes)\n",
    "There are four attributes in each measure:\n",
    "    1. \"gd\": the measure melody token\n",
    "    2. \"px\": the measure pitch token\n",
    "    3, \"rx\": the measure rhythm token\n",
    "    4. \"len_x\": the number of the pitch in this measure\n",
    "We did not provide this process code, but we provide the processed \"npy\" files with the data link in Github.\n",
    "'''\n",
    "data_path = [\n",
    "    \"data/irish-dis-measure-vae-train-whole.npy\",\n",
    "    \"data/irish-dis-measure-vae-validate-whole.npy\",\n",
    "    \"data/irish-dis-measure-vae-test-whole.npy\"\n",
    "]\n",
    "lr = 1e-4\n",
    "decay = 0.9999\n",
    "##############################\n",
    "##############  for vae init ##############\n",
    "vae_hidden_dims = 1024\n",
    "vae_zp_dims = 128\n",
    "vae_zr_dims = 128\n",
    "vae_beta = 0.1\n",
    "vae_input_dims = 130\n",
    "vae_pitch_dims = 129\n",
    "vae_rhythm_dims = 3\n",
    "vae_seq_len = 6 * 4\n",
    "vae_beat_num = 4\n",
    "vae_tick_num = 6\n",
    "############################\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input data dis-measure-vae\n",
    "def processed_data_tensor(data):\n",
    "    print(\"processed data:\")\n",
    "    gd = [] \n",
    "    px = []\n",
    "    rx = []\n",
    "    len_x = []\n",
    "    nrx = []\n",
    "    total = 0\n",
    "    for i, d in enumerate(data):\n",
    "        gd.append([list(dd[0]) for dd in d])\n",
    "        px.append([list(dd[1]) for dd in d])\n",
    "        rx.append([list(dd[2]) for dd in d])\n",
    "        len_x.append([dd[3] for dd in d])\n",
    "        if len(gd[-1][-1]) != vae_seq_len:\n",
    "            gd[-1][-1].extend([128] * (vae_seq_len - len(gd[-1][-1])))\n",
    "            px[-1][-1].extend([128] * (vae_seq_len - len(px[-1][-1])))\n",
    "            rx[-1][-1].extend([2] * (vae_seq_len - len(rx[-1][-1])))\n",
    "    for i,d in enumerate(len_x):\n",
    "        for j,dd in enumerate(d):\n",
    "            if len_x[i][j] == 0:\n",
    "                gd[i][j][0] = 60\n",
    "                px[i][j][0] = 60\n",
    "                rx[i][j][0] = 1\n",
    "                len_x[i][j] = 1\n",
    "                total += 1\n",
    "    gd = np.array(gd)\n",
    "    px = np.array(px)\n",
    "    rx = np.array(rx)\n",
    "    len_x = np.array(len_x)\n",
    "    for d in rx:\n",
    "        nnrx = []\n",
    "        for dd in d:\n",
    "            temp = np.zeros((vae_seq_len, vae_rhythm_dims))\n",
    "            lins = np.arange(0, len(dd))\n",
    "            temp[lins, dd - 1] = 1\n",
    "            nnrx.append(temp)\n",
    "        nrx.append(nnrx)\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! zeros:\", total)\n",
    "    print(gd.size(),px.size(),rx.size(),len_x.size(),nrx.size())\n",
    "    return TensorDataset(px, rx, len_x, nrx, gd)\n",
    "\n",
    "train_set = np.load(data_path[0],allow_pickle = True)\n",
    "train_loader = 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",
    "\n",
    "validate_set = np.load(data_path[1],allow_pickle = True)\n",
    "validate_loader = 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",
    "validate_data = []\n",
    "for i,d in enumerate(validate_loader):\n",
    "    validate_data.append(d)\n",
    "print(len(validate_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load VAE model\n",
    "vae_model = SketchVAE(\n",
    "    vae_input_dims, vae_pitch_dims, vae_rhythm_dims, vae_hidden_dims, \n",
    "    vae_zp_dims, vae_zr_dims, vae_seq_len, vae_beat_num, vae_tick_num, 4000)\n",
    "dic = torch.load(\"model_backup/sketchvae-param.pt\")\n",
    "\n",
    "for name in list(dic.keys()):\n",
    "    dic[name.replace('module.', '')] = dic.pop(name)\n",
    "vae_model.load_state_dict(dic)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print('Using: ', torch.cuda.get_device_name(torch.cuda.current_device()))\n",
    "    vae_model.cuda()\n",
    "else:\n",
    "    print('Using: CPU')\n",
    "vae_model.eval()\n",
    "print(vae_model.training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import model\n",
    "save_path = \"model_backup/\"\n",
    "save_period = 5\n",
    "\n",
    "# think about traning with mse\n",
    "model = InpaintingNet(input_dims, pf_hidden_dims, g_h_dims, pf_num, inpaint_len, vae_model, False, 2000, True)\n",
    "optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),lr=lr)\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')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_raw_x(raw_x, n_past, n_inpaint, n_future):\n",
    "    raw_px, raw_rx, raw_len_x, raw_nrx, raw_gd = raw_x\n",
    "    past_px = raw_px[:,:n_past,:]\n",
    "    inpaint_px = raw_px[:,n_past:n_past + n_inpaint,:]\n",
    "    future_px = raw_px[:,n_future:,:]\n",
    "    past_rx = raw_rx[:,:n_past,:]\n",
    "    inpaint_rx = raw_rx[:,n_past:n_past + n_inpaint,:]\n",
    "    future_rx = raw_rx[:,n_future:,:]\n",
    "    past_len_x = raw_len_x[:,:n_past]\n",
    "    inpaint_len_x = raw_len_x[:,n_past:n_past + n_inpaint]\n",
    "    future_len_x = raw_len_x[:,n_future:]\n",
    "    past_nrx = raw_nrx[:,:n_past,:]\n",
    "    inpaint_nrx = raw_nrx[:,n_past:n_past + n_inpaint,:]\n",
    "    future_nrx = raw_nrx[:,n_future:,:]\n",
    "    past_gd = raw_gd[:,:n_past,:]\n",
    "    inpaint_gd = raw_gd[:,n_past:n_past + n_inpaint,:]\n",
    "    future_gd = raw_gd[:,n_future:,:]\n",
    "    re = [\n",
    "        past_px, past_rx, past_len_x, past_nrx, past_gd,\n",
    "        inpaint_px, inpaint_rx, inpaint_len_x, inpaint_nrx, inpaint_gd,\n",
    "        future_px, future_rx, future_len_x, future_nrx, future_gd,\n",
    "    ]\n",
    "    return re\n",
    "def get_acc(recon, gd):\n",
    "    recon = recon.cpu().detach().numpy()\n",
    "    gd = gd.cpu().detach().numpy()\n",
    "    return np.sum(recon == gd) / recon.size\n",
    "\n",
    "device = torch.device(torch.cuda.current_device())\n",
    "save_period = 2\n",
    "losses = []\n",
    "step = 0\n",
    "n_past = 6\n",
    "n_future = 10\n",
    "n_inpaint = 4\n",
    "iteration = 0\n",
    "# save_period = 200\n",
    "print(vae_model.training)\n",
    "for epoch in range(n_epochs):\n",
    "    model.train()\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,tr_data in enumerate(train_loader):\n",
    "        model.train()\n",
    "#         print(\"begin_time\", time.process_time())\n",
    "        j = i % len(validate_data)\n",
    "        raw_x = process_raw_x(tr_data, n_past, n_inpaint, n_future)\n",
    "        for k in range(len(raw_x)):\n",
    "            raw_x[k] = raw_x[k].to(device = device,non_blocking = True)\n",
    "        past_px, past_rx, past_len_x, past_nrx, past_gd, \\\n",
    "        inpaint_px, inpaint_rx, inpaint_len_x, inpaint_nrx, inpaint_gd,\\\n",
    "        future_px, future_rx, future_len_x, future_nrx, future_gd = raw_x\n",
    "        inpaint_gd_whole = inpaint_gd.contiguous().view(-1)\n",
    "        past_x = [past_px, past_rx, past_len_x, past_nrx, past_gd]\n",
    "        inpaint_x = [inpaint_px, inpaint_rx, inpaint_len_x, inpaint_nrx, inpaint_gd]\n",
    "        future_x = [future_px, future_rx, future_len_x, future_nrx, future_gd]\n",
    "        \n",
    "        # validate\n",
    "        v_raw_x = process_raw_x(validate_data[j], n_past, n_inpaint, n_future)\n",
    "        for k in range(len(v_raw_x)):\n",
    "            v_raw_x[k] = v_raw_x[k].to(device = device,non_blocking = True)\n",
    "        v_past_px, v_past_rx, v_past_len_x, v_past_nrx, v_past_gd, \\\n",
    "        v_inpaint_px, v_inpaint_rx, v_inpaint_len_x, v_inpaint_nrx, v_inpaint_gd,\\\n",
    "        v_future_px, v_future_rx, v_future_len_x, v_future_nrx, v_future_gd = v_raw_x\n",
    "        v_inpaint_gd_whole = v_inpaint_gd.contiguous().view(-1)\n",
    "        v_past_x = [v_past_px, v_past_rx, v_past_len_x, v_past_nrx, v_past_gd]\n",
    "        v_inpaint_x = [v_inpaint_px, v_inpaint_rx, v_inpaint_len_x, v_inpaint_nrx, v_inpaint_gd]\n",
    "        v_future_x = [v_future_px, v_future_rx, v_future_len_x, v_future_nrx, v_future_gd]\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        recon_x, iteration, use_teacher = model(past_x, future_x, inpaint_x)\n",
    "        # for NLL Loss\n",
    "        loss = F.cross_entropy(recon_x.view(-1, recon_x.size(-1)), inpaint_gd_whole, reduction = \"mean\") \n",
    "        acc = get_acc(recon_x.view(-1, recon_x.size(-1)).argmax(-1), inpaint_gd_whole)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "   \n",
    "        total += 1\n",
    "        mean_loss += loss.item()\n",
    "        mean_acc += acc\n",
    "\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            v_recon_x, _,_ = model(v_past_x, v_future_x, v_inpaint_x)\n",
    "            v_loss = F.cross_entropy(v_recon_x.view(-1, v_recon_x.size(-1)), v_inpaint_gd_whole, reduction = \"mean\") \n",
    "            v_acc = get_acc(v_recon_x.view(-1, v_recon_x.size(-1)).argmax(-1), v_inpaint_gd_whole)\n",
    "            v_mean_loss += v_loss.item()\n",
    "            v_mean_acc += v_acc\n",
    "        print(\"batch %d loss: %.5f acc: %.5f | v_loss: %.5f v_acc: %.5f |  iteration: %d teacher: %d\" \\\n",
    "              % (i,loss.item(), acc, v_loss.item(), v_acc, iteration, use_teacher),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 | v_loss: %.5f v_acc: %.5f \"  % (epoch,mean_loss, mean_acc, v_mean_loss, v_mean_acc),flush = True)\n",
    "    losses.append([mean_loss, v_mean_loss])\n",
    "    if (epoch + 1) % save_period == 0:\n",
    "        filename = \"inpaintNet-sketchvae-\" + 'loss_' + str(mean_loss) + \"_\" + str(epoch+1) + '_' + str(iteration) + \".pt\"\n",
    "        torch.save(model.cpu().state_dict(),save_path + filename)\n",
    "        model.cuda()\n",
    "    np.save(s_dir + \"inpaintNet_sketchvae_log.npy\",losses)\n"
   ]
  }
 ],
 "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
}