Train_exp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import random\n",
    "import shutil\n",
    "import logging\n",
    "import argparse\n",
    "import subprocess\n",
    "from time import time\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from test import test\n",
    "from lib.config import Config\n",
    "from utils.evaluator import Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, exp_dir, cfg, val_loader, train_state=None):\n",
    "    # Get initial train state\n",
    "    optimizer = cfg.get_optimizer(model.parameters())\n",
    "    scheduler = cfg.get_lr_scheduler(optimizer)\n",
    "    starting_epoch = 1\n",
    "\n",
    "    if train_state is not None:\n",
    "        model.load_state_dict(train_state['model'])\n",
    "        optimizer.load_state_dict(train_state['optimizer'])\n",
    "        scheduler.load_state_dict(train_state['lr_scheduler'])\n",
    "        starting_epoch = train_state['epoch'] + 1\n",
    "        scheduler.step(starting_epoch)\n",
    "\n",
    "    # Train the model\n",
    "    criterion_parameters = cfg.get_loss_parameters()\n",
    "    criterion = model.loss\n",
    "    total_step = len(train_loader)\n",
    "    ITER_LOG_INTERVAL = cfg['iter_log_interval']\n",
    "    ITER_TIME_WINDOW = cfg['iter_time_window']\n",
    "    MODEL_SAVE_INTERVAL = cfg['model_save_interval']\n",
    "    t0 = time()\n",
    "    total_iter = 0\n",
    "    iter_times = []\n",
    "    logging.info(\"Starting training.\")\n",
    "    for epoch in range(starting_epoch, num_epochs + 1):\n",
    "        epoch_t0 = time()\n",
    "        logging.info(\"Beginning epoch {}\".format(epoch))\n",
    "        accum_loss = 0\n",
    "        for i, (images, labels, img_idxs) in enumerate(train_loader):\n",
    "            total_iter += 1\n",
    "            iter_t0 = time()\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(images, epoch=epoch)\n",
    "            loss, loss_dict_i = criterion(outputs, labels, **criterion_parameters)\n",
    "            accum_loss += loss.item()\n",
    "\n",
    "            # Backward and optimize\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            iter_times.append(time() - iter_t0)\n",
    "            if len(iter_times) > 100:\n",
    "                iter_times = iter_times[-ITER_TIME_WINDOW:]\n",
    "            if (i + 1) % ITER_LOG_INTERVAL == 0:\n",
    "                loss_str = ', '.join(\n",
    "                    ['{}: {:.4f}'.format(loss_name, loss_dict_i[loss_name]) for loss_name in loss_dict_i])\n",
    "                logging.info(\"Epoch [{}/{}], Step [{}/{}], Loss: {:.4f} ({}), s/iter: {:.4f}, lr: {:.1e}\".format(\n",
    "                    epoch,\n",
    "                    num_epochs,\n",
    "                    i + 1,\n",
    "                    total_step,\n",
    "                    accum_loss / (i + 1),\n",
    "                    loss_str,\n",
    "                    np.mean(iter_times),\n",
    "                    optimizer.param_groups[0][\"lr\"],\n",
    "                ))\n",
    "        logging.info(\"Epoch time: {:.4f}\".format(time() - epoch_t0))\n",
    "        if epoch % MODEL_SAVE_INTERVAL == 0 or epoch == num_epochs:\n",
    "            model_path = os.path.join(exp_dir, \"models\", \"model_{:03d}.pt\".format(epoch))\n",
    "            save_train_state(model_path, model, optimizer, scheduler, epoch)\n",
    "        if val_loader is not None:\n",
    "            evaluator = Evaluator(val_loader.dataset, exp_root)\n",
    "            evaluator, val_loss = test(\n",
    "                model,\n",
    "                val_loader,\n",
    "                evaluator,\n",
    "                None,\n",
    "                cfg,\n",
    "                view=False,\n",
    "                epoch=-1,\n",
    "                verbose=False,\n",
    "            )\n",
    "            _, results = evaluator.eval(label=None, only_metrics=True)\n",
    "            logging.info(\"Epoch [{}/{}], Val loss: {:.4f}\".format(epoch, num_epochs, val_loss))\n",
    "            model.train()\n",
    "        scheduler.step()\n",
    "    logging.info(\"Training time: {:.4f}\".format(time() - t0))\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_train_state(path, model, optimizer, lr_scheduler, epoch):\n",
    "    train_state = {\n",
    "        'model': model.state_dict(),\n",
    "        'optimizer': optimizer.state_dict(),\n",
    "        'lr_scheduler': lr_scheduler.state_dict(),\n",
    "        'epoch': epoch\n",
    "    }\n",
    "\n",
    "    torch.save(train_state, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description=\"Train PolyLaneNet\")\n",
    "    parser.add_argument(\"--exp_name\", default=\"default\", help=\"Experiment name\", required=True)\n",
    "    parser.add_argument(\"--cfg\", default=\"config.yaml\", help=\"Config file\", required=True)\n",
    "    parser.add_argument(\"--resume\", action=\"store_true\", help=\"Resume training\")\n",
    "    parser.add_argument(\"--validate\", action=\"store_true\", help=\"Validate model during training\")\n",
    "    parser.add_argument(\"--deterministic\",\n",
    "                        action=\"store_true\",\n",
    "                        help=\"set cudnn.deterministic = True and cudnn.benchmark = False\")\n",
    "\n",
    "    return parser.parse_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_code_state():\n",
    "    state = \"Git hash: {}\".format(\n",
    "        subprocess.run(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE).stdout.decode('utf-8'))\n",
    "    state += '\\n*************\\nGit diff:\\n*************\\n'\n",
    "    state += subprocess.run(['git', 'diff'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
    "\n",
    "    return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_exp_dir(exps_dir, exp_name, cfg_path):\n",
    "    dirs = [\"models\"]\n",
    "    exp_root = os.path.join(exps_dir, exp_name)\n",
    "\n",
    "    for dirname in dirs:\n",
    "        os.makedirs(os.path.join(exp_root, dirname), exist_ok=True)\n",
    "\n",
    "    shutil.copyfile(cfg_path, os.path.join(exp_root, 'config.yaml'))\n",
    "    with open(os.path.join(exp_root, 'code_state.txt'), 'w') as file:\n",
    "        file.write(get_code_state())\n",
    "\n",
    "    return exp_root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_exp_train_state(exp_root):\n",
    "    models_dir = os.path.join(exp_root, \"models\")\n",
    "    models = os.listdir(models_dir)\n",
    "    last_epoch, last_modelname = sorted(\n",
    "        [(int(name.split(\"_\")[1].split(\".\")[0]), name) for name in models],\n",
    "        key=lambda x: x[0],\n",
    "    )[-1]\n",
    "    train_state = torch.load(os.path.join(models_dir, last_modelname))\n",
    "\n",
    "    return train_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_on_exception(exc_type, exc_value, exc_traceback):\n",
    "    logging.error(\"Uncaught exception\", exc_info=(exc_type, exc_value, exc_traceback))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description=\"Train PolyLaneNet\")\n",
    "parser.add_argument(\"--exp_name\", default=\"default\", help=\"Experiment name\", required=True)\n",
    "parser.add_argument(\"--cfg\", default=\"config.yaml\", help=\"Config file\", required=True)\n",
    "parser.add_argument(\"--resume\", action=\"store_true\", help=\"Resume training\")\n",
    "parser.add_argument(\"--validate\", action=\"store_true\", help=\"Validate model during training\")\n",
    "parser.add_argument(\"--deterministic\",\n",
    "                        action=\"store_true\",\n",
    "                        help=\"set cudnn.deterministic = True and cudnn.benchmark = False\")\n",
    "args = parser.parse_args(args=['--exp_name', 'tusimple', '--cfg', './cfgs/tusimple.yaml'])\n",
    "cfg = Config(args.cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up seeds\n",
    "torch.manual_seed(cfg['seed'])\n",
    "np.random.seed(cfg['seed'])\n",
    "random.seed(cfg['seed'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.deterministic:\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up experiment\n",
    "if not args.resume:\n",
    "    exp_root = setup_exp_dir(cfg['exps_dir'], args.exp_name, args.cfg)\n",
    "else:\n",
    "    exp_root = os.path.join(cfg['exps_dir'], os.path.basename(os.path.normpath(args.exp_name)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.basicConfig(\n",
    "        format=\"[%(asctime)s] [%(levelname)s] %(message)s\",\n",
    "        level=logging.INFO,\n",
    "        handlers=[\n",
    "            logging.FileHandler(os.path.join(exp_root, \"log.txt\")),\n",
    "            logging.StreamHandler(),\n",
    "        ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.excepthook = log_on_exception"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2020-06-13 11:36:12,082] [INFO] Experiment name: tusimple\n",
      "[2020-06-13 11:36:12,084] [INFO] Config:\n",
      "# Training settings\n",
      "exps_dir: 'experiments'\n",
      "iter_log_interval: 1\n",
      "iter_time_window: 100\n",
      "model_save_interval: 1\n",
      "seed: 1\n",
      "backup:\n",
      "model:\n",
      "  name: PolyRegression\n",
      "  parameters:\n",
      "    num_outputs: 35 # (5 lanes) * (1 conf + 2 (upper & lower) + 4 poly coeffs)\n",
      "    pretrained: false\n",
      "    backbone: 'efficientnet-b0'\n",
      "    pred_category: false\n",
      "    curriculum_steps: [0, 0, 0, 0]\n",
      "loss_parameters:\n",
      "  conf_weight: 1\n",
      "  lower_weight: 1\n",
      "  upper_weight: 1\n",
      "  cls_weight: 0\n",
      "  poly_weight: 300\n",
      "batch_size: 16\n",
      "epochs: 2695\n",
      "optimizer:\n",
      "  name: Adam\n",
      "  parameters:\n",
      "    lr: 3.0e-4\n",
      "lr_scheduler:\n",
      "  name: CosineAnnealingLR\n",
      "  parameters:\n",
      "    T_max: 385\n",
      "\n",
      "# Testing settings\n",
      "test_parameters:\n",
      "  conf_threshold: 0.5\n",
      "\n",
      "# Dataset settings\n",
      "datasets:\n",
      "  train:\n",
      "    type: LaneDataset\n",
      "    parameters:\n",
      "      dataset: tusimple\n",
      "      split: train\n",
      "      img_size: [360, 640]\n",
      "      normalize: true\n",
      "      aug_chance: 0.9090909090909091 # 10/11\n",
      "      augmentations:\n",
      "       - name: Affine\n",
      "         parameters:\n",
      "           rotate: !!python/tuple [-10, 10]\n",
      "       - name: HorizontalFlip\n",
      "         parameters:\n",
      "           p: 0.5\n",
      "       - name: CropToFixedSize\n",
      "         parameters:\n",
      "           width: 1152\n",
      "           height: 648\n",
      "      root: \"/home/ubuntu/data/tusimple\"\n",
      "\n",
      "  test: &test\n",
      "    type: LaneDataset\n",
      "    parameters:\n",
      "      dataset: tusimple\n",
      "      split: val\n",
      "      max_lanes: 5\n",
      "      img_size: [360, 640]\n",
      "      root: \"/home/ubuntu/data/tusimple\"\n",
      "      normalize: true\n",
      "      augmentations: []\n",
      "\n",
      "  # val = test\n",
      "  val:\n",
      "    <<: *test\n",
      "\n",
      "[2020-06-13 11:36:12,085] [INFO] Args:\n",
      "Namespace(cfg='./cfgs/tusimple.yaml', deterministic=False, exp_name='tusimple', resume=False, validate=False)\n"
     ]
    }
   ],
   "source": [
    "logging.info(\"Experiment name: {}\".format(args.exp_name))\n",
    "logging.info(\"Config:\\n\" + str(cfg))\n",
    "logging.info(\"Args:\\n\" + str(args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total annos 3268\n",
      "Transforming annotations...\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "# Get data sets\n",
    "train_dataset = cfg.get_dataset(\"train\")\n",
    "\n",
    "#この行の理解が重要そう"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Device configuration\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyper parameters\n",
    "num_epochs = cfg[\"epochs\"]\n",
    "batch_size = cfg[\"batch_size\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model\n",
    "model = cfg.get_model().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_state = None\n",
    "if args.resume:\n",
    "    train_state = get_exp_train_state(exp_root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data loader\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_dataset,\n",
    "                                               batch_size=batch_size,\n",
    "                                               shuffle=True,\n",
    "                                               num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.validate:\n",
    "    val_dataset = cfg.get_dataset(\"val\")\n",
    "    val_loader = torch.utils.data.DataLoader(dataset=val_dataset,\n",
    "                                                 batch_size=batch_size,\n",
    "                                                 shuffle=False,\n",
    "                                                 num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train regressor\n",
    "try:\n",
    "    model = train(\n",
    "            model,\n",
    "            train_loader,\n",
    "            exp_root,\n",
    "            cfg,\n",
    "            val_loader=val_loader if args.validate else None,\n",
    "            train_state=train_state,\n",
    "        )\n",
    "except KeyboardInterrupt:\n",
    "    logging.info(\"Training session terminated.\")\n",
    "test_epoch = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if cfg['backup'] is not None:\n",
    "    subprocess.run(['rclone', 'copy', exp_root, '{}/{}'.format(cfg['backup'], args.exp_name)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Eval model after training\n",
    "test_dataset = cfg.get_dataset(\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset,\n",
    "                                              batch_size=batch_size,\n",
    "                                              shuffle=False,\n",
    "                                              num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluator = Evaluator(test_loader.dataset, exp_root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.basicConfig(\n",
    "        format=\"[%(asctime)s] [%(levelname)s] %(message)s\",\n",
    "        level=logging.INFO,\n",
    "        handlers=[\n",
    "            logging.FileHandler(os.path.join(exp_root, \"test_log.txt\")),\n",
    "            logging.StreamHandler(),\n",
    "        ],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.info('Code state:\\n {}'.format(get_code_state()))\n",
    "_, mean_loss = test(model, test_loader, evaluator, exp_root, cfg, epoch=test_epoch, view=False)\n",
    "logging.info(\"Mean test loss: {:.4f}\".format(mean_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluator.exp_name = args.exp_name\n",
    "\n",
    "eval_str, _ = evaluator.eval(label='{}_{}'.format(os.path.basename(args.exp_name), test_epoch))\n",
    "\n",
    "logging.info(eval_str)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_pytorch_p36)",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}