XT_REMI_transposed_BPE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eac75c41-86b4-4250-bce1-e23dbbe061f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import secrets\n",
    "import tqdm\n",
    "\n",
    "import zipfile\n",
    "\n",
    "\n",
    "import json\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "import math\n",
    "\n",
    "from torch.optim.lr_scheduler import LambdaLR \n",
    "from torch.optim import Adam \n",
    "from torch.cuda.amp import autocast, GradScaler # Mixed Precision\n",
    "\n",
    "from x_transformer2 import TransformerWrapper, Decoder, AutoregressiveWrapper\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "52331240-37b7-411a-802e-9a0deae13ade",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LrStepTracker:\n",
    "    \"\"\"\n",
    "    ----------\n",
    "    Author: Ryan Marshall\n",
    "    Modified: Damon Gwinn\n",
    "    ----------\n",
    "    Class for custom learn rate scheduler (to be used by torch.optim.lr_scheduler.LambdaLR).\n",
    "    Learn rate for each step (batch) given the warmup steps is:\n",
    "        lr = [ 1/sqrt(d_model) ] * min[ 1/sqrt(step) , step * (warmup_steps)^-1.5 ]\n",
    "    This is from Attention is All you Need (https://arxiv.org/abs/1706.03762)\n",
    "    ----------\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model_dim=1024, warmup_steps=4000, init_steps=0):\n",
    "        # Store Values\n",
    "        self.warmup_steps = warmup_steps\n",
    "        self.model_dim = model_dim\n",
    "        self.init_steps = init_steps\n",
    "\n",
    "        # Begin Calculations\n",
    "        self.invsqrt_dim = (1 / math.sqrt(model_dim))\n",
    "        self.invsqrt_warmup = (1 / (warmup_steps * math.sqrt(warmup_steps)))\n",
    "\n",
    "    # step\n",
    "    def step(self, step):\n",
    "        \"\"\"\n",
    "        ----------\n",
    "        Author: Ryan Marshall\n",
    "        Modified: Damon Gwinn\n",
    "        ----------\n",
    "        Method to pass to LambdaLR. Increments the step and computes the new learn rate.\n",
    "        ----------\n",
    "        \"\"\"\n",
    "\n",
    "        step += self.init_steps\n",
    "        if(step <= self.warmup_steps):\n",
    "            return self.invsqrt_dim * self.invsqrt_warmup * step\n",
    "        else:\n",
    "            invsqrt_step = (1 / math.sqrt(step))\n",
    "            return self.invsqrt_dim * invsqrt_step\n",
    "\n",
    "\n",
    "\n",
    "# get_lr\n",
    "def get_lr(optimizer):\n",
    "    \"\"\"\n",
    "    ----------\n",
    "    Author: Damon Gwinn\n",
    "    ----------\n",
    "    Hack to get the current learn rate of the model\n",
    "    ----------\n",
    "    \"\"\"\n",
    "\n",
    "    for param_group in optimizer.param_groups:\n",
    "        return param_group['lr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a7dc7466-d48c-40cf-9fd4-eab6a327ae32",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens_train = list()\n",
    "train_path = \"data_remi_norm_bpe/train\"\n",
    "for (dirpath, dirnames, filenames) in os.walk(train_path):\n",
    "    # Filter files with only .json\n",
    "    tokens_train += [os.path.join(dirpath, file) for file in filenames if file.endswith(\".json\")]\n",
    "\n",
    "tokens_train_aug = list()\n",
    "train_path = \"data_remi_norm_bpe_aug/train\"\n",
    "for (dirpath, dirnames, filenames) in os.walk(train_path):\n",
    "    # Filter files with only .json\n",
    "    tokens_train_aug += [os.path.join(dirpath, file) for file in filenames if file.endswith(\".json\")]\n",
    "\n",
    "tokens_val = list()\n",
    "val_path = \"data_remi_norm_bpe/val\"\n",
    "for (dirpath, dirnames, filenames) in os.walk(val_path):\n",
    "    # Filter files with only .json\n",
    "    tokens_val += [os.path.join(dirpath, file) for file in filenames if file.endswith(\".json\")]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "30b8d911-7e5f-4af9-8b0a-13e2e2da419c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train len: 21919, Train Aug len: 136653, Val len: 1159\n"
     ]
    }
   ],
   "source": [
    "print(f\"Train len: {len(tokens_train)}, Val len: {len(tokens_val)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3e74f8f4-ee50-4c9b-8c5a-b03bb4a13e8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21919/21919 [10:50<00:00, 33.69it/s]\n",
      "100%|██████████| 1159/1159 [00:01<00:00, 643.21it/s]\n"
     ]
    }
   ],
   "source": [
    "train_data = torch.Tensor([0, 0, 0, 0])\n",
    "\n",
    "for f in tqdm.tqdm(tokens_train):\n",
    "    a = json.load(open(f))['ids'][0]\n",
    "    a.append(2001) # Append END_TOKEN\n",
    "    train_data = torch.cat((train_data, torch.Tensor(a)))\n",
    "\n",
    "val_data = torch.Tensor([0, 0, 0, 0])\n",
    "\n",
    "for f in tqdm.tqdm(tokens_val):\n",
    "    a = json.load(open(f))['ids'][0]\n",
    "    a.append(2001) # Append END_TOKEN\n",
    "    val_data = torch.cat((val_data, torch.Tensor(a)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "17b60aee-8595-4f02-9a64-d54d3ce8abd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save or load the data\n",
    "torch.save(train_data, 'train_data_r_norm_bpe_tensor.pt')\n",
    "torch.save(val_data, 'val_data_r_norm_bpe_tensor.pt')\n",
    "#train_data = torch.load('train_data_r_norm_bpe_tensor.pt')\n",
    "#val_data = torch.load('val_data_r_norm_bpe_tensor.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "42022d7b-c47c-4a0f-990f-9f7e930d9d00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the parameters and the data loaders\n",
    "lr_init = None\n",
    "batch_size = 2\n",
    "max_sequence = 1024\n",
    "model_dim = 1024\n",
    "model_depth = 32\n",
    "epochs = 10\n",
    "gradient_accum = 16\n",
    "\n",
    "num_batches = (len(train_data) // max_sequence // batch_size) * epochs\n",
    "\n",
    "VALIDATE_EVERY  = 500\n",
    "SAVE_EVERY = 5000\n",
    "SAVE_STATS_EVERY = 500\n",
    "\n",
    "# helpers\n",
    "\n",
    "def cycle(loader):\n",
    "    while True:\n",
    "        for data in loader:\n",
    "            yield data\n",
    "\n",
    "# Dataloader\n",
    "\n",
    "class MusicDataset(Dataset):\n",
    "    def __init__(self, data, seq_len):\n",
    "        super().__init__()\n",
    "        self.data = data\n",
    "        self.seq_len = seq_len\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        # random sampling\n",
    "        \n",
    "        idx = secrets.randbelow(self.data.size(0) - self.seq_len - 1)\n",
    "        full_seq = self.data[idx: idx + self.seq_len + 1].long()\n",
    "        \n",
    "        return full_seq.cuda()\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data.size(0)\n",
    "\n",
    "train_dataset = MusicDataset(train_data, max_sequence)\n",
    "val_dataset   = MusicDataset(val_data, max_sequence)\n",
    "train_loader  = cycle(DataLoader(train_dataset, batch_size = batch_size))\n",
    "val_loader    = cycle(DataLoader(val_dataset, batch_size = batch_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "02eadeff-bb1e-4280-9f48-dc0b8c351a62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function for getting a moving average of the last 500 losses/ accuracies\n",
    "def last_500(lst):\n",
    "    if len(lst) > 500:\n",
    "        return lst[-500:]\n",
    "    else:\n",
    "        return lst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fac26216-3bf0-42f6-b8c0-c2c169910a99",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): AutoregressiveWrapper(\n",
       "    (net): TransformerWrapper(\n",
       "      (token_emb): TokenEmbedding(\n",
       "        (emb): Embedding(2002, 1024)\n",
       "      )\n",
       "      (pos_emb): AbsolutePositionalEmbedding(\n",
       "        (emb): Embedding(1024, 1024)\n",
       "      )\n",
       "      (post_emb_norm): Identity()\n",
       "      (emb_dropout): Dropout(p=0.0, inplace=False)\n",
       "      (project_emb): Identity()\n",
       "      (attn_layers): Decoder(\n",
       "        (layers): ModuleList(\n",
       "          (0): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (1): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (2): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (3): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (4): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (5): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (6): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (7): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (8): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (9): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (10): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (11): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (12): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (13): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (14): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (15): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (16): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (17): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (18): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (19): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (20): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (21): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (22): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (23): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (24): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (25): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (26): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (27): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (28): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (29): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (30): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (31): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (32): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (33): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (34): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (35): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (36): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (37): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (38): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (39): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (40): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (41): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (42): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (43): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (44): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (45): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (46): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (47): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (48): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (49): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (50): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (51): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (52): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (53): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (54): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (55): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (56): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (57): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (58): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (59): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (60): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (61): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (62): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): Attention(\n",
       "              (to_q): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_k): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (to_v): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "              (attend): Attend(\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (to_out): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "          (63): ModuleList(\n",
       "            (0): ModuleList(\n",
       "              (0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (1-2): 2 x None\n",
       "            )\n",
       "            (1): FeedForward(\n",
       "              (ff): Sequential(\n",
       "                (0): Sequential(\n",
       "                  (0): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "                  (1): GELU(approximate='none')\n",
       "                )\n",
       "                (1): Identity()\n",
       "                (2): Dropout(p=0.15, inplace=False)\n",
       "                (3): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              )\n",
       "            )\n",
       "            (2): Residual()\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "      (to_logits): Linear(in_features=1024, out_features=2002, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize the model\n",
    "model = TransformerWrapper(\n",
    "    num_tokens = 2002,\n",
    "    max_seq_len = max_sequence,\n",
    "    attn_layers = Decoder(\n",
    "        dim = model_dim,\n",
    "        depth = model_depth,\n",
    "        heads = 16,\n",
    "        layer_dropout = 0.15,   # stochastic depth - dropout entire layer\n",
    "        attn_dropout = 0.1,    # dropout post-attention\n",
    "        ff_dropout = 0.15       # feedforward dropout\n",
    "    ))\n",
    "model = AutoregressiveWrapper(model)\n",
    "model = torch.nn.DataParallel(model)\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ea2e3523-e63e-4391-b3c7-8d27e822fdf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the learning rate params\n",
    "ADAM_BETA_1             = 0.9\n",
    "ADAM_BETA_2             = 0.98\n",
    "ADAM_EPSILON            = 10e-9\n",
    "\n",
    "LR_DEFAULT_START        = 1.0\n",
    "\n",
    "if lr_init is None: \n",
    "    init_step = 0 \n",
    "    lr = LR_DEFAULT_START \n",
    "    lr_stepper = LrStepTracker(1024, 4000, 0)\n",
    "else: \n",
    "    lr = lr_init\n",
    "\n",
    "opt = Adam(model.parameters(), lr=lr, betas=(ADAM_BETA_1, ADAM_BETA_2), eps=ADAM_EPSILON)\n",
    "\n",
    "if lr_init is None: \n",
    "    lr_scheduler = LambdaLR(opt, lr_stepper.step) #LambdaLR(optimizer, lr_lambda=lf,last_epoch = start_epoch)\n",
    "else:\n",
    "    lr_scheduler = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ab2263d7-c16a-4686-97b5-f600a5c0b5c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the scaler for mixed-precision training\n",
    "scaler = GradScaler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "264efdc7-5833-45f1-a9b6-368e9285b69c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the stats variables and the save directory\n",
    "ckpt_dir = \"./ckpt\"\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "\n",
    "train_accs = []\n",
    "val_accs = []\n",
    "\n",
    "train_losses_a = 0\n",
    "val_losses_a = 0\n",
    "\n",
    "train_accs_a = 0\n",
    "val_accs_a = 0\n",
    "\n",
    "last_val_accs_a = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "08fbf76d-1767-4305-8ab9-e249a8a4f3e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load saved model\n",
    "'''\n",
    "checkpoint = torch.load(\"./ckpt/latest_bpe.pth\", map_location=\"cpu\")\n",
    "model.load_state_dict(checkpoint['state_dict'])\n",
    "opt.load_state_dict(checkpoint['optimizer'])\n",
    "lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
    "scaler.load_state_dict(checkpoint['scaler'])\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f031c76-4fad-4c82-ae15-333c78761e44",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training:   9%|▉         | 30322/328390 [1:26:32<12:29:26,  6.63it/s, stats=TL: 1.202, TA: 0.687, VL: 1.448, VA: 0.648 -----]"
     ]
    }
   ],
   "source": [
    "# Training loop (saving periodically)\n",
    "with tqdm.tqdm(range(num_batches), mininterval=10., desc='Training') as pbar:\n",
    "    for i in pbar:\n",
    "        model.train()\n",
    "\n",
    "        # Using mixed-precision\n",
    "        with autocast():\n",
    "            loss, acc = model(next(train_loader))\n",
    "            # Backward pass with autocasting and gradient scaling\n",
    "        scaler.scale(loss).backward()\n",
    "            \n",
    "        if i % SAVE_STATS_EVERY == 0:\n",
    "            with open('train_acc.txt', 'a') as file:\n",
    "                file.write(f\"{acc.mean().item()}\\n\")\n",
    "            with open('train_loss.txt', 'a') as file:\n",
    "                file.write(f\"{loss.mean().item()}\\n\")\n",
    "        \n",
    "        # Keep track of the losses and accs\n",
    "        train_losses.append(loss.mean().item())\n",
    "        train_accs.append(acc.mean().item())\n",
    "        train_losses_a = sum(last_5000(train_losses))/len(last_5000(train_losses))\n",
    "        train_accs_a = sum(last_5000(train_accs))/len(last_5000(train_accs))\n",
    "\n",
    "        # Accumulate gradients\n",
    "        if ((i + 1) % gradient_accum == 0) or (i + 1 == num_batches):\n",
    "        # Gradient scaling step and optimizer update\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 2)\n",
    "            scaler.step(opt)\n",
    "            scaler.update()\n",
    "            opt.zero_grad()\n",
    "            lr_scheduler.step() \n",
    "        # Show stats\n",
    "        pbar.set_postfix(stats=f\" TL: {round(train_losses_a, 3)}, TA: {round(train_accs_a, 3)}, VL: {round(val_losses_a, 3)}, VA: {round(val_accs_a, 3)} -----\")\n",
    "    \n",
    "        # Validation loop\n",
    "        if i % VALIDATE_EVERY == 0:\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                with autocast():\n",
    "                    val_loss, val_acc = model(next(val_loader))\n",
    "                \n",
    "                val_losses.append(val_loss.mean().item())\n",
    "                val_accs.append(val_acc.mean().item())\n",
    "                val_losses_a = sum(last_5000(val_losses))/len(last_5000(val_losses))\n",
    "                val_accs_a = sum(last_5000(val_accs))/len(last_5000(val_accs))\n",
    "                #print(f\"val acc:{val_accs_a}\")\n",
    "                # Save stats\n",
    "                with open('val_acc.txt', 'a') as file:\n",
    "                    file.write(f\"{val_acc.mean().item()}\\n\")\n",
    "                with open('val_loss.txt', 'a') as file:\n",
    "                    file.write(f\"{val_loss.mean().item()}\\n\")\n",
    "        # Save if current val accuracy is better than previous one\n",
    "        if i % SAVE_EVERY == 0:\n",
    "            if len(val_accs) > 1:\n",
    "                if val_accs_a > last_val_accs_a:\n",
    "                    torch.save({\n",
    "                        'state_dict': model.state_dict(),\n",
    "                        'optimizer': opt.state_dict(),\n",
    "                        'lr_scheduler': lr_scheduler.state_dict(),\n",
    "                        'scaler': scaler.state_dict()\n",
    "                    }, os.path.join(ckpt_dir, 'latest_bpe.pth'))\n",
    "                    #print(f\"last acc before:{last_val_accs_a}\")\n",
    "                    last_val_accs_a = val_accs_a\n",
    "                    #print(f\"last acc:{last_val_accs_a}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2c456d24-78da-4ecb-89fa-f3dbc8010e34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Last save\n",
    "torch.save({\n",
    "    'state_dict': model.state_dict(),\n",
    "    'optimizer': opt.state_dict(),\n",
    "    'lr_scheduler': lr_scheduler.state_dict(),\n",
    "    'scaler': scaler.state_dict(),\n",
    "}, os.path.join(ckpt_dir, 'latest_bpe.pth'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}