4-5-SketchNet-eval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "In this file we eval the sketchnet\n",
    "Notice the eval of stage-1 and stage-2 is the same\n",
    "So we include as one file\n",
    "We provide a sceneario for simply sketching\n",
    "You can choose one another song B and take any part of it as a sketch information to the original song A, to see what the generation will be\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 SketchNet.sketchnet import SketchNet\n",
    "from utils.helpers import *\n",
    "import time\n",
    "###############################\n",
    "# initial parameters\n",
    "s_dir = \"\"\n",
    "zp_dims = 128\n",
    "zr_dims = 128\n",
    "pf_dims = 512\n",
    "gen_dims = 1024\n",
    "combine_dims = 512\n",
    "combine_head = 4\n",
    "combine_num = 4\n",
    "pf_num = 2\n",
    "inpaint_len = 4\n",
    "seq_len = 16\n",
    "total_len = 16\n",
    "batch_size = 1\n",
    "n_epochs = 15\n",
    "data_path = [\n",
    "    s_dir + \"data/irish-dis-measure-vae-train-whole.npy\",\n",
    "    s_dir + \"data/irish-dis-measure-vae-validate-whole.npy\",\n",
    "    s_dir + \"data/irish-dis-measure-vae-validate-repetition.npy\",\n",
    "    s_dir + \"data/irish-dis-measure-vae-validate-non-repetition.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",
    "validate_set = np.load(data_path[3],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 = s_dir +\"model_backup/\"\n",
    "save_period = 5\n",
    "\n",
    "# think about traning with mse\n",
    "model = SketchNet(\n",
    "    zp_dims, zr_dims, \n",
    "    pf_dims, gen_dims, combine_dims,\n",
    "    pf_num, combine_num, combine_head,\n",
    "    inpaint_len, total_len, \n",
    "    vae_model, True\n",
    ")\n",
    "dic = torch.load(save_path + \"sketchnet-param.pt\")\n",
    "for name in list(dic.keys()):\n",
    "    dic[name.replace('module.', '')] = dic.pop(name)\n",
    "model.load_state_dict(dic)\n",
    "model.set_stage(\"sketch\")\n",
    "\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",
    "print(model)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for normal eval process\n",
    "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",
    "model.set_stage(\"sketch\")\n",
    "device = torch.device(torch.cuda.current_device())\n",
    "save_period = 5\n",
    "losses = []\n",
    "step = 0\n",
    "n_past = 6\n",
    "n_future = 10\n",
    "n_inpaint = 4\n",
    "iteration = 0\n",
    "output = []\n",
    "v_mean_loss = 0.0\n",
    "v_mean_acc = 0.0\n",
    "total = 0\n",
    "print(vae_model.training)\n",
    "for i in range(len(validate_data)):\n",
    "    v_raw_x = process_raw_x(validate_data[i], 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",
    "    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",
    "        v_result = v_recon_x.argmax(-1)\n",
    "    total += 1\n",
    "    output.append(\n",
    "        {\n",
    "         \"past\": v_past_gd.cpu().detach().numpy(),\n",
    "         \"future\":v_future_gd.cpu().detach().numpy(),\n",
    "         \"inpaint\":v_result.cpu().detach().numpy(),\n",
    "         \"gd\":v_inpaint_gd.cpu().detach().numpy(),\n",
    "         \"acc\":v_acc,\n",
    "         \"nll\":v_loss.item()\n",
    "        }\n",
    "    )\n",
    "    print(i)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for sketch eval proces\n",
    "# sketch control generation\n",
    "import random\n",
    "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",
    "def gen_sketch(index, px, len_x, nrx, total):\n",
    "    # pitch_fisrt, rhythm_last\n",
    "    sketch_cond = []\n",
    "    for i in index:\n",
    "        if i < total:\n",
    "            sketch_cond.append([px[:, i], len_x[:, i]])\n",
    "        else:\n",
    "            sketch_cond.append(nrx[:, i - total])\n",
    "    return index, sketch_cond\n",
    "            \n",
    "            \n",
    "    \n",
    "\n",
    "model.set_stage(\"sketch\")\n",
    "device = torch.device(torch.cuda.current_device())\n",
    "save_period = 5\n",
    "losses = []\n",
    "step = 0\n",
    "n_past = 6\n",
    "n_future = 10\n",
    "n_inpaint = 4\n",
    "iteration = 0\n",
    "output = []\n",
    "v_mean_loss = 0.0\n",
    "v_mean_acc = 0.0\n",
    "total = 0\n",
    "print(vae_model.training)\n",
    "for i in range(len(validate_data)):\n",
    "    v_raw_x = process_raw_x(validate_data[i], n_past, n_inpaint, n_future)\n",
    "    # choose another song any number you like\n",
    "    p = random.randint(0, len(validate_data) - 1)\n",
    "    t_raw_x = process_raw_x(validate_data[p], 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",
    "        t_raw_x[k] = t_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",
    "    t_past_px, t_past_rx, t_past_len_x, t_past_nrx, t_past_gd, \\\n",
    "    t_inpaint_px, t_inpaint_rx, t_inpaint_len_x, t_inpaint_nrx, t_inpaint_gd,\\\n",
    "    t_future_px, t_future_rx, t_future_len_x, t_future_nrx, t_future_gd = t_raw_x\n",
    "    t_inpaint_gd_whole = t_inpaint_gd.contiguous().view(-1)\n",
    "    t_past_x = [t_past_px, t_past_rx, t_past_len_x, t_past_nrx, t_past_gd]\n",
    "    t_inpaint_x = [t_inpaint_px, t_inpaint_rx, t_inpaint_len_x, t_inpaint_nrx, t_inpaint_gd]\n",
    "    t_future_x = [t_future_px, t_future_rx, t_future_len_x, t_future_nrx, t_future_gd] \n",
    "    \n",
    "    v_sketch_index, v_sketch_cond = gen_sketch([0, 5, 7], t_inpaint_px, t_inpaint_len_x, t_inpaint_nrx, 4)\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        v_recon_x = model.sketch_generation(v_past_x, v_future_x, v_inpaint_x, v_sketch_index, v_sketch_cond)\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",
    "        v_result = v_recon_x.argmax(-1)\n",
    "    total += 1\n",
    "    output.append(\n",
    "        {\n",
    "         \"past1\": v_past_gd.cpu().detach().numpy(),\n",
    "         \"past2\": t_past_gd.cpu().detach().numpy(),\n",
    "         \"future1\":v_future_gd.cpu().detach().numpy(),\n",
    "         \"future2\":t_future_gd.cpu().detach().numpy(),\n",
    "         \"inpaint\":v_result.cpu().detach().numpy(),\n",
    "         \"gd1\":v_inpaint_gd.cpu().detach().numpy(),\n",
    "         \"gd2\":t_inpaint_gd.cpu().detach().numpy(),\n",
    "         \"acc\":v_acc,\n",
    "         \"nll\":v_loss.item()\n",
    "        }\n",
    "    )\n",
    "    print(i)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"res-exp-sketchnet_control_mixture_2.npy\", output)"
   ]
  }
 ],
 "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
}