Generating_figure.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import argparse\n",
    "import os\n",
    "import pprint\n",
    "import shutil\n",
    "\n",
    "import torch\n",
    "# import torch.nn.parallel\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim\n",
    "import torchvision.transforms as transforms\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "# import _init_paths\n",
    "from lib.core.config import config\n",
    "from lib.core.config import update_config\n",
    "from lib.core.config import update_dir\n",
    "from lib.core.config import get_model_name\n",
    "from lib.core.loss import JointsMSELoss\n",
    "from lib.core.function import train\n",
    "from lib.core.function import validate\n",
    "from lib.utils.utils import get_optimizer\n",
    "from lib.utils.utils import save_checkpoint\n",
    "from lib.utils.utils import create_logger\n",
    "\n",
    "import lib.dataset as dataset\n",
    "import lib.models as models\n",
    "\n",
    "\n",
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description='Train keypoints network')\n",
    "    # general\n",
    "    parser.add_argument('--cfg',\n",
    "                        help='experiment configure file name',\n",
    "                        required=True,\n",
    "                        type=str)\n",
    "\n",
    "    args, rest = parser.parse_known_args()\n",
    "    # update config\n",
    "    update_config(args.cfg)\n",
    "\n",
    "    # training\n",
    "    parser.add_argument('--frequent',\n",
    "                        help='frequency of logging',\n",
    "                        default=config.PRINT_FREQ,\n",
    "                        type=int)\n",
    "    parser.add_argument('--gpus',\n",
    "                        help='gpus',\n",
    "                        type=str)\n",
    "    parser.add_argument('--workers',\n",
    "                        help='num of dataloader workers',\n",
    "                        type=int)\n",
    "\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    return args\n",
    "\n",
    "\n",
    "def reset_config(config, args):\n",
    "    if args.gpus:\n",
    "        config.GPUS = args.gpus\n",
    "    if args.workers:\n",
    "        config.WORKERS = args.workers\n",
    "\n",
    "\n",
    "def main():\n",
    "    #args = parse_args()\n",
    "    with open('experiments/simulated/128x128_d256x3_adam_lr1e-3.yaml') as file:\n",
    "        args = yaml.load(file, Loader=yaml.FullLoader)\n",
    "    config = dotdict(args)\n",
    "    \n",
    "    #reset_config(config, args)\n",
    "\n",
    "    logger, final_output_dir, tb_log_dir = create_logger(\n",
    "        config, args.cfg, 'train')\n",
    "\n",
    "    logger.info(pprint.pformat(args))\n",
    "    logger.info(pprint.pformat(config))\n",
    "\n",
    "    # cudnn related setting\n",
    "    cudnn.benchmark = config.CUDNN.BENCHMARK\n",
    "    torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC\n",
    "    torch.backends.cudnn.enabled = config.CUDNN.ENABLED\n",
    "\n",
    "    model = eval('models.'+config.MODEL.NAME+'.get_neuron_net')(\n",
    "        config, is_train=True\n",
    "    )\n",
    "\n",
    "    # copy model file\n",
    "    this_dir = os.path.dirname(__file__)\n",
    "    shutil.copy2(\n",
    "        os.path.join(this_dir, '../lib/models', config.MODEL.NAME + '.py'),\n",
    "        final_output_dir)\n",
    "\n",
    "    writer_dict = {\n",
    "        'writer': SummaryWriter(log_dir=tb_log_dir),\n",
    "        'train_global_steps': 0,\n",
    "        'valid_global_steps': 0,\n",
    "        'vis_global_steps': 0,\n",
    "    }\n",
    "\n",
    "    dump_input = torch.rand((config.TRAIN.BATCH_SIZE,\n",
    "                             3,\n",
    "                             config.MODEL.IMAGE_SIZE[1],\n",
    "                             config.MODEL.IMAGE_SIZE[0]))\n",
    "    writer_dict['writer'].add_graph(model, (dump_input, ), verbose=False)\n",
    "\n",
    "    gpus = [int(i) for i in config.GPUS.split(',')]\n",
    "    model = torch.nn.DataParallel(model, device_ids=gpus).cuda()\n",
    "\n",
    "    # define loss function (criterion) and optimizer\n",
    "    criterion = JointsMSELoss(\n",
    "        use_target_weight=config.LOSS.USE_TARGET_WEIGHT\n",
    "    ).cuda()\n",
    "\n",
    "    optimizer = get_optimizer(config, model)\n",
    "\n",
    "    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(\n",
    "        optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR\n",
    "    )\n",
    "\n",
    "    # Data loading code\n",
    "    normalize = transforms.Normalize(mean=[0.1440, 0.1440, 0.1440],\n",
    "                                     std=[19.0070, 19.0070, 19.0070])\n",
    "    train_dataset = eval('dataset.'+config.DATASET.DATASET)(\n",
    "        config,\n",
    "        config.DATASET.ROOT,\n",
    "        config.DATASET.TRAIN_SET,\n",
    "        True,\n",
    "        transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    )\n",
    "    valid_dataset = eval('dataset.'+config.DATASET.DATASET)(\n",
    "        config,\n",
    "        config.DATASET.ROOT,\n",
    "        config.DATASET.TEST_SET,\n",
    "        False,\n",
    "        transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    )\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=config.TRAIN.BATCH_SIZE*len(gpus),\n",
    "        shuffle=config.TRAIN.SHUFFLE,\n",
    "        num_workers=config.WORKERS,\n",
    "        pin_memory=True\n",
    "    )\n",
    "    valid_loader = torch.utils.data.DataLoader(\n",
    "        valid_dataset,\n",
    "        batch_size=config.TEST.BATCH_SIZE*len(gpus),\n",
    "        shuffle=False,\n",
    "        num_workers=config.WORKERS,\n",
    "        pin_memory=True\n",
    "    )\n",
    "\n",
    "    best_perf = 0.0\n",
    "    best_model = False\n",
    "    for epoch in range(config.TRAIN.BEGIN_EPOCH, config.TRAIN.END_EPOCH):\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        # train for one epoch\n",
    "        train(config, train_loader, model, criterion, optimizer, epoch,\n",
    "              final_output_dir, tb_log_dir, writer_dict)\n",
    "\n",
    "\n",
    "        # evaluate on validation set\n",
    "        perf_indicator = validate(config, valid_loader, valid_dataset, model,\n",
    "                                  criterion, final_output_dir, tb_log_dir,\n",
    "                                  writer_dict)\n",
    "\n",
    "        if perf_indicator > best_perf:\n",
    "            best_perf = perf_indicator\n",
    "            best_model = True\n",
    "        else:\n",
    "            best_model = False\n",
    "\n",
    "        logger.info('=> saving checkpoint to {}'.format(final_output_dir))\n",
    "        save_checkpoint({\n",
    "            'epoch': epoch + 1,\n",
    "            'model': get_model_name(config),\n",
    "            'state_dict': model.state_dict(),\n",
    "            'perf': perf_indicator,\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "        }, best_model, final_output_dir)\n",
    "\n",
    "    final_model_state_file = os.path.join(final_output_dir,\n",
    "                                          'final_state.pth.tar')\n",
    "    logger.info('saving final model state to {}'.format(\n",
    "        final_model_state_file))\n",
    "    torch.save(model.module.state_dict(), final_model_state_file)\n",
    "    writer_dict['writer'].close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "       \n",
    "with open('experiments/simulated/128x128_d256x3_adam_lr1e-3.yaml') as file:\n",
    "    args = yaml.load(file, Loader=yaml.FullLoader)\n",
    "\n",
    "class dotdict(dict):\n",
    "    \"\"\"dot.notation access to dictionary attributes\"\"\"\n",
    "    __getattr__ = dict.get\n",
    "    __setattr__ = dict.__setitem__\n",
    "    __delattr__ = dict.__delitem__\n",
    "\n",
    "args = dotdict(args)\n",
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyrooz",
   "language": "python",
   "name": "pyrooz"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}