Super_Piano_6_Performer_GPU.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Performer-Pytorch.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyNMiBCV+zp01Da43Z4AVlup",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/asigalov61/SuperPiano/blob/master/Super_Piano_6_Performer_GPU.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ii5xWJ1sRc6x"
      },
      "source": [
        "https://github.com/lucidrains/performer-pytorch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2mfyXs8qBTFA"
      },
      "source": [
        "!pip install performer-pytorch\r\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W7XQB1KbKPrE"
      },
      "source": [
        "%%writefile setup.sh \r\n",
        "\r\n",
        "# install apex to be able to use mix precision\r\n",
        "export CUDA_HOME=/usr/local/cuda-10.1\r\n",
        "git clone https://github.com/NVIDIA/apex\r\n",
        "pip install -v -q --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./apex"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-UND0x9oKQhN"
      },
      "source": [
        "!sh setup.sh"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jDx1DeBNBNs7"
      },
      "source": [
        "from performer_pytorch import PerformerLM\r\n",
        "from performer_pytorch.autoregressive_wrapper import AutoregressiveWrapper\r\n",
        "\r\n",
        "import random\r\n",
        "import tqdm\r\n",
        "import gzip\r\n",
        "import numpy as np\r\n",
        "import torch\r\n",
        "import torch.optim as optim\r\n",
        "from torch.nn import functional as F\r\n",
        "from torch.utils.data import DataLoader, Dataset\r\n",
        "from torch.cuda.amp import autocast, GradScaler"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yi-8-u0IBBjp"
      },
      "source": [
        "# constants\r\n",
        "\r\n",
        "NUM_BATCHES = int(1e5)\r\n",
        "BATCH_SIZE = 16\r\n",
        "GRADIENT_ACCUMULATE_EVERY = 4\r\n",
        "LEARNING_RATE = 1e-4\r\n",
        "VALIDATE_EVERY  = 100\r\n",
        "GENERATE_EVERY  = 500\r\n",
        "GENERATE_LENGTH = 4096\r\n",
        "SEQ_LEN = 4096\r\n",
        "\r\n",
        "# helpers\r\n",
        "\r\n",
        "def cycle(loader):\r\n",
        "    while True:\r\n",
        "        for data in loader:\r\n",
        "            yield data\r\n",
        "\r\n",
        "def decode_token(token):\r\n",
        "    return str(chr(max(32, token)))\r\n",
        "\r\n",
        "def decode_tokens(tokens):\r\n",
        "    return ''.join(list(map(decode_token, tokens)))\r\n",
        "\r\n",
        "# instantiate model\r\n",
        "\r\n",
        "model = PerformerLM(\r\n",
        "    num_tokens = 256,\r\n",
        "    dim = 512,\r\n",
        "    depth = 6,\r\n",
        "    max_seq_len = SEQ_LEN,\r\n",
        "    heads = 8,\r\n",
        "    causal = True,\r\n",
        "    reversible = True,\r\n",
        "    nb_features = 256,\r\n",
        "    use_scalenorm = True,\r\n",
        "    local_attn_heads = (8, 8, 8, 6, 4, 2)\r\n",
        ")\r\n",
        "\r\n",
        "model = AutoregressiveWrapper(model)\r\n",
        "model.cuda()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S0YzpbD5ClFh"
      },
      "source": [
        "# prepare music data\r\n",
        "\r\n",
        "with open('/content/INT_DATASET.TXT') as file:\r\n",
        "    X = np.fromstring(file.read(int(3e6)), dtype=np.uint8)\r\n",
        "    trX, vaX = np.split(X, [int(2e6)])\r\n",
        "    data_train, data_val = torch.from_numpy(trX), torch.from_numpy(vaX)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yMEjUQapComn"
      },
      "source": [
        "class TextSamplerDataset(Dataset):\r\n",
        "    def __init__(self, data, seq_len):\r\n",
        "        super().__init__()\r\n",
        "        self.data = data\r\n",
        "        self.seq_len = seq_len\r\n",
        "\r\n",
        "    def __getitem__(self, index):\r\n",
        "        rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))\r\n",
        "        full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()\r\n",
        "        return full_seq.cuda()\r\n",
        "\r\n",
        "    def __len__(self):\r\n",
        "        return self.data.size(0) // self.seq_len\r\n",
        "\r\n",
        "train_dataset = TextSamplerDataset(data_train, SEQ_LEN)\r\n",
        "val_dataset   = TextSamplerDataset(data_val, SEQ_LEN)\r\n",
        "train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE))\r\n",
        "val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE))\r\n",
        "\r\n",
        "# optimizer\r\n",
        "\r\n",
        "optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\r\n",
        "scaler = GradScaler()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RGMJN1DrCwqg"
      },
      "source": [
        "# training\r\n",
        "\r\n",
        "for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):\r\n",
        "    model.train()\r\n",
        "\r\n",
        "    for __ in range(GRADIENT_ACCUMULATE_EVERY):\r\n",
        "        with autocast():\r\n",
        "            loss = model(next(train_loader), return_loss = True)\r\n",
        "        scaler.scale(loss).backward()\r\n",
        "\r\n",
        "    print(f'training loss: {loss.item()}')\r\n",
        "\r\n",
        "    scaler.unscale_(optim)\r\n",
        "    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\r\n",
        "    scaler.step(optim)\r\n",
        "    scaler.update()\r\n",
        "    optim.zero_grad()\r\n",
        "\r\n",
        "    if i % VALIDATE_EVERY == 0:\r\n",
        "        model.eval()\r\n",
        "        with torch.no_grad():\r\n",
        "            loss = model(next(val_loader), return_loss = True)\r\n",
        "            print(f'validation loss: {loss.item()}')\r\n",
        "\r\n",
        "    if i % GENERATE_EVERY == 0 and i != 0:\r\n",
        "        model.eval()\r\n",
        "        inp = random.choice(val_dataset)[:-1]\r\n",
        "        prime = decode_tokens(inp)\r\n",
        "        print(f'%s \\n\\n %s', (prime, '*' * 100))\r\n",
        "\r\n",
        "        sample = model.generate(inp, GENERATE_LENGTH)\r\n",
        "        output_str = decode_tokens(sample)\r\n",
        "        print(output_str)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ML6iQwFhGXES"
      },
      "source": [
        "model.eval()\r\n",
        "inp = random.choice(val_dataset)[:-1]\r\n",
        "prime = decode_tokens(inp)\r\n",
        "print(f'%s \\n\\n %s', (prime, '*' * 100))\r\n",
        "\r\n",
        "sample = model.generate(inp, GENERATE_LENGTH)\r\n",
        "output_str = decode_tokens(sample)\r\n",
        "print(output_str)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}