3-1-MusicInpaintNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "In this file we train the Music InpaintNet = MeasureVAE + InpaintRNN\n",
    "\"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 InpaintRNN.inpaintrnn import InpaintingNet\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "from torch.distributions import Normal\n",
    "from MeasureVAE.measure_vae import MeasureVAE\n",
    "from utils.helpers import *\n",
    "import time\n",
    "###############################\n",
    "# initial parameters\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",
    "---------!!!NOTICE!!!---------\n",
    "Here we cut the irish_train/validate/test.npy into 16-measure as the group\n",
    "And we save them as dict file \"irish-measure-vae-xxx-whole.npy\"\n",
    "There are two keys in the dict file:\n",
    "    1. \"data\": the latent variable output from the input 16-measure, it is deprecated.\n",
    "    2. \"gd\": the 16-mesaure melody token, we only need this\n",
    "The processing method is easily obtained by using np.split\n",
    "We did not provide this process code, but we provide the processed \"npy\" files with the data link in Github.\n",
    "'''\n",
    "whole_data_path = [\n",
    "    \"data/irish-measure-vae-train-whole.npy\",\n",
    "    \"data/irish-measure-vae-validate-whole.npy\",\n",
    "    \"data/irish-measure-vae-test-whole.npy\"\n",
    "]\n",
    "lr = 1e-4\n",
    "save_period = 5\n",
    "decay = 0.9999\n",
    "##############################\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "train_set = np.load(whole_data_path[0], allow_pickle = True)\n",
    "validate_set = np.load(whole_data_path[1],allow_pickle = True)\n",
    "\n",
    "train_x = torch.from_numpy(train_set.item()[\"data\"]).float()\n",
    "train_gd = torch.from_numpy(train_set.item()[\"gd\"]).long()\n",
    "validate_x = torch.from_numpy(validate_set.item()[\"data\"]).float()\n",
    "validate_gd = torch.from_numpy(validate_set.item()[\"gd\"]).long()\n",
    "\n",
    "print(train_x.size())\n",
    "print(train_gd.size())\n",
    "print(validate_x.size())\n",
    "print(validate_gd.size())\n",
    "\n",
    "# You can see here we only use the \"gd\" without \"data\"\n",
    "train_set = TensorDataset(train_gd,train_gd)\n",
    "validate_set = TensorDataset(validate_gd,validate_gd)\n",
    "\n",
    "train_loader = DataLoader(dataset = train_set, batch_size = batch_size, \n",
    "                          shuffle = True, num_workers = 2, pin_memory = True, drop_last = True)\n",
    "validate_loader = DataLoader(dataset = validate_set, batch_size = batch_size, \n",
    "                          shuffle = False, num_workers = 2, pin_memory = True, drop_last = True)\n",
    "\n",
    "validate_x = []\n",
    "validate_gd = []\n",
    "for x,y in validate_loader:\n",
    "    validate_x.append(x)\n",
    "    validate_gd.append(y)\n",
    "\n",
    "validate_loader = []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load VAE model\n",
    "vae_num_notes = 130\n",
    "vae_note_embedding_dim=10\n",
    "vae_metadata_embedding_dim=2\n",
    "vae_num_encoder_layers=2\n",
    "vae_encoder_hidden_size=512\n",
    "vae_encoder_dropout_prob=0.5\n",
    "vae_has_metadata=False\n",
    "vae_latent_space_dim=256\n",
    "vae_num_decoder_layers=2\n",
    "vae_decoder_hidden_size=512\n",
    "vae_decoder_dropout_prob=0.5\n",
    "vae_batch_size=256\n",
    "vae_num_epochs=30\n",
    "vae_train=False\n",
    "vae_plot=False\n",
    "vae_log=True\n",
    "vae_lr=1e-4\n",
    "vae_seq_len = 6 * 4\n",
    "\n",
    "vae_model = MeasureVAE(\n",
    "        num_notes = vae_num_notes,\n",
    "        note_embedding_dim=vae_note_embedding_dim,\n",
    "        metadata_embedding_dim=vae_metadata_embedding_dim,\n",
    "        num_encoder_layers=vae_num_encoder_layers,\n",
    "        encoder_hidden_size=vae_encoder_hidden_size,\n",
    "        encoder_dropout_prob=vae_encoder_dropout_prob,\n",
    "        latent_space_dim=vae_latent_space_dim,\n",
    "        num_decoder_layers=vae_num_decoder_layers,\n",
    "        decoder_hidden_size=vae_decoder_hidden_size,\n",
    "        decoder_dropout_prob=vae_decoder_dropout_prob,\n",
    "        has_metadata=vae_has_metadata\n",
    ")\n",
    "dic = torch.load(\"model_backup/measure-vae-param.pt\")\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import model\n",
    "save_path = \"model_backup/\"\n",
    "\n",
    "\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": [
    "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",
    "for epoch in range(n_epochs):\n",
    "    model.train()\n",
    "    print(\"epoch: %d\\n__________________________________________\" % (epoch), flush = True)\n",
    "    mean_loss = 0.0\n",
    "    v_mean_loss = 0.0\n",
    "    total = 0\n",
    "    for i,tr_data in enumerate(train_loader):\n",
    "        model.train()\n",
    "        j = i % len(validate_x)\n",
    "        raw_x,raw_gd = tr_data\n",
    "        past_x = raw_x[:,:n_past,:]\n",
    "        inpaint_x = raw_x[:,n_past:n_past + n_inpaint,:]\n",
    "        future_x = raw_x[:,n_future:,:]\n",
    "        \n",
    "        inpaint_gd =  raw_gd[:,n_past:n_past + n_inpaint,:]\n",
    "        inpaint_gd = inpaint_gd.contiguous().view(-1)\n",
    "        \n",
    "        past_x = past_x.to(device = device,non_blocking = True)\n",
    "        inpaint_x = inpaint_x.to(device = device,non_blocking = True)\n",
    "        future_x = future_x.to(device = device,non_blocking = True)\n",
    "        inpaint_gd = inpaint_gd.to(device = device,non_blocking = True)\n",
    "        \n",
    "        # validate\n",
    "        v_raw_x = validate_x[j]\n",
    "        v_raw_gd =  validate_gd[j]\n",
    "        v_past_x = v_raw_x[:,:n_past,:]\n",
    "        v_inpaint_x = v_raw_x[:,n_past:n_past + n_inpaint,:]\n",
    "        v_future_x = v_raw_x[:,n_future:,:] \n",
    "        v_inpaint_gd = v_raw_gd[:,n_past:n_past + n_inpaint,:]\n",
    "        v_inpaint_gd = v_inpaint_gd.contiguous().view(-1)\n",
    "        \n",
    "        v_past_x = v_past_x.to(device = device, non_blocking = True)\n",
    "        v_inpaint_x = v_inpaint_x.to(device = device, non_blocking = True)\n",
    "        v_future_x = v_future_x.to(device = device, non_blocking = True)\n",
    "        v_inpaint_gd = v_inpaint_gd.to(device = device, non_blocking = True)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        recon_x, iteration = model(past_x, future_x, inpaint_x)\n",
    "        loss = F.cross_entropy(recon_x.view(-1, recon_x.size(-1)), inpaint_gd, reduction = \"mean\") \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total += 1\n",
    "        mean_loss += loss.item()\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, reduction = \"mean\") \n",
    "            v_mean_loss += v_loss.item()\n",
    "        print(\"batch %d loss: %.5f | v_loss: %.5f |  iteration: %d\"  % (i,loss.item(), v_loss.item(), iteration),flush = True)\n",
    "    mean_loss /= total\n",
    "    v_mean_loss /= total\n",
    "    print(\"epoch %d loss: %.5f | v_loss: %.5f \"  % (epoch,mean_loss, v_mean_loss),flush = True)\n",
    "    losses.append([mean_loss, v_mean_loss])\n",
    "    if (epoch + 1) % save_period == 0:\n",
    "        filename = \"inpaintNet-\" + '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(\"inpaintNet_measure_log.npy\",losses)\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
}