trainX.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from utils.datasets import LoadDataset\n",
    "from utils import torch_utils\n",
    "import time\n",
    "import rtest\n",
    "from tqdm import tnrange, tqdm_notebook, tqdm\n",
    "from models import *\n",
    "import os\n",
    "from tqdm import tnrange, tqdm_notebook\n",
    "import tensorboardX\n",
    "from prefetch_generator import BackgroundGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using CUDA device0 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', total_memory=16280MB)\n",
      "           device1 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', total_memory=16280MB)\n",
      "           device2 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', total_memory=16280MB)\n",
      "           device3 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', total_memory=16280MB)\n"
     ]
    }
   ],
   "source": [
    "# specify visible GPU\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '2,3,4,5'\n",
    "device = torch_utils.select_device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.backends.cudnn.benchmark = True\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "torch.cuda.manual_seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_SIZE = 416\n",
    "EPOCHES = 240\n",
    "BATCH_SIZE = 64\n",
    "START_EPOCH = 0\n",
    "CFG = 'cfg/yolov3-tiny.cfg'\n",
    "DATA_CFG = 'cfg/coco.data'\n",
    "NUM_WORKERS = 30\n",
    "FREEZE_BACKBONE = False\n",
    "FROM_SCRATCH = True\n",
    "NAME = 'train_yolo-tiny_{}'.format(int(time.time()))\n",
    "\n",
    "LOG_PATH = 'logs/'\n",
    "weights = 'weights/'\n",
    "latest = os.path.join(weights, 'latest.pt')\n",
    "best = os.path.join(weights, 'best.pt')\n",
    "TBoard = tensorboardX.SummaryWriter(log_dir=os.path.join(LOG_PATH, NAME))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/data/hktxt/e/CV/coco/trainvalno5k.txt'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_path = parse_data_cfg(DATA_CFG)['train'];train_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset\n",
    "dataset = LoadDataset(train_path, img_size=IMG_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model\n",
    "model = Darknet(CFG, IMG_SIZE).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataloader\n",
    "dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, shuffle=False, pin_memory=True, collate_fn=dataset.collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1833"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check Dataloader\n",
    "for i, (imgs, targets, _, _) in enumerate(dataloader):\n",
    "    #print(targets.shape)\n",
    "    plot_images(imgs=imgs, targets=targets, fname='train_batch0.jpg')\n",
    "    if i == 0:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "cutoff = -1  # backbone reaches to cutoff layer\n",
    "START_EPOCH = 0\n",
    "best_loss = float('inf')\n",
    "\n",
    "if not FROM_SCRATCH:\n",
    "    if '-tiny.cfg' in CFG:\n",
    "        cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')\n",
    "    else:\n",
    "        cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Let's use 4 GPUs!\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.device_count() > 1:\n",
    "    print(\"Let's use\", torch.cuda.device_count(), \"GPUs!\")\n",
    "    model = nn.DataParallel(model)\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "layer                                     name  gradient   parameters                shape         mu      sigma\n",
      "    0                   module.0.conv_0.weight      True          432        [16, 3, 3, 3]    0.00273      0.109\n",
      "    1             module.0.batch_norm_0.weight      True           16                 [16]      0.569      0.303\n",
      "    2               module.0.batch_norm_0.bias      True           16                 [16]          0          0\n",
      "    3                   module.2.conv_2.weight      True         4608       [32, 16, 3, 3]   -0.00111     0.0481\n",
      "    4             module.2.batch_norm_2.weight      True           32                 [32]      0.549      0.305\n",
      "    5               module.2.batch_norm_2.bias      True           32                 [32]          0          0\n",
      "    6                   module.4.conv_4.weight      True        18432       [64, 32, 3, 3]   4.86e-05     0.0341\n",
      "    7             module.4.batch_norm_4.weight      True           64                 [64]      0.505      0.283\n",
      "    8               module.4.batch_norm_4.bias      True           64                 [64]          0          0\n",
      "    9                   module.6.conv_6.weight      True        73728      [128, 64, 3, 3]    6.5e-05      0.024\n",
      "   10             module.6.batch_norm_6.weight      True          128                [128]      0.484      0.289\n",
      "   11               module.6.batch_norm_6.bias      True          128                [128]          0          0\n",
      "   12                   module.8.conv_8.weight      True       294912     [256, 128, 3, 3]  -4.65e-06      0.017\n",
      "   13             module.8.batch_norm_8.weight      True          256                [256]      0.526      0.289\n",
      "   14               module.8.batch_norm_8.bias      True          256                [256]          0          0\n",
      "   15                 module.10.conv_10.weight      True  1.17965e+06     [512, 256, 3, 3]  -1.09e-05      0.012\n",
      "   16           module.10.batch_norm_10.weight      True          512                [512]      0.503      0.296\n",
      "   17             module.10.batch_norm_10.bias      True          512                [512]          0          0\n",
      "   18                 module.12.conv_12.weight      True  4.71859e+06    [1024, 512, 3, 3]   5.19e-06    0.00851\n",
      "   19           module.12.batch_norm_12.weight      True         1024               [1024]      0.493      0.291\n",
      "   20             module.12.batch_norm_12.bias      True         1024               [1024]          0          0\n",
      "   21                 module.13.conv_13.weight      True       262144    [256, 1024, 1, 1]  -5.87e-05      0.018\n",
      "   22           module.13.batch_norm_13.weight      True          256                [256]      0.488      0.286\n",
      "   23             module.13.batch_norm_13.bias      True          256                [256]          0          0\n",
      "   24                 module.14.conv_14.weight      True  1.17965e+06     [512, 256, 3, 3]  -7.33e-06      0.012\n",
      "   25           module.14.batch_norm_14.weight      True          512                [512]      0.493      0.291\n",
      "   26             module.14.batch_norm_14.bias      True          512                [512]          0          0\n",
      "   27                 module.15.conv_15.weight      True       130560     [255, 512, 1, 1]    8.9e-05     0.0255\n",
      "   28                   module.15.conv_15.bias      True          255                [255]   -0.00128     0.0244\n",
      "   29                 module.18.conv_18.weight      True        32768     [128, 256, 1, 1]  -0.000175      0.036\n",
      "   30           module.18.batch_norm_18.weight      True          128                [128]      0.475      0.278\n",
      "   31             module.18.batch_norm_18.bias      True          128                [128]          0          0\n",
      "   32                 module.21.conv_21.weight      True       884736     [256, 384, 3, 3]  -5.51e-06    0.00983\n",
      "   33           module.21.batch_norm_21.weight      True          256                [256]      0.516      0.302\n",
      "   34             module.21.batch_norm_21.bias      True          256                [256]          0          0\n",
      "   35                 module.22.conv_22.weight      True        65280     [255, 256, 1, 1]  -0.000114     0.0361\n",
      "   36                   module.22.conv_22.bias      True          255                [255]   -0.00116      0.036\n",
      "Model Summary: 37 layers, 8.85237e+06 parameters, 8.85237e+06 gradients\n"
     ]
    }
   ],
   "source": [
    "model_info(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimizer\n",
    "lr0 = 0.001  # initial learning rate\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=lr0, momentum=0.9, weight_decay=0.0005)\n",
    "\n",
    "# Scheduler (reduce lr at epochs 218, 245, i.e. batches 400k, 450k)\n",
    "scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[218, 245], gamma=0.1,\n",
    "                                                 last_epoch=START_EPOCH - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "resume. load latest.pt\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Compute efficiency:1.00, epoch:79/240: 100%|██████████| 1833/1833 [20:43<00:00,  1.17it/s, loss=1.03]\n",
      "100%|██████████| 78/78 [01:28<00:00,  1.02s/it]\n",
      "Compute efficiency:0.87, epoch:80/240: 100%|██████████| 1833/1833 [20:57<00:00,  5.05it/s, loss=0.962]\n",
      "100%|██████████| 78/78 [01:29<00:00,  1.53it/s]\n",
      "Compute efficiency:0.95, epoch:81/240: 100%|██████████| 1833/1833 [20:18<00:00,  4.03it/s, loss=0.921]\n",
      "100%|██████████| 78/78 [01:21<00:00,  1.69it/s]\n",
      "Compute efficiency:0.94, epoch:82/240: 100%|██████████| 1833/1833 [20:12<00:00,  4.93it/s, loss=0.927]\n",
      "100%|██████████| 78/78 [01:14<00:00,  1.94it/s]\n",
      "Compute efficiency:0.93, epoch:83/240: 100%|██████████| 1833/1833 [20:18<00:00,  6.01it/s, loss=0.956]\n",
      "100%|██████████| 78/78 [01:14<00:00,  1.62it/s]\n",
      "Compute efficiency:0.95, epoch:84/240: 100%|██████████| 1833/1833 [20:35<00:00,  4.64it/s, loss=1.03]\n",
      "100%|██████████| 78/78 [01:23<00:00,  2.10it/s]\n",
      "Compute efficiency:0.94, epoch:85/240: 100%|██████████| 1833/1833 [20:50<00:00,  4.94it/s, loss=1.02]\n",
      "100%|██████████| 78/78 [01:23<00:00,  2.00it/s]\n",
      "Compute efficiency:0.94, epoch:86/240: 100%|██████████| 1833/1833 [20:24<00:00,  4.72it/s, loss=0.944]\n",
      "100%|██████████| 78/78 [01:17<00:00,  2.14it/s]\n",
      "Compute efficiency:0.93, epoch:87/240: 100%|██████████| 1833/1833 [20:06<00:00,  5.11it/s, loss=0.949]\n",
      "100%|██████████| 78/78 [01:14<00:00,  1.87it/s]\n",
      "Compute efficiency:0.95, epoch:88/240: 100%|██████████| 1833/1833 [20:50<00:00,  4.48it/s, loss=0.921]\n",
      "100%|██████████| 78/78 [01:32<00:00,  1.51it/s]\n",
      "Compute efficiency:0.93, epoch:89/240: 100%|██████████| 1833/1833 [20:36<00:00,  4.82it/s, loss=0.922]\n",
      "100%|██████████| 78/78 [01:40<00:00,  1.04it/s]\n",
      "Compute efficiency:0.94, epoch:90/240: 100%|██████████| 1833/1833 [21:07<00:00,  4.64it/s, loss=0.954]\n",
      "100%|██████████| 78/78 [01:35<00:00,  1.49it/s]\n",
      "Compute efficiency:0.95, epoch:91/240: 100%|██████████| 1833/1833 [20:56<00:00,  4.67it/s, loss=1.1]\n",
      "100%|██████████| 78/78 [01:42<00:00,  1.29it/s]\n",
      "Compute efficiency:0.95, epoch:92/240: 100%|██████████| 1833/1833 [20:20<00:00,  4.78it/s, loss=0.907]\n",
      "100%|██████████| 78/78 [01:17<00:00,  1.96it/s]\n",
      "Compute efficiency:0.94, epoch:93/240: 100%|██████████| 1833/1833 [20:21<00:00,  4.56it/s, loss=0.933]\n",
      "100%|██████████| 78/78 [01:20<00:00,  1.87it/s]\n",
      "Compute efficiency:0.70, epoch:94/240:  74%|███████▍  | 1355/1833 [15:31<03:49,  2.08it/s, loss=1.67]"
     ]
    }
   ],
   "source": [
    "resume = True\n",
    "transfer = False\n",
    "\n",
    "if resume:  # Load previously saved model\n",
    "    print('resume. load latest.pt')\n",
    "    if transfer:  # Transfer learning\n",
    "        chkpt = torch.load(weights + 'yolov3-spp.pt', map_location=device)\n",
    "        model.load_state_dict({k: v for k, v in chkpt['model'].items() if v.numel() > 1 and v.shape[0] != 255},\n",
    "                              strict=False)\n",
    "        for p in model.parameters():\n",
    "            p.requires_grad = True if p.shape[0] == nf else False\n",
    "\n",
    "    else:  # resume from latest.pt\n",
    "        chkpt = torch.load(latest, map_location=device)  # load checkpoint\n",
    "        model.load_state_dict(chkpt['model'])\n",
    "\n",
    "    START_EPOCH = chkpt['epoch'] + 1\n",
    "    if chkpt['optimizer'] is not None:\n",
    "        optimizer.load_state_dict(chkpt['optimizer'])\n",
    "        best_loss = chkpt['best_loss']\n",
    "    del chkpt\n",
    "\n",
    "\n",
    "# Start training\n",
    "train_start = time.time()\n",
    "nB = len(dataloader) # num of batches\n",
    "n_burnin = min(round(nB / 5 + 1), 1000)\n",
    "accumulate = 1\n",
    "multi_scale = False\n",
    "\n",
    "#print('Start training with Batch_size: {}, Eopch: {}, Batches: {}'.format(BATCH_SIZE, EPOCHES, len(dataloader)))\n",
    "for epoch in range(START_EPOCH, EPOCHES):\n",
    "    model.train()\n",
    "    #print(('\\n%8s%12s' + '%10s' * 7) % ('Epoch', 'Batch', 'xy', 'wh', 'conf', 'cls', 'total', 'nTargets', 'time'))\n",
    "\n",
    "    # Update scheduler\n",
    "    scheduler.step()\n",
    "\n",
    "    mloss = torch.zeros(5).to(device) # mean losses\n",
    "    #for i, (imgs, targets, _, _) in enumerate(tqdm(dataloader, desc='{}/{}'.format(epoch, EPOCHES-1))):\n",
    "    with tqdm(enumerate(BackgroundGenerator(dataloader)),total=len(dataloader)) as pbar:\n",
    "        start_time = time.time()\n",
    "        for i, (imgs, targets, _, _) in pbar:\n",
    "            #pbar.set_description('{}/{}'.format(epoch, EPOCHES))\n",
    "            imgs = imgs.to(device)\n",
    "            targets = targets.to(device)\n",
    "            nt = len(targets)\n",
    "\n",
    "            # SGD burn-in\n",
    "            if epoch == 0 and i <= n_burnin:\n",
    "                lr = lr0 * (i / n_burnin) ** 4\n",
    "                for x in optimizer.param_groups:\n",
    "                    x['lr'] = lr\n",
    "                    \n",
    "            prepare_time = start_time-time.time()\n",
    "            \n",
    "            # Run model\n",
    "            pred = model(imgs)\n",
    "\n",
    "            # Build targets\n",
    "            target_list = build_targets(model, targets)\n",
    "\n",
    "            # Compute loss\n",
    "            loss, loss_items = compute_loss(pred, target_list)\n",
    "\n",
    "            # Compute gradient\n",
    "            loss.backward()\n",
    "\n",
    "            # Accumulate gradient for x batches before optimizing\n",
    "            if (i + 1) % accumulate == 0 or (i + 1) == nB:\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "            # Update running mean of tracked metrics\n",
    "            mloss = (mloss * i + loss_items) / (i + 1)\n",
    "\n",
    "            # Print batch results\n",
    "            s = ('%8s%12s' + '%10.3g' * 7) % (\n",
    "                '%g/%g' % (epoch, EPOCHES - 1),\n",
    "                '%g/%g' % (i, nB - 1), *mloss, nt, time.time() - train_start)\n",
    "            \n",
    "            process_time = start_time-time.time()-prepare_time\n",
    "            pbar.set_description(\"Compute efficiency:{:.2f}, epoch:{}/{}\".format(\n",
    "                process_time/(process_time+prepare_time), epoch+1, EPOCHES))\n",
    "            start_time = time.time()\n",
    "            pbar.set_postfix(loss=loss.item())\n",
    "            # add graph is cost time\n",
    "            #if epoch == 0:\n",
    "                #TBoard.add_graph(model, (imgs,))\n",
    "            #step_num = epoch * len(dataloader) + i\n",
    "            #imgs_batch = torchvision.utils.make_grid(imgs, nrow=5)\n",
    "            #TBoard.add_image('images', imgs_batch, step_num)\n",
    "            #TBoard.add_graph(cnn, (b_x, ))\n",
    "\n",
    "            \n",
    "\n",
    "    # Calculate mAP\n",
    "    if epoch >= 30:\n",
    "        with torch.no_grad():\n",
    "            results = rtest.test(CFG, DATA_CFG, batch_size=BATCH_SIZE, img_size=IMG_SIZE, model=model, conf_thres=0.1);\n",
    "        \n",
    "        TBoard.add_scalar('gMetrics/train_loss', loss.item(), epoch)\n",
    "        TBoard.add_scalar('gMetrics/P', results[0], epoch)\n",
    "        TBoard.add_scalar('gMetrics/R', results[1], epoch)\n",
    "        TBoard.add_scalar('gMetrics/mAP', results[2], epoch)\n",
    "        TBoard.add_scalar('gMetrics/F1', results[3], epoch)\n",
    "        TBoard.add_scalar('gMetrics/test_loss', results[4], epoch)\n",
    "\n",
    "    # Write epoch results\n",
    "        with open('results.txt', 'a') as file:\n",
    "            file.write(s + '%11.3g' * 5 % results + '\\n')  # P, R, mAP, F1, test_loss\n",
    "\n",
    "        # Update best loss\n",
    "        test_loss = results[4]\n",
    "        if test_loss < best_loss:\n",
    "            best_loss = test_loss\n",
    "        \n",
    "\n",
    "    # Save training results\n",
    "    save = True\n",
    "    if save and epoch >= 30:\n",
    "        # Create checkpoint\n",
    "        chkpt = {'epoch': epoch,\n",
    "                 'best_loss': best_loss,\n",
    "                 'model': model.module.state_dict() if type(\n",
    "                     model) is nn.parallel.DistributedDataParallel else model.state_dict(),\n",
    "                 'optimizer': optimizer.state_dict()}\n",
    "\n",
    "        # Save latest checkpoint\n",
    "        torch.save(chkpt, latest)\n",
    "\n",
    "        # Save best checkpoint\n",
    "        if best_loss == test_loss:\n",
    "            torch.save(chkpt, best)\n",
    "\n",
    "        # Save backup every 10 epochs (optional)\n",
    "        if epoch > 0 and epoch % 10 == 0:\n",
    "            torch.save(chkpt, weights + 'backup%g.pt' % epoch)\n",
    "\n",
    "        # Delete checkpoint\n",
    "        del chkpt\n",
    "TBoard.close()\n",
    "end = time.time()\n",
    "print('Training finished! using time: {}'.format(end - train_start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}