transfer_learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.utils.data\n",
    "import tqdm\n",
    "import muspy\n",
    "import src.utils as utils\n",
    "import src.representation as representation\n",
    "import src.dataset as dataset\n",
    "import src.music_x_transformers as music_x_transformers\n",
    "import pathlib\n",
    "import src.advUtils as advUtils\n",
    "import shutil\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load configurations\n",
    "train_args = utils.load_json(\".\\pre_trained_models\\mmt_sod_ape_training_logs.json\")\n",
    "encoding = representation.load_encoding(\"encoding.json\")\n",
    "\n",
    "sos = encoding[\"type_code_map\"][\"start-of-song\"]\n",
    "eos = encoding[\"type_code_map\"][\"end-of-song\"]\n",
    "beat_0 = encoding[\"beat_code_map\"][0]\n",
    "beat_4 = encoding[\"beat_code_map\"][4]\n",
    "beat_16 = encoding[\"beat_code_map\"][16]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load training/testing/demo Data\n",
    "data_set = advUtils.convert_extract_load(train_args,encoding, json_dir = \"./data/test/json\",repr_dir=\"./data/test/repr\")\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    data_set,\n",
    "    shuffle=True,\n",
    "    num_workers=1,\n",
    "    collate_fn=dataset.MusicDataset.collate,\n",
    ")\n",
    "\n",
    "test_loader = data_loader\n",
    "train_loader = data_loader\n",
    "valid_loader = data_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating the model...\n",
      "Loaded the pretrained model weights\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MusicXTransformer(\n",
       "  (decoder): MusicAutoregressiveWrapper(\n",
       "    (net): MusicTransformerWrapper(\n",
       "      (token_emb): ModuleList(\n",
       "        (0): TokenEmbedding(\n",
       "          (emb): Embedding(5, 512)\n",
       "        )\n",
       "        (1): TokenEmbedding(\n",
       "          (emb): Embedding(257, 512)\n",
       "        )\n",
       "        (2): TokenEmbedding(\n",
       "          (emb): Embedding(13, 512)\n",
       "        )\n",
       "        (3): TokenEmbedding(\n",
       "          (emb): Embedding(129, 512)\n",
       "        )\n",
       "        (4): TokenEmbedding(\n",
       "          (emb): Embedding(33, 512)\n",
       "        )\n",
       "        (5): TokenEmbedding(\n",
       "          (emb): Embedding(65, 512)\n",
       "        )\n",
       "      )\n",
       "      (pos_emb): AbsolutePositionalEmbedding(\n",
       "        (emb): Embedding(1024, 512)\n",
       "      )\n",
       "      (emb_dropout): Dropout(p=0.2, inplace=False)\n",
       "      (project_emb): Identity()\n",
       "      (attn_layers): Decoder(\n",
       "        (layers): ModuleList(\n",
       "          (0): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (dropout): Dropout(p=0.2, inplace=False)\n",
       "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (1): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (1): GELU()\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.2, inplace=False)\n",
       "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (2): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (dropout): Dropout(p=0.2, inplace=False)\n",
       "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (3): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (1): GELU()\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.2, inplace=False)\n",
       "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (4): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (dropout): Dropout(p=0.2, inplace=False)\n",
       "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (5): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (1): GELU()\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.2, inplace=False)\n",
       "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (6): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (dropout): Dropout(p=0.2, inplace=False)\n",
       "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (7): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (1): GELU()\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.2, inplace=False)\n",
       "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (8): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (dropout): Dropout(p=0.2, inplace=False)\n",
       "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (9): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (1): GELU()\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.2, inplace=False)\n",
       "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (10): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
       "              (dropout): Dropout(p=0.2, inplace=False)\n",
       "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (11): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "              (1): None\n",
       "              (2): None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
       "                  (1): GELU()\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.2, inplace=False)\n",
       "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      (to_logits): ModuleList(\n",
       "        (0): Linear(in_features=512, out_features=5, bias=True)\n",
       "        (1): Linear(in_features=512, out_features=257, bias=True)\n",
       "        (2): Linear(in_features=512, out_features=13, bias=True)\n",
       "        (3): Linear(in_features=512, out_features=129, bias=True)\n",
       "        (4): Linear(in_features=512, out_features=33, bias=True)\n",
       "        (5): Linear(in_features=512, out_features=65, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load model\n",
    "device = torch.device(\"cpu\")\n",
    "print(f\"Creating the model...\")\n",
    "model = music_x_transformers.MusicXTransformer(\n",
    "    dim=train_args[\"dim\"],\n",
    "    encoding=encoding,\n",
    "    depth=train_args[\"layers\"],\n",
    "    heads=train_args[\"heads\"],\n",
    "    max_seq_len=train_args[\"max_seq_len\"],\n",
    "    max_beat=train_args[\"max_beat\"],\n",
    "    rotary_pos_emb=train_args[\"rel_pos_emb\"],\n",
    "    use_abs_pos_emb=train_args[\"abs_pos_emb\"],\n",
    "    emb_dropout=train_args[\"dropout\"],\n",
    "    attn_dropout=train_args[\"dropout\"],\n",
    "    ff_dropout=train_args[\"dropout\"],\n",
    ").to(device)\n",
    "\n",
    "origModel = music_x_transformers.MusicXTransformer(\n",
    "    dim=train_args[\"dim\"],\n",
    "    encoding=encoding,\n",
    "    depth=train_args[\"layers\"],\n",
    "    heads=train_args[\"heads\"],\n",
    "    max_seq_len=train_args[\"max_seq_len\"],\n",
    "    max_beat=train_args[\"max_beat\"],\n",
    "    rotary_pos_emb=train_args[\"rel_pos_emb\"],\n",
    "    use_abs_pos_emb=train_args[\"abs_pos_emb\"],\n",
    "    emb_dropout=train_args[\"dropout\"],\n",
    "    attn_dropout=train_args[\"dropout\"],\n",
    "    ff_dropout=train_args[\"dropout\"],\n",
    ").to(device)\n",
    "\n",
    "firstTrain = model\n",
    "origModel.load_state_dict(torch.load(\"./pre_trained_models/mmt_sod_ape_best_model.pt\", map_location=device))\n",
    "model.load_state_dict(torch.load(\"./pre_trained_models/mmt_sod_ape_best_model.pt\", map_location=device))\n",
    "firstTrain.load_state_dict(torch.load(\"./experiment/checkpoints/model_1000.pt\", map_location=device))\n",
    "\n",
    "print(f\"Loaded the pretrained model weights\")\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(n,sample_dir,model=model,modes=[\"unconditioned\"],seq_len=1024,temperature=1,filter_logits=\"top_k\",filter_thresh=0.9):\n",
    "    with torch.no_grad():\n",
    "        data_iter = iter(test_loader)\n",
    "        for i in tqdm.tqdm(range(n), ncols=80):\n",
    "            batch = next(data_iter)\n",
    "            print(\"Generating based on\",batch['name'])\n",
    "            for mode in modes:\n",
    "                if(mode==\"unconditioned\"):\n",
    "                    tgt_start = torch.zeros((1, 1, 6), dtype=torch.long, device=device)\n",
    "                    tgt_start[:, 0, 0] = sos\n",
    "                elif(mode==\"instrument_informed\"):\n",
    "                    prefix_len = int(np.argmax(batch[\"seq\"][0, :, 1] >= beat_0))\n",
    "                    tgt_start = batch[\"seq\"][:1, :prefix_len].to(device)\n",
    "                elif(mode==\"4_beat\"):\n",
    "                    cond_len = int(np.argmax(batch[\"seq\"][0, :, 1] >= beat_4))\n",
    "                    tgt_start = batch[\"seq\"][:1, :cond_len].to(device)\n",
    "                elif(mode==\"16_beat\"):\n",
    "                    cond_len = int(np.argmax(batch[\"seq\"][0, :, 1] >= beat_16))\n",
    "                    tgt_start = batch[\"seq\"][:1, :cond_len].to(device)\n",
    "                # Generate new samples\n",
    "                generated = model.generate(\n",
    "                    tgt_start,\n",
    "                    seq_len,\n",
    "                    eos_token=eos,\n",
    "                    temperature=temperature,\n",
    "                    filter_logits_fn=filter_logits,\n",
    "                    filter_thres=filter_thresh,\n",
    "                    monotonicity_dim=(\"type\", \"beat\"),\n",
    "                )\n",
    "                generated_np = torch.cat((tgt_start, generated), 1).cpu().numpy()\n",
    "\n",
    "                # Save the results\n",
    "                advUtils.save_result(\n",
    "                    f\"{i}_{mode}\", generated_np[0], sample_dir, encoding,savecsv=False,savetxt=False,savenpy=False,savepng=False,savejson=False\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Test generation\n",
    "#generate(2,\"./samples\",seq_len=50,modes=[\"unconditioned\",\"instrument_informed\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Transfer Learning</h1>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 19944950\n",
      "Number of trainable parameters: 257536\n"
     ]
    }
   ],
   "source": [
    "#Freeze All network layers\n",
    "for m in model.parameters():\n",
    "    m.requires_grad = False\n",
    "\n",
    "layers_to_unfreeze = [\"decoder.net.to_logits.5.weight\",\"decoder.net.to_logits.4.weight\",\"decoder.net.to_logits.3.weight\",\"decoder.net.to_logits.2.weight\",\"decoder.net.to_logits.1.weight\",\"decoder.net.to_logits.0.weight\",\"decoder.net.norm.weight\"]\n",
    "\n",
    "for m in model.named_parameters():\n",
    "    state_dict = model.state_dict()\n",
    "    if(m[0] in layers_to_unfreeze):\n",
    "      newLayer = torch.rand(m[1].shape,requires_grad=True)\n",
    "      if(m[0]==\"decoder.net.norm.weight\"):\n",
    "         newLayer=torch.add(newLayer,1)\n",
    "      else:\n",
    "         newLayer=torch.add(newLayer,-0.5)\n",
    "      state_dict[m[0]] = newLayer\n",
    "      model.load_state_dict(state_dict)\n",
    "\n",
    "for m in model.named_parameters():\n",
    "    if(m[0] in layers_to_unfreeze):\n",
    "       m[1].requires_grad=True\n",
    "\n",
    "n_parameters = sum(p.numel() for p in model.parameters())\n",
    "n_trainables = sum(\n",
    "   p.numel() for p in model.parameters() if p.requires_grad\n",
    ")\n",
    "print(f\"Number of parameters: {n_parameters}\")\n",
    "print(f\"Number of trainable parameters: {n_trainables}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<bound method Module.state_dict of MusicXTransformer(\n",
      "  (decoder): MusicAutoregressiveWrapper(\n",
      "    (net): MusicTransformerWrapper(\n",
      "      (token_emb): ModuleList(\n",
      "        (0): TokenEmbedding(\n",
      "          (emb): Embedding(5, 512)\n",
      "        )\n",
      "        (1): TokenEmbedding(\n",
      "          (emb): Embedding(257, 512)\n",
      "        )\n",
      "        (2): TokenEmbedding(\n",
      "          (emb): Embedding(13, 512)\n",
      "        )\n",
      "        (3): TokenEmbedding(\n",
      "          (emb): Embedding(129, 512)\n",
      "        )\n",
      "        (4): TokenEmbedding(\n",
      "          (emb): Embedding(33, 512)\n",
      "        )\n",
      "        (5): TokenEmbedding(\n",
      "          (emb): Embedding(65, 512)\n",
      "        )\n",
      "      )\n",
      "      (pos_emb): AbsolutePositionalEmbedding(\n",
      "        (emb): Embedding(1024, 512)\n",
      "      )\n",
      "      (emb_dropout): Dropout(p=0.2, inplace=False)\n",
      "      (project_emb): Identity()\n",
      "      (attn_layers): Decoder(\n",
      "        (layers): ModuleList(\n",
      "          (0): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): Attention(\n",
      "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (dropout): Dropout(p=0.2, inplace=False)\n",
      "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (1): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): FeedForward(\n",
      "              (net): Sequential(\n",
      "                (0): Sequential(\n",
      "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
      "                  (1): GELU()\n",
      "                )\n",
      "                (1): Identity()\n",
      "                (2): Dropout(p=0.2, inplace=False)\n",
      "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
      "              )\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (2): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): Attention(\n",
      "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (dropout): Dropout(p=0.2, inplace=False)\n",
      "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (3): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): FeedForward(\n",
      "              (net): Sequential(\n",
      "                (0): Sequential(\n",
      "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
      "                  (1): GELU()\n",
      "                )\n",
      "                (1): Identity()\n",
      "                (2): Dropout(p=0.2, inplace=False)\n",
      "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
      "              )\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (4): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): Attention(\n",
      "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (dropout): Dropout(p=0.2, inplace=False)\n",
      "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (5): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): FeedForward(\n",
      "              (net): Sequential(\n",
      "                (0): Sequential(\n",
      "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
      "                  (1): GELU()\n",
      "                )\n",
      "                (1): Identity()\n",
      "                (2): Dropout(p=0.2, inplace=False)\n",
      "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
      "              )\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (6): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): Attention(\n",
      "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (dropout): Dropout(p=0.2, inplace=False)\n",
      "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (7): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): FeedForward(\n",
      "              (net): Sequential(\n",
      "                (0): Sequential(\n",
      "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
      "                  (1): GELU()\n",
      "                )\n",
      "                (1): Identity()\n",
      "                (2): Dropout(p=0.2, inplace=False)\n",
      "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
      "              )\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (8): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): Attention(\n",
      "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (dropout): Dropout(p=0.2, inplace=False)\n",
      "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (9): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): FeedForward(\n",
      "              (net): Sequential(\n",
      "                (0): Sequential(\n",
      "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
      "                  (1): GELU()\n",
      "                )\n",
      "                (1): Identity()\n",
      "                (2): Dropout(p=0.2, inplace=False)\n",
      "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
      "              )\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (10): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): Attention(\n",
      "              (to_q): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_k): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (to_v): Linear(in_features=512, out_features=512, bias=False)\n",
      "              (dropout): Dropout(p=0.2, inplace=False)\n",
      "              (to_out): Linear(in_features=512, out_features=512, bias=True)\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "          (11): ModuleList(\n",
      "            (0): ModuleList(\n",
      "              (0): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "              (1): None\n",
      "              (2): None\n",
      "            )\n",
      "            (1): FeedForward(\n",
      "              (net): Sequential(\n",
      "                (0): Sequential(\n",
      "                  (0): Linear(in_features=512, out_features=2048, bias=True)\n",
      "                  (1): GELU()\n",
      "                )\n",
      "                (1): Identity()\n",
      "                (2): Dropout(p=0.2, inplace=False)\n",
      "                (3): Linear(in_features=2048, out_features=512, bias=True)\n",
      "              )\n",
      "            )\n",
      "            (2): Residual()\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
      "      (to_logits): ModuleList(\n",
      "        (0): Linear(in_features=512, out_features=5, bias=True)\n",
      "        (1): Linear(in_features=512, out_features=257, bias=True)\n",
      "        (2): Linear(in_features=512, out_features=13, bias=True)\n",
      "        (3): Linear(in_features=512, out_features=129, bias=True)\n",
      "        (4): Linear(in_features=512, out_features=33, bias=True)\n",
      "        (5): Linear(in_features=512, out_features=65, bias=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")>\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "zerosModel = model\n",
    "for name, param in zerosModel.named_parameters():\n",
    "    if(name in layers_to_unfreeze):\n",
    "        zerosModel.state_dict()[name].fill_(0)\n",
    "print(zerosModel.state_dict)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                     | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating based on ['albinoni_sonate_da_chiesa_6_(c)icking-archive']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|██████████████████████▌                      | 1/2 [00:48<00:48, 48.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating based on ['anglebert_fugue_3_(c)mccoy']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 2/2 [01:49<00:00, 54.67s/it]\n",
      "  0%|                                                     | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating based on ['anglebert_fugue_3_(c)mccoy']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|██████████████████████▌                      | 1/2 [00:21<00:21, 21.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating based on ['albinoni_sonate_da_chiesa_6_(c)icking-archive']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 2/2 [00:34<00:00, 17.04s/it]\n"
     ]
    }
   ],
   "source": [
    "#comparative generation:\n",
    "generate(2,\"./samples\",seq_len=50,model=origModel,modes=[\"unconditioned\",\"instrument_informed\",\"4_beat\",\"16_beat\"])\n",
    "generate(2,\"./samples2\",seq_len=50,model=firstTrain,modes=[\"unconditioned\",\"instrument_informed\",\"4_beat\",\"16_beat\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Training the model\n",
    "def get_lr_multiplier(\n",
    "    step, warmup_steps, decay_end_steps, decay_end_multiplier\n",
    "):\n",
    "    \"\"\"Return the learning rate multiplier with a warmup and decay schedule.\n",
    "\n",
    "    The learning rate multiplier starts from 0 and linearly increases to 1\n",
    "    after `warmup_steps`. After that, it linearly decreases to\n",
    "    `decay_end_multiplier` until `decay_end_steps` is reached.\n",
    "\n",
    "    \"\"\"\n",
    "    if step < warmup_steps:\n",
    "        return (step + 1) / warmup_steps\n",
    "    if step > decay_end_steps:\n",
    "        return decay_end_multiplier\n",
    "    position = (step - warmup_steps) / (decay_end_steps - warmup_steps)\n",
    "    return 1 - (1 - decay_end_multiplier) * position\n",
    "\n",
    "\n",
    "def train(out_dir,model):\n",
    "    \"\"\"Main function.\"\"\"\n",
    "    # Parse the command-line arguments\n",
    "\n",
    "    # Make sure the output directory exists\n",
    "    pathlib.Path(out_dir).mkdir(exist_ok=True)\n",
    "    pathlib.Path(out_dir+\"/checkpoints\").mkdir(exist_ok=True)\n",
    "    # Get the specified device\n",
    "    device = torch.device(\"cpu\")\n",
    "    \n",
    "\n",
    "    # Summarize the model\n",
    "    n_parameters = sum(p.numel() for p in model.parameters())\n",
    "    n_trainables = sum(\n",
    "        p.numel() for p in model.parameters() if p.requires_grad\n",
    "    )\n",
    "    print(f\"Number of parameters: {n_parameters}\")\n",
    "    print(f\"Number of trainable parameters: {n_trainables}\")\n",
    "\n",
    "    # Create the optimizer\n",
    "    optimizer = torch.optim.Adam(model.parameters(), train_args[\"learning_rate\"])\n",
    "    scheduler = torch.optim.lr_scheduler.LambdaLR(\n",
    "        optimizer,\n",
    "        lr_lambda=lambda step: get_lr_multiplier(\n",
    "            step,\n",
    "            train_args[\"lr_warmup_steps\"],\n",
    "            train_args[\"lr_decay_steps\"],\n",
    "            train_args[\"lr_decay_multiplier\"],\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    # Create a file to record losses\n",
    "    loss_csv = open(out_dir+\"/loss.csv\", \"w\")\n",
    "    loss_csv.write(\n",
    "        \"step,train_loss,valid_loss,type_loss,beat_loss,position_loss,\"\n",
    "        \"pitch_loss,duration_loss,instrument_loss\\n\"\n",
    "    )\n",
    "    # Initialize variables\n",
    "    step = 0\n",
    "    min_val_loss = float(\"inf\")\n",
    "    if train_args[\"early_stopping\"]:\n",
    "        count_early_stopping = 0\n",
    "\n",
    "    # Iterate for the specified number of steps\n",
    "    train_iterator = iter(train_loader)\n",
    "    while step < train_args[\"steps\"]:\n",
    "\n",
    "        # Training\n",
    "        print(f\"Training...\")\n",
    "        model.train()\n",
    "        recent_losses = []\n",
    "\n",
    "        for batch in (pbar := tqdm.tqdm(range(train_args[\"valid_steps\"]), ncols=80)):\n",
    "            # Get next batch\n",
    "            try:\n",
    "                batch = next(train_iterator)\n",
    "            except StopIteration:\n",
    "                # Reinitialize dataset iterator\n",
    "                train_iterator = iter(train_loader)\n",
    "                batch = next(train_iterator)\n",
    "\n",
    "            # Get input and output pair\n",
    "            seq = batch[\"seq\"].to(device)\n",
    "            mask = batch[\"mask\"].to(device)\n",
    "\n",
    "            # Update the model parameters\n",
    "            optimizer.zero_grad()\n",
    "            loss = model(seq, mask=mask)\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(\n",
    "                model.parameters(), train_args[\"grad_norm_clip\"]\n",
    "            )\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "\n",
    "            # Compute the moving average of the loss\n",
    "            recent_losses.append(float(loss))\n",
    "            if len(recent_losses) > 10:\n",
    "                del recent_losses[0]\n",
    "            train_loss = np.mean(recent_losses)\n",
    "            pbar.set_postfix(loss=f\"{train_loss:8.4f}\")\n",
    "\n",
    "            step += 1\n",
    "\n",
    "        # Release GPU memory right away\n",
    "        del seq, mask\n",
    "\n",
    "        # Validation\n",
    "        print(f\"Validating...\")\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            total_loss = 0\n",
    "            total_losses = [0] * 6\n",
    "            count = 0\n",
    "            for batch in valid_loader:\n",
    "                # Get input and output pair\n",
    "                seq = batch[\"seq\"].to(device)\n",
    "                mask = batch[\"mask\"].to(device)\n",
    "\n",
    "                # Pass through the model\n",
    "                loss, losses = model(seq, return_list=True, mask=mask)\n",
    "\n",
    "                # Accumulate validation loss\n",
    "                count += len(batch)\n",
    "                total_loss += len(batch) * float(loss)\n",
    "                for idx in range(6):\n",
    "                    total_losses[idx] += float(losses[idx])\n",
    "        val_loss = total_loss / count\n",
    "        individual_losses = [l / count for l in total_losses]\n",
    "        print(f\"Validation loss: {val_loss:.4f}\")\n",
    "        print(\n",
    "            f\"Individual losses: type={individual_losses[0]:.4f}, \"\n",
    "            f\"beat: {individual_losses[1]:.4f}, \"\n",
    "            f\"position: {individual_losses[2]:.4f}, \"\n",
    "            f\"pitch: {individual_losses[3]:.4f}, \"\n",
    "            f\"duration: {individual_losses[4]:.4f}, \"\n",
    "            f\"instrument: {individual_losses[5]:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Release GPU memory right away\n",
    "        del seq, mask\n",
    "\n",
    "        # Write losses to file\n",
    "        loss_csv.write(\n",
    "            f\"{step},{train_loss},{val_loss},{individual_losses[0]},\"\n",
    "            f\"{individual_losses[1]},{individual_losses[2]},\"\n",
    "            f\"{individual_losses[3]},{individual_losses[4]},\"\n",
    "            f\"{individual_losses[5]}\\n\"\n",
    "        )\n",
    "\n",
    "        # Save the model\n",
    "        checkpoint_filename = out_dir+\"/checkpoints/\"+f\"model_{step}.pt\"\n",
    "        torch.save(model.state_dict(), checkpoint_filename)\n",
    "        print(f\"Saved the model to: {checkpoint_filename}\")\n",
    "\n",
    "        # Copy the model if it is the best model so far\n",
    "        if val_loss < min_val_loss:\n",
    "            min_val_loss = val_loss\n",
    "            shutil.copyfile(\n",
    "                checkpoint_filename,\n",
    "                out_dir+\"/checkpoints/\"+\"best_model.pt\",\n",
    "            )\n",
    "            # Reset the early stopping counter if we found a better model\n",
    "            if train_args[\"early_stopping\"]:\n",
    "                count_early_stopping = 0\n",
    "        elif train_args[\"early_stopping\"]:\n",
    "            # Increment the early stopping counter if no improvement is found\n",
    "            count_early_stopping += 1\n",
    "\n",
    "        # Early stopping\n",
    "        if (\n",
    "            train_args[\"early_stopping\"]\n",
    "            and count_early_stopping > train_args[\"early_stopping_tolerance\"]\n",
    "        ):\n",
    "            print(\n",
    "                \"Stopped the training for no improvements in \"\n",
    "                f\"{train_args['early_stopping_tolerance']} rounds.\"\n",
    "            )\n",
    "            break\n",
    "\n",
    "    # Log minimum validation loss\n",
    "    print(f\"Minimum validation loss achieved: {min_val_loss}\")\n",
    "\n",
    "    # Save the optimizer states\n",
    "    optimizer_filename = out_dir+\"/checkpoints/\"+f\"optimizer_{step}.pt\"\n",
    "    torch.save(optimizer.state_dict(), optimizer_filename)\n",
    "    print(f\"Saved the optimizer state to: {optimizer_filename}\")\n",
    "\n",
    "    # Save the scheduler states\n",
    "    scheduler_filename = out_dir+\"checkpoints/\"+f\"scheduler_{step}.pt\"\n",
    "    torch.save(scheduler.state_dict(), scheduler_filename)\n",
    "    print(f\"Saved the scheduler state to: {scheduler_filename}\")\n",
    "\n",
    "    # Close the file\n",
    "    loss_csv.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 19944950\n",
      "Number of trainable parameters: 257536\n",
      "Training...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████| 1000/1000 [9:07:57<00:00, 32.88s/it, loss=58.6785]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validating...\n",
      "Validation loss: 54.9613\n",
      "Individual losses: type=1.8032, beat: 3.1263, position: 1.6815, pitch: 3.2006, duration: 1.6555, instrument: 2.2732\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'builtin_function_or_method' object has no attribute 'info'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[61], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mexperiment\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[60], line 153\u001b[0m, in \u001b[0;36mtrain\u001b[1;34m(out_dir, model)\u001b[0m\n\u001b[0;32m    151\u001b[0m checkpoint_filename \u001b[38;5;241m=\u001b[39m out_dir\u001b[38;5;241m+\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/checkpoints/\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m+\u001b[39m\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstep\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.pt\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    152\u001b[0m torch\u001b[38;5;241m.\u001b[39msave(model\u001b[38;5;241m.\u001b[39mstate_dict(), checkpoint_filename)\n\u001b[1;32m--> 153\u001b[0m \u001b[38;5;28;43mprint\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minfo\u001b[49m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSaved the model to: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcheckpoint_filename\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    155\u001b[0m \u001b[38;5;66;03m# Copy the model if it is the best model so far\u001b[39;00m\n\u001b[0;32m    156\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m val_loss \u001b[38;5;241m<\u001b[39m min_val_loss:\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'builtin_function_or_method' object has no attribute 'info'"
     ]
    }
   ],
   "source": [
    "train(\"experiment\",model)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mtmt",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}