train_nb.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# %load train.py\n",
    "#!/usr/bin/env python\n",
    "\n",
    "\n",
    "from __future__ import division\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import cuda\n",
    "\n",
    "import onmt\n",
    "import onmt.Models\n",
    "import onmt.ModelConstructor\n",
    "import onmt.modules\n",
    "from onmt.Utils import aeq, use_gpu\n",
    "import opts\n",
    "import sys\n",
    "\n",
    "from comet_ml import Experiment\n",
    "\n",
    "\n",
    "\n",
    "def report_func(epoch, batch, num_batches,\n",
    "                start_time, lr, report_stats):\n",
    "    \"\"\"\n",
    "    This is the user-defined batch-level traing progress\n",
    "    report function.\n",
    "\n",
    "    Args:\n",
    "        epoch(int): current epoch count.\n",
    "        batch(int): current batch count.\n",
    "        num_batches(int): total number of batches.\n",
    "        start_time(float): last report time.\n",
    "        lr(float): current learning rate.\n",
    "        report_stats(Statistics): old Statistics instance.\n",
    "    Returns:\n",
    "        report_stats(Statistics): updated Statistics instance.\n",
    "    \"\"\"\n",
    "    if batch % opt.report_every == -1 % opt.report_every:\n",
    "        report_stats.output(epoch, batch+1, num_batches, start_time)\n",
    "        if opt.exp_host:\n",
    "            report_stats.log(\"progress\", experiment, lr)\n",
    "        report_stats = onmt.Statistics()\n",
    "\n",
    "    return report_stats\n",
    "\n",
    "\n",
    "def make_train_data_iter(train_data, opt):\n",
    "    \"\"\"\n",
    "    This returns user-defined train data iterator for the trainer\n",
    "    to iterate over during each train epoch. We implement simple\n",
    "    ordered iterator strategy here, but more sophisticated strategy\n",
    "    like curriculum learning is ok too.\n",
    "    \"\"\"\n",
    "    return onmt.IO.OrderedIterator(\n",
    "                dataset=train_data, batch_size=opt.batch_size,\n",
    "                device=opt.gpuid[0] if opt.gpuid else -1,\n",
    "                repeat=False)\n",
    "\n",
    "\n",
    "def make_valid_data_iter(valid_data, opt):\n",
    "    \"\"\"\n",
    "    This returns user-defined validate data iterator for the trainer\n",
    "    to iterate over during each validate epoch. We implement simple\n",
    "    ordered iterator strategy here, but more sophisticated strategy\n",
    "    is ok too.\n",
    "    \"\"\"\n",
    "    return onmt.IO.OrderedIterator(\n",
    "                dataset=valid_data, batch_size=opt.batch_size,\n",
    "                device=opt.gpuid[0] if opt.gpuid else -1,\n",
    "                train=False, sort=True)\n",
    "\n",
    "\n",
    "def make_loss_compute(model, tgt_vocab, dataset, opt):\n",
    "    \"\"\"\n",
    "    This returns user-defined LossCompute object, which is used to\n",
    "    compute loss in train/validate process. You can implement your\n",
    "    own *LossCompute class, by subclassing LossComputeBase.\n",
    "    \"\"\"\n",
    "    if opt.copy_attn:\n",
    "        compute = onmt.modules.CopyGeneratorLossCompute(\n",
    "            model.generator, tgt_vocab, dataset, opt.copy_attn_force)\n",
    "    else:\n",
    "        compute = onmt.Loss.NMTLossCompute(model.generator, tgt_vocab)\n",
    "\n",
    "    if use_gpu(opt):\n",
    "        compute.cuda()\n",
    "\n",
    "    return compute\n",
    "\n",
    "\n",
    "def train_model(model, train_data, valid_data, fields, optim):\n",
    "\n",
    "    train_iter = make_train_data_iter(train_data, opt)\n",
    "    valid_iter = make_valid_data_iter(valid_data, opt)\n",
    "\n",
    "    train_loss = make_loss_compute(model, fields[\"tgt\"].vocab,\n",
    "                                   train_data, opt)\n",
    "    valid_loss = make_loss_compute(model, fields[\"tgt\"].vocab,\n",
    "                                   valid_data, opt)\n",
    "    \n",
    "\n",
    "    trunc_size = opt.truncated_decoder  # Badly named...\n",
    "    shard_size = opt.max_generator_batches\n",
    "\n",
    "    trainer = onmt.Trainer(model, train_iter, valid_iter,\n",
    "                           train_loss, valid_loss, optim,\n",
    "                           trunc_size, shard_size)\n",
    "\n",
    "    for epoch in range(opt.start_epoch, opt.epochs + 1):\n",
    "        print('')\n",
    "\n",
    "        # 1. Train for one epoch on the training set.\n",
    "        train_stats = trainer.train(epoch, report_func)\n",
    "        print('Train perplexity: %g' % train_stats.ppl())\n",
    "        print('Train accuracy: %g' % train_stats.accuracy())\n",
    "        if comet:\n",
    "            experiment_comet.log_metric(\"Train acc: {}\".format(train_stats.accuracy()))\n",
    "\n",
    "        # 2. Validate on the validation set.\n",
    "        valid_stats = trainer.validate()\n",
    "        print('Validation perplexity: %g' % valid_stats.ppl())\n",
    "        print('Validation accuracy: %g' % valid_stats.accuracy())\n",
    "        if comet:\n",
    "            experiment_comet.log-metric(\"Test acc: {}\".format(valid_stats.accuracy()))\n",
    "\n",
    "        # 3. Log to remote server.\n",
    "        if opt.exp_host:\n",
    "            train_stats.log(\"train\", experiment, optim.lr)\n",
    "            valid_stats.log(\"valid\", experiment, optim.lr)\n",
    "\n",
    "        # 4. Update the learning rate\n",
    "        trainer.epoch_step(valid_stats.ppl(), epoch)\n",
    "\n",
    "        # 5. Drop a checkpoint if needed.\n",
    "        if epoch >= opt.start_checkpoint_at:\n",
    "            trainer.drop_checkpoint(opt, epoch, fields, valid_stats)\n",
    "\n",
    "\n",
    "def check_save_model_path():\n",
    "    save_model_path = os.path.abspath(opt.save_model)\n",
    "    model_dirname = os.path.dirname(save_model_path)\n",
    "    if not os.path.exists(model_dirname):\n",
    "        os.makedirs(model_dirname)\n",
    "\n",
    "\n",
    "def tally_parameters(model):\n",
    "    n_params = sum([p.nelement() for p in model.parameters()])\n",
    "    print('* number of parameters: %d' % n_params)\n",
    "    enc = 0\n",
    "    dec = 0\n",
    "    for name, param in model.named_parameters():\n",
    "        if 'encoder' in name:\n",
    "            enc += param.nelement()\n",
    "        elif 'decoder' or 'generator' in name:\n",
    "            dec += param.nelement()\n",
    "    print('encoder: ', enc)\n",
    "    print('decoder: ', dec)\n",
    "\n",
    "\n",
    "def load_fields(train, valid, checkpoint):\n",
    "    fields = onmt.IO.load_fields(\n",
    "                torch.load(opt.data + '.vocab.pt'))\n",
    "    fields = dict([(k, f) for (k, f) in fields.items()\n",
    "                  if k in train.examples[0].__dict__])\n",
    "    train.fields = fields\n",
    "    valid.fields = fields\n",
    "\n",
    "    if opt.train_from:\n",
    "        print('Loading vocab from checkpoint at %s.' % opt.train_from)\n",
    "        fields = onmt.IO.load_fields(checkpoint['vocab'])\n",
    "\n",
    "    print(' * vocabulary size. source = %d; target = %d' %\n",
    "          (len(fields['src'].vocab), len(fields['tgt'].vocab)))\n",
    "\n",
    "    return fields\n",
    "\n",
    "\n",
    "def collect_features(train, fields):\n",
    "    # TODO: account for target features.\n",
    "    # Also, why does fields need to have the structure it does?\n",
    "    src_features = onmt.IO.collect_features(fields)\n",
    "    aeq(len(src_features), train.n_src_feats)\n",
    "\n",
    "    return src_features\n",
    "\n",
    "\n",
    "def build_model(model_opt, opt, fields, checkpoint):\n",
    "    print('Building model...')\n",
    "    model = onmt.ModelConstructor.make_base_model(model_opt, fields,\n",
    "                                                  use_gpu(opt), checkpoint)\n",
    "    if len(opt.gpuid) > 1:\n",
    "        print('Multi gpu training: ', opt.gpuid)\n",
    "        model = nn.DataParallel(model, device_ids=opt.gpuid, dim=1)\n",
    "    print(model)\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def build_optim(model, checkpoint):\n",
    "    if opt.train_from:\n",
    "        print('Loading optimizer from checkpoint.')\n",
    "        optim = checkpoint['optim']\n",
    "        optim.optimizer.load_state_dict(\n",
    "            checkpoint['optim'].optimizer.state_dict())\n",
    "    else:\n",
    "        # what members of opt does Optim need?\n",
    "        optim = onmt.Optim(\n",
    "            opt.optim, opt.learning_rate, opt.max_grad_norm,\n",
    "            lr_decay=opt.learning_rate_decay,\n",
    "            start_decay_at=opt.start_decay_at,\n",
    "            opt=opt\n",
    "        )\n",
    "\n",
    "    optim.set_parameters(model.parameters())\n",
    "\n",
    "    return optim\n",
    "\n",
    "\n",
    "def launch():\n",
    "    # Load train and validate data.\n",
    "    print(\"Loading train and validate data from '%s'\" % opt.data)\n",
    "    train = torch.load(opt.data + '.train.pt')\n",
    "    valid = torch.load(opt.data + '.valid.pt')\n",
    "    print(' * number of training sentences: %d' % len(train))\n",
    "    print(' * maximum batch size: %d' % opt.batch_size)\n",
    "\n",
    "    # Load checkpoint if we resume from a previous training.\n",
    "    if opt.train_from:\n",
    "        print('Loading checkpoint from %s' % opt.train_from)\n",
    "        checkpoint = torch.load(opt.train_from,\n",
    "                                map_location=lambda storage, loc: storage)\n",
    "        model_opt = checkpoint['opt']\n",
    "        # I don't like reassigning attributes of opt: it's not clear\n",
    "        opt.start_epoch = checkpoint['epoch'] + 1\n",
    "    else:\n",
    "        checkpoint = None\n",
    "        model_opt = opt\n",
    "\n",
    "    # Load fields generated from preprocess phase.\n",
    "    fields = load_fields(train, valid, checkpoint)\n",
    "\n",
    "    # Collect features.\n",
    "    src_features = collect_features(train, fields)\n",
    "    for j, feat in enumerate(src_features):\n",
    "        print(' * src feature %d size = %d' % (j, len(fields[feat].vocab)))\n",
    "\n",
    "    # Build model.\n",
    "    model = build_model(model_opt, opt, fields, checkpoint)\n",
    "    tally_parameters(model)\n",
    "    check_save_model_path()\n",
    "\n",
    "    # Build optimizer.\n",
    "    optim = build_optim(model, checkpoint)\n",
    "\n",
    "    # Do training.\n",
    "    train_model(model, train, valid, fields, optim)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/share/apps/python3/3.5.3/intel/lib/python3.5/site-packages/ipykernel/__main__.py', '--data', '../data/sumdata/train/textsum']\n",
      "['/share/apps/python3/3.5.3/intel/lib/python3.5/site-packages/ipykernel/__main__.py', '-data', '../data/sumdata/train/textsum']\n"
     ]
    }
   ],
   "source": [
    "comet = False\n",
    "\n",
    "if comet:\n",
    "    experiment_comet = Experiment(api_key=os.environ.get(\"COMET_KEY\"))\n",
    "\n",
    "\n",
    "parser = argparse.ArgumentParser(\n",
    "    description='train.py',\n",
    "    formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
    "\n",
    "# opts.py\n",
    "opts.add_md_help_argument(parser)\n",
    "opts.model_opts(parser)\n",
    "opts.train_opts(parser)\n",
    "print(sys.argv)\n",
    "\n",
    "# Ugly black magic hack to have it work on notebook\n",
    "sys.argv = sys.argv[:1] + [\"-data\", \"../data/sumdata/train/textsum\"]\n",
    "print(sys.argv)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "opt = parser.parse_args()\n",
    "if opt.word_vec_size != -1:\n",
    "    opt.src_word_vec_size = opt.word_vec_size\n",
    "    opt.tgt_word_vec_size = opt.word_vec_size\n",
    "\n",
    "if opt.layers != -1:\n",
    "    opt.enc_layers = opt.layers\n",
    "    opt.dec_layers = opt.layers\n",
    "\n",
    "opt.brnn = (opt.encoder_type == \"brnn\")\n",
    "if opt.seed > 0:\n",
    "    torch.manual_seed(opt.seed)\n",
    "\n",
    "if opt.rnn_type == \"SRU\" and not opt.gpuid:\n",
    "    raise AssertionError(\"Using SRU requires -gpuid set.\")\n",
    "\n",
    "if torch.cuda.is_available() and not opt.gpuid:\n",
    "    print(\"WARNING: You have a CUDA device, should run with -gpuid 0\")\n",
    "\n",
    "if opt.gpuid:\n",
    "    cuda.set_device(opt.gpuid[0])\n",
    "    if opt.seed > 0:\n",
    "        torch.cuda.manual_seed(opt.seed)\n",
    "\n",
    "if len(opt.gpuid) > 1:\n",
    "    sys.stderr.write(\"Sorry, multigpu isn't supported yet, coming soon!\\n\")\n",
    "    sys.exit(1)\n",
    "\n",
    "\n",
    "# Set up the Crayon logging server.\n",
    "if opt.exp_host != \"\":\n",
    "    from pycrayon import CrayonClient\n",
    "    cc = CrayonClient(hostname=opt.exp_host)\n",
    "\n",
    "    experiments = cc.get_experiment_names()\n",
    "    print(experiments)\n",
    "    if opt.exp in experiments:\n",
    "        cc.remove_experiment(opt.exp)\n",
    "    experiment = cc.create_experiment(opt.exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading train and validate data from '../data/sumdata/train/textsum'\n",
      " * number of training sentences: 3736183\n",
      " * maximum batch size: 64\n",
      " * vocabulary size. source = 50002; target = 50004\n",
      "Building model...\n",
      "Intializing model parameters.\n",
      "NMTModel (\n",
      "  (encoder): RNNEncoder (\n",
      "    (embeddings): Embeddings (\n",
      "      (make_embedding): Sequential (\n",
      "        (emb_luts): Elementwise (\n",
      "          (0): Embedding(50002, 500, padding_idx=1)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (rnn): LSTM(500, 500, num_layers=2, dropout=0.3)\n",
      "  )\n",
      "  (decoder): InputFeedRNNDecoder (\n",
      "    (embeddings): Embeddings (\n",
      "      (make_embedding): Sequential (\n",
      "        (emb_luts): Elementwise (\n",
      "          (0): Embedding(50004, 500, padding_idx=1)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (dropout): Dropout (p = 0.3)\n",
      "    (rnn): StackedLSTM (\n",
      "      (dropout): Dropout (p = 0.3)\n",
      "      (layers): ModuleList (\n",
      "        (0): LSTMCell(1000, 500)\n",
      "        (1): LSTMCell(500, 500)\n",
      "      )\n",
      "    )\n",
      "    (attn): GlobalAttention (\n",
      "      (linear_in): Linear (500 -> 500)\n",
      "      (linear_out): Linear (1000 -> 500)\n",
      "      (sm): Softmax ()\n",
      "      (tanh): Tanh ()\n",
      "    )\n",
      "  )\n",
      "  (generator): Sequential (\n",
      "    (0): Linear (500 -> 50004)\n",
      "    (1): LogSoftmax ()\n",
      "  )\n",
      ")\n",
      "* number of parameters: 84821004\n",
      "encoder:  29009000\n",
      "decoder:  55812004\n",
      "\n",
      "Epoch  1,    50/58378; acc:   6.26; ppl: 443268.91; 581 src tok/s; 177 tgt tok/s;    167 s elapsed\n"
     ]
    }
   ],
   "source": [
    "launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}