train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cc4ac7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import os.path as osp\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from accelerate import Accelerator\n",
    "from accelerate.utils import LoggerType\n",
    "\n",
    "from transformers import AdamW\n",
    "from transformers import AlbertConfig, AlbertModel\n",
    "from accelerate import DistributedDataParallelKwargs\n",
    "\n",
    "from model import MultiTaskModel\n",
    "from dataloader import build_dataloader\n",
    "from utils import length_to_mask, scan_checkpoint\n",
    "\n",
    "from datasets import load_from_disk\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68d0c7ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "import pickle\n",
    "\n",
    "config_path = \"Configs/config.yml\" # you can change it to anything else\n",
    "config = yaml.safe_load(open(config_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23a7f165",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(config['dataset_params']['token_maps'], 'rb') as handle:\n",
    "    token_maps = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "158bf338",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TransfoXLTokenizer, TransfoXLModel\n",
    "tokenizer = TransfoXLTokenizer.from_pretrained(config['dataset_params']['tokenizer'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e60819c",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss() # F0 loss (regression)\n",
    "\n",
    "best_loss = float('inf')  # best test loss\n",
    "start_epoch = 0  # start from epoch 0 or last checkpoint epoch\n",
    "loss_train_record = list([])\n",
    "loss_test_record = list([])\n",
    "\n",
    "num_steps = config['num_steps']\n",
    "log_interval = config['log_interval']\n",
    "save_interval = config['save_interval']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4fa9e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    \n",
    "    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)\n",
    "    \n",
    "    curr_steps = 0\n",
    "    \n",
    "    dataset = load_from_disk(config[\"data_folder\"])\n",
    "\n",
    "    log_dir = config['log_dir']\n",
    "    if not osp.exists(log_dir): os.makedirs(log_dir, exist_ok=True)\n",
    "    shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path)))\n",
    "    \n",
    "    batch_size = config[\"batch_size\"]\n",
    "    train_loader = build_dataloader(dataset, \n",
    "                                    batch_size=batch_size, \n",
    "                                    num_workers=0, \n",
    "                                    dataset_config=config['dataset_params'])\n",
    "\n",
    "    albert_base_configuration = AlbertConfig(**config['model_params'])\n",
    "    \n",
    "    bert = AlbertModel(albert_base_configuration)\n",
    "    bert = MultiTaskModel(bert, \n",
    "                          num_vocab=1 + max([m['token'] for m in token_maps.values()]), \n",
    "                          num_tokens=config['model_params']['vocab_size'],\n",
    "                          hidden_size=config['model_params']['hidden_size'])\n",
    "    \n",
    "    load = True\n",
    "    try:\n",
    "        files = os.listdir(log_dir)\n",
    "        ckpts = []\n",
    "        for f in os.listdir(log_dir):\n",
    "            if f.startswith(\"step_\"): ckpts.append(f)\n",
    "\n",
    "        iters = [int(f.split('_')[-1].split('.')[0]) for f in ckpts if os.path.isfile(os.path.join(log_dir, f))]\n",
    "        iters = sorted(iters)[-1]\n",
    "    except:\n",
    "        iters = 0\n",
    "        load = False\n",
    "    \n",
    "    optimizer = AdamW(bert.parameters(), lr=1e-4)\n",
    "    \n",
    "    accelerator = Accelerator(mixed_precision=config['mixed_precision'], split_batches=True, kwargs_handlers=[ddp_kwargs])\n",
    "    \n",
    "    if load:\n",
    "        checkpoint = torch.load(log_dir + \"/step_\" + str(iters) + \".t7\", map_location='cpu')\n",
    "        state_dict = checkpoint['net']\n",
    "        from collections import OrderedDict\n",
    "        new_state_dict = OrderedDict()\n",
    "        for k, v in state_dict.items():\n",
    "            name = k[7:] # remove `module.`\n",
    "            new_state_dict[name] = v\n",
    "\n",
    "        bert.load_state_dict(new_state_dict, strict=False)\n",
    "        \n",
    "        accelerator.print('Checkpoint loaded.')\n",
    "        optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "    \n",
    "    bert, optimizer, train_loader = accelerator.prepare(\n",
    "        bert, optimizer, train_loader\n",
    "    )\n",
    "\n",
    "    accelerator.print('Start training...')\n",
    "\n",
    "    running_loss = 0\n",
    "    \n",
    "    for _, batch in enumerate(train_loader):        \n",
    "        curr_steps += 1\n",
    "        \n",
    "        words, labels, phonemes, input_lengths, masked_indices = batch\n",
    "        text_mask = length_to_mask(torch.Tensor(input_lengths))# .to(device)\n",
    "        \n",
    "        tokens_pred, words_pred = bert(phonemes, attention_mask=(~text_mask).int())\n",
    "        \n",
    "        loss_vocab = 0\n",
    "        for _s2s_pred, _text_input, _text_length, _masked_indices in zip(words_pred, words, input_lengths, masked_indices):\n",
    "            loss_vocab += criterion(_s2s_pred[:_text_length], \n",
    "                                        _text_input[:_text_length])\n",
    "        loss_vocab /= words.size(0)\n",
    "        \n",
    "        loss_token = 0\n",
    "        sizes = 1\n",
    "        for _s2s_pred, _text_input, _text_length, _masked_indices in zip(tokens_pred, labels, input_lengths, masked_indices):\n",
    "            if len(_masked_indices) > 0:\n",
    "                _text_input = _text_input[:_text_length][_masked_indices]\n",
    "                loss_tmp = criterion(_s2s_pred[:_text_length][_masked_indices], \n",
    "                                            _text_input[:_text_length]) \n",
    "                loss_token += loss_tmp\n",
    "                sizes += 1\n",
    "        loss_token /= sizes\n",
    "\n",
    "        loss = loss_vocab + loss_token\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        accelerator.backward(loss)\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        iters = iters + 1\n",
    "        if (iters+1)%log_interval == 0:\n",
    "            accelerator.print ('Step [%d/%d], Loss: %.5f, Vocab Loss: %.5f, Token Loss: %.5f'\n",
    "                    %(iters+1, num_steps, running_loss / log_interval, loss_vocab, loss_token))\n",
    "            running_loss = 0\n",
    "            \n",
    "        if (iters+1)%save_interval == 0:\n",
    "            accelerator.print('Saving..')\n",
    "\n",
    "            state = {\n",
    "                'net':  bert.state_dict(),\n",
    "                'step': iters,\n",
    "                'optimizer': optimizer.state_dict(),\n",
    "            }\n",
    "\n",
    "            accelerator.save(state, log_dir + '/step_' + str(iters + 1) + '.t7')\n",
    "\n",
    "        if curr_steps > num_steps:\n",
    "            return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c4b8958",
   "metadata": {},
   "outputs": [],
   "source": [
    "from accelerate import notebook_launcher\n",
    "while True:\n",
    "    notebook_launcher(train, args=(), num_processes=3, use_port=33389)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dcf4988",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "BERT",
   "language": "python",
   "name": "bert"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}