added frozen layers

Author: yonigozlan

demo/demo.py

@@ -460,7 +460,9 @@ def main():
                         )
                         if len(pose_preds) >= 1:
                             for kpt in pose_preds:
-                                if len(kpt) == 41:
+                                if len(kpt) == 58:
+                                    draw_pose_infinity_coco(kpt, image_bgr)
+                                elif len(kpt) == 41:
                                     draw_pose_infinity(kpt, image_bgr)
                                 else:
                                     draw_pose(kpt, image_bgr)  # draw the poses

experiments/infinity_coco/hrnet/w48_384x288_adam_lr1e-3.yaml

@@ -5,8 +5,8 @@ CUDNN:
   ENABLED: true
 DATA_DIR: ""
 GPUS: (0,)
-OUTPUT_DIR: "output_infinity_coco"
-LOG_DIR: "log_infinity_coco"
+OUTPUT_DIR: "output_infinity_coco_frozen"
+LOG_DIR: "log_infinity_coco_frozen"
 WORKERS: 2
 PRINT_FREQ: 10
 
@@ -36,6 +36,18 @@ MODEL:
     - 96
   SIGMA: 3
   EXTRA:
+    FREEZE_LAYERS: true
+    FROZEN_LAYERS:
+      - "conv1"
+      - "bn1"
+      - "conv2"
+      - "bn2"
+      - "layer1"
+      - "transition1"
+      - "stage2"
+      - "transition2"
+      - "stage3"
+      - "transition3"
     PRETRAINED_LAYERS:
       - "conv1"
       - "bn1"
@@ -91,7 +103,7 @@ MODEL:
 LOSS:
   USE_TARGET_WEIGHT: true
 TRAIN:
-  BATCH_SIZE_PER_GPU: 40
+  BATCH_SIZE_PER_GPU: 2
   SHUFFLE: true
   BEGIN_EPOCH: 0
   END_EPOCH: 200

lib/core/function.py

@@ -4,35 +4,46 @@
 # Written by Bin Xiao (Bin.Xiao@microsoft.com)
 # ------------------------------------------------------------------------------
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
- 
-import time
+from __future__ import absolute_import, division, print_function
+
 import logging
 import os
+import time
 
 import numpy as np
 import torch
-
 from core.evaluate import accuracy
 from core.inference import get_final_preds
 from utils.transforms import flip_back
 from utils.vis import save_debug_images
 
-
 logger = logging.getLogger(__name__)
 
 
-def train(config, train_loader, model, criterion, optimizer, epoch,
-          output_dir, tb_log_dir, writer_dict):
+def train(
+    config,
+    train_loader,
+    model,
+    criterion,
+    optimizer,
+    epoch,
+    output_dir,
+    tb_log_dir,
+    writer_dict,
+):
     batch_time = AverageMeter()
     data_time = AverageMeter()
     losses = AverageMeter()
     acc = AverageMeter()
 
     # switch to train mode
     model.train()
+    # freeze specified layers
+    extra = config.MODEL.EXTRA
+    if "FREEZE_LAYERS" in extra and extra["FREEZE_LAYERS"]:
+        frozen_layers = extra.FROZEN_LAYERS
+        for layer in frozen_layers:
+            eval("model.module." + layer + ".requires_grad_(False)")
 
     end = time.time()
     for i, (input, target, target_weight, meta) in enumerate(train_loader):
@@ -63,39 +74,55 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
         # measure accuracy and record loss
         losses.update(loss.item(), input.size(0))
 
-        _, avg_acc, cnt, pred = accuracy(output.detach().cpu().numpy(),
-                                         target.detach().cpu().numpy())
+        _, avg_acc, cnt, pred = accuracy(
+            output.detach().cpu().numpy(), target.detach().cpu().numpy()
+        )
         acc.update(avg_acc, cnt)
 
         # measure elapsed time
         batch_time.update(time.time() - end)
         end = time.time()
 
         if i % config.PRINT_FREQ == 0:
-            msg = 'Epoch: [{0}][{1}/{2}]\t' \
-                  'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
-                  'Speed {speed:.1f} samples/s\t' \
-                  'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
-                  'Loss {loss.val:.5f} ({loss.avg:.5f})\t' \
-                  'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
-                      epoch, i, len(train_loader), batch_time=batch_time,
-                      speed=input.size(0)/batch_time.val,
-                      data_time=data_time, loss=losses, acc=acc)
+            msg = (
+                "Epoch: [{0}][{1}/{2}]\t"
+                "Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t"
+                "Speed {speed:.1f} samples/s\t"
+                "Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t"
+                "Loss {loss.val:.5f} ({loss.avg:.5f})\t"
+                "Accuracy {acc.val:.3f} ({acc.avg:.3f})".format(
+                    epoch,
+                    i,
+                    len(train_loader),
+                    batch_time=batch_time,
+                    speed=input.size(0) / batch_time.val,
+                    data_time=data_time,
+                    loss=losses,
+                    acc=acc,
+                )
+            )
             logger.info(msg)
 
-            writer = writer_dict['writer']
-            global_steps = writer_dict['train_global_steps']
-            writer.add_scalar('train_loss', losses.val, global_steps)
-            writer.add_scalar('train_acc', acc.val, global_steps)
-            writer_dict['train_global_steps'] = global_steps + 1
-
-            prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
-            save_debug_images(config, input, meta, target, pred*4, output,
-                              prefix)
-
-
-def validate(config, val_loader, val_dataset, model, criterion, output_dir,
-             tb_log_dir, writer_dict=None):
+            writer = writer_dict["writer"]
+            global_steps = writer_dict["train_global_steps"]
+            writer.add_scalar("train_loss", losses.val, global_steps)
+            writer.add_scalar("train_acc", acc.val, global_steps)
+            writer_dict["train_global_steps"] = global_steps + 1
+
+            prefix = "{}_{}".format(os.path.join(output_dir, "train"), i)
+            save_debug_images(config, input, meta, target, pred * 4, output, prefix)
+
+
+def validate(
+    config,
+    val_loader,
+    val_dataset,
+    model,
+    criterion,
+    output_dir,
+    tb_log_dir,
+    writer_dict=None,
+):
     batch_time = AverageMeter()
     losses = AverageMeter()
     acc = AverageMeter()
@@ -104,10 +131,7 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
     model.eval()
 
     num_samples = len(val_dataset)
-    all_preds = np.zeros(
-        (num_samples, config.MODEL.NUM_JOINTS, 3),
-        dtype=np.float32
-    )
+    all_preds = np.zeros((num_samples, config.MODEL.NUM_JOINTS, 3), dtype=np.float32)
     all_boxes = np.zeros((num_samples, 6))
     image_path = []
     filenames = []
@@ -132,15 +156,14 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
                 else:
                     output_flipped = outputs_flipped
 
-                output_flipped = flip_back(output_flipped.cpu().numpy(),
-                                           val_dataset.flip_pairs)
+                output_flipped = flip_back(
+                    output_flipped.cpu().numpy(), val_dataset.flip_pairs
+                )
                 output_flipped = torch.from_numpy(output_flipped.copy()).cuda()
 
-
                 # feature is not aligned, shift flipped heatmap for higher accuracy
                 if config.TEST.SHIFT_HEATMAP:
-                    output_flipped[:, :, :, 1:] = \
-                        output_flipped.clone()[:, :, :, 0:-1]
+                    output_flipped[:, :, :, 1:] = output_flipped.clone()[:, :, :, 0:-1]
 
                 output = (output + output_flipped) * 0.5
 
@@ -152,51 +175,47 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
             num_images = input.size(0)
             # measure accuracy and record loss
             losses.update(loss.item(), num_images)
-            _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
-                                             target.cpu().numpy())
+            _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(), target.cpu().numpy())
 
             acc.update(avg_acc, cnt)
 
             # measure elapsed time
             batch_time.update(time.time() - end)
             end = time.time()
 
-            c = meta['center'].numpy()
-            s = meta['scale'].numpy()
-            score = meta['score'].numpy()
+            c = meta["center"].numpy()
+            s = meta["scale"].numpy()
+            score = meta["score"].numpy()
 
-            preds, maxvals = get_final_preds(
-                config, output.clone().cpu().numpy(), c, s)
+            preds, maxvals = get_final_preds(config, output.clone().cpu().numpy(), c, s)
 
-            all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
-            all_preds[idx:idx + num_images, :, 2:3] = maxvals
+            all_preds[idx : idx + num_images, :, 0:2] = preds[:, :, 0:2]
+            all_preds[idx : idx + num_images, :, 2:3] = maxvals
             # double check this all_boxes parts
-            all_boxes[idx:idx + num_images, 0:2] = c[:, 0:2]
-            all_boxes[idx:idx + num_images, 2:4] = s[:, 0:2]
-            all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)
-            all_boxes[idx:idx + num_images, 5] = score
-            image_path.extend(meta['image'])
+            all_boxes[idx : idx + num_images, 0:2] = c[:, 0:2]
+            all_boxes[idx : idx + num_images, 2:4] = s[:, 0:2]
+            all_boxes[idx : idx + num_images, 4] = np.prod(s * 200, 1)
+            all_boxes[idx : idx + num_images, 5] = score
+            image_path.extend(meta["image"])
 
             idx += num_images
 
             if i % config.PRINT_FREQ == 0:
-                msg = 'Test: [{0}/{1}]\t' \
-                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
-                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
-                      'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
-                          i, len(val_loader), batch_time=batch_time,
-                          loss=losses, acc=acc)
+                msg = (
+                    "Test: [{0}/{1}]\t"
+                    "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
+                    "Loss {loss.val:.4f} ({loss.avg:.4f})\t"
+                    "Accuracy {acc.val:.3f} ({acc.avg:.3f})".format(
+                        i, len(val_loader), batch_time=batch_time, loss=losses, acc=acc
+                    )
+                )
                 logger.info(msg)
 
-                prefix = '{}_{}'.format(
-                    os.path.join(output_dir, 'val'), i
-                )
-                save_debug_images(config, input, meta, target, pred*4, output,
-                                  prefix)
+                prefix = "{}_{}".format(os.path.join(output_dir, "val"), i)
+                save_debug_images(config, input, meta, target, pred * 4, output, prefix)
 
         name_values, perf_indicator = val_dataset.evaluate(
-            config, all_preds, output_dir, all_boxes, image_path,
-            filenames, imgnums
+            config, all_preds, output_dir, all_boxes, image_path, filenames, imgnums
         )
 
         model_name = config.MODEL.NAME
@@ -207,32 +226,16 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
             _print_name_value(name_values, model_name)
 
         if writer_dict:
-            writer = writer_dict['writer']
-            global_steps = writer_dict['valid_global_steps']
-            writer.add_scalar(
-                'valid_loss',
-                losses.avg,
-                global_steps
-            )
-            writer.add_scalar(
-                'valid_acc',
-                acc.avg,
-                global_steps
-            )
+            writer = writer_dict["writer"]
+            global_steps = writer_dict["valid_global_steps"]
+            writer.add_scalar("valid_loss", losses.avg, global_steps)
+            writer.add_scalar("valid_acc", acc.avg, global_steps)
             if isinstance(name_values, list):
                 for name_value in name_values:
-                    writer.add_scalars(
-                        'valid',
-                        dict(name_value),
-                        global_steps
-                    )
+                    writer.add_scalars("valid", dict(name_value), global_steps)
             else:
-                writer.add_scalars(
-                    'valid',
-                    dict(name_values),
-                    global_steps
-                )
-            writer_dict['valid_global_steps'] = global_steps + 1
+                writer.add_scalars("valid", dict(name_values), global_steps)
+            writer_dict["valid_global_steps"] = global_steps + 1
 
     return perf_indicator
 
@@ -242,24 +245,23 @@ def _print_name_value(name_value, full_arch_name):
     names = name_value.keys()
     values = name_value.values()
     num_values = len(name_value)
-    logger.info(
-        '| Arch ' +
-        ' '.join(['| {}'.format(name) for name in names]) +
-        ' |'
-    )
-    logger.info('|---' * (num_values+1) + '|')
+    logger.info("| Arch " + " ".join(["| {}".format(name) for name in names]) + " |")
+    logger.info("|---" * (num_values + 1) + "|")
 
     if len(full_arch_name) > 15:
-        full_arch_name = full_arch_name[:8] + '...'
+        full_arch_name = full_arch_name[:8] + "..."
     logger.info(
-        '| ' + full_arch_name + ' ' +
-        ' '.join(['| {:.3f}'.format(value) for value in values]) +
-         ' |'
+        "| "
+        + full_arch_name
+        + " "
+        + " ".join(["| {:.3f}".format(value) for value in values])
+        + " |"
     )
 
 
 class AverageMeter(object):
     """Computes and stores the average and current value"""
+
     def __init__(self):
         self.reset()
 

tools/test.py

@@ -5,62 +5,49 @@
 # Written by Bin Xiao (Bin.Xiao@microsoft.com)
 # ------------------------------------------------------------------------------
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+from __future__ import absolute_import, division, print_function
 
 import argparse
 import os
 import pprint
 
+import _init_paths
+import dataset
 import torch
-import torch.nn.parallel
 import torch.backends.cudnn as cudnn
+import torch.nn.parallel
 import torch.optim
 import torch.utils.data
 import torch.utils.data.distributed
 import torchvision.transforms as transforms
-
-import _init_paths
-from config import cfg
-from config import update_config
-from core.loss import JointsMSELoss
+from config import cfg, update_config
 from core.function import validate
+from core.loss import JointsMSELoss
 from utils.utils import create_logger
 
-import dataset
 import models
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description='Train keypoints network')
+    parser = argparse.ArgumentParser(description="Train keypoints network")
     # general
-    parser.add_argument('--cfg',
-                        help='experiment configure file name',
-                        required=True,
-                        type=str)
-
-    parser.add_argument('opts',
-                        help="Modify config options using the command-line",
-                        default=None,
-                        nargs=argparse.REMAINDER)
-
-    parser.add_argument('--modelDir',
-                        help='model directory',
-                        type=str,
-                        default='')
-    parser.add_argument('--logDir',
-                        help='log directory',
-                        type=str,
-                        default='')
-    parser.add_argument('--dataDir',
-                        help='data directory',
-                        type=str,
-                        default='')
-    parser.add_argument('--prevModelDir',
-                        help='prev Model directory',
-                        type=str,
-                        default='')
+    parser.add_argument(
+        "--cfg", help="experiment configure file name", required=True, type=str
+    )
+
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+
+    parser.add_argument("--modelDir", help="model directory", type=str, default="")
+    parser.add_argument("--logDir", help="log directory", type=str, default="")
+    parser.add_argument("--dataDir", help="data directory", type=str, default="")
+    parser.add_argument(
+        "--prevModelDir", help="prev Model directory", type=str, default=""
+    )
 
     args = parser.parse_args()
     return args
@@ -70,8 +57,7 @@ def main():
     args = parse_args()
     update_config(cfg, args)
 
-    logger, final_output_dir, tb_log_dir = create_logger(
-        cfg, args.cfg, 'valid')
+    logger, final_output_dir, tb_log_dir = create_logger(cfg, args.cfg, "valid")
 
     logger.info(pprint.pformat(args))
     logger.info(cfg)
@@ -81,50 +67,49 @@ def main():
     torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
     torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED
 
-    model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
-        cfg, is_train=False
-    )
+    model = eval("models." + cfg.MODEL.NAME + ".get_pose_net")(cfg, is_train=False)
 
     if cfg.TEST.MODEL_FILE:
-        logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
+        logger.info("=> loading model from {}".format(cfg.TEST.MODEL_FILE))
         model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
     else:
-        model_state_file = os.path.join(
-            final_output_dir, 'final_state.pth'
-        )
-        logger.info('=> loading model from {}'.format(model_state_file))
+        model_state_file = os.path.join(final_output_dir, "final_state.pth")
+        logger.info("=> loading model from {}".format(model_state_file))
         model.load_state_dict(torch.load(model_state_file))
 
     model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
-
     # define loss function (criterion) and optimizer
-    criterion = JointsMSELoss(
-        use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT
-    ).cuda()
+    criterion = JointsMSELoss(use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT).cuda()
 
     # Data loading code
     normalize = transforms.Normalize(
         mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
     )
-    valid_dataset = eval('dataset.'+cfg.DATASET.DATASET)(
-        cfg, cfg.DATASET.ROOT, cfg.DATASET.TEST_SET, False,
-        transforms.Compose([
-            transforms.ToTensor(),
-            normalize,
-        ])
+    valid_dataset = eval("dataset." + cfg.DATASET.DATASET)(
+        cfg,
+        cfg.DATASET.ROOT,
+        cfg.DATASET.TEST_SET,
+        False,
+        transforms.Compose(
+            [
+                transforms.ToTensor(),
+                normalize,
+            ]
+        ),
     )
     valid_loader = torch.utils.data.DataLoader(
         valid_dataset,
-        batch_size=cfg.TEST.BATCH_SIZE_PER_GPU*len(cfg.GPUS),
+        batch_size=cfg.TEST.BATCH_SIZE_PER_GPU * len(cfg.GPUS),
         shuffle=False,
         num_workers=cfg.WORKERS,
-        pin_memory=True
+        pin_memory=True,
     )
 
     # evaluate on validation set
-    validate(cfg, valid_loader, valid_dataset, model, criterion,
-             final_output_dir, tb_log_dir)
+    validate(
+        cfg, valid_loader, valid_dataset, model, criterion, final_output_dir, tb_log_dir
+    )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()