autogluon integration

dmlc · Jerryzcn · Jul 8, 2020 · May 6, 2020 · May 6, 2020 · May 6, 2020
commit 304ce5b03e8dce15711d3546c204757d489a8c54
diff --git a/gluoncv/model_zoo/rcnn/rcnn.py b/gluoncv/model_zoo/rcnn/rcnn.py
@@ -326,8 +326,8 @@ def custom_rcnn_fpn(pretrained_base=True, base_network_name='resnet18_v1b', norm
         fpn_inputs_names = ['layers1_relu8_fwd', 'layers2_relu11_fwd', 'layers3_relu68_fwd',
                             'layers4_relu8_fwd']
     elif base_network_name == 'resnest50':
-        from ...model_zoo.resnest import resnet50
-        base_network = resnet50(pretrained=pretrained_base, dilated=False,
+        from ...model_zoo.resnest import resnest50
+        base_network = resnest50(pretrained=pretrained_base, dilated=False,
                                 use_global_stats=use_global_stats, norm_layer=norm_layer,
                                 norm_kwargs=norm_kwargs)
         fpn_inputs_names = ['layers1_relu11_fwd', 'layers2_relu15_fwd', 'layers3_relu23_fwd',

diff --git a/gluoncv/pipelines/estimators/rcnn/faster_rcnn.py b/gluoncv/pipelines/estimators/rcnn/faster_rcnn.py
@@ -89,6 +89,25 @@ def _get_dataloader(net, train_dataset, val_dataset, train_transform, val_transf
     return train_loader, val_loader
 
 
+def _get_testloader(net, test_dataset, num_devices, config):
+    if config.meta_arch == 'faster_rcnn':
+        """Get faster rcnn dataloader."""
+        test_bfn = Tuple(*[Append() for _ in range(3)])
+        short = net.short[-1] if isinstance(net.short, (tuple, list)) else net.short
+        # validation use 1 sample per device
+        test_loader = gluon.data.DataLoader(
+            test_dataset.transform(FasterRCNNDefaultValTransform(short, net.max_size)),
+            num_devices,
+            False,
+            batchify_fn=test_bfn,
+            last_batch='keep',
+            num_workers=config.num_workers
+        )
+        return test_loader
+    else:
+        raise NotImplementedError('%s not implemented.' % config.meta_arch)
+
+
 def _get_dataset(dataset, args):
     if dataset.lower() == 'voc':
         train_dataset = gdata.VOCDetection(
@@ -120,7 +139,7 @@ class FasterRCNNEstimator:
     TODO: use base estimators.
     """
 
-    def __init__(self, config, logger=None):
+    def __init__(self, config, logger=None, reporter=None):
         """
         Constructs Faster R-CNN estimators.
 
@@ -130,9 +149,12 @@ def __init__(self, config, logger=None):
             Configuration object containing information for constructing Faster R-CNN estimators.
         logger : logger object, default is None
             If not `None`, will use default logging object.
+        reporter : reporter object, default is None
+
         """
         super(FasterRCNNEstimator, self).__init__()
         self._logger = logger if logger is not None else logging.getLogger(__name__)
+        self._reporter = reporter
         self._cfg = config
         # training contexts
         if self._cfg.horovod:
@@ -173,7 +195,7 @@ def __init__(self, config, logger=None):
                 norm_kwargs = None
                 sym_norm_layer = None
                 sym_norm_kwargs = None
-            classes = self._cfg.classes
+            classes = self.train_dataset.CLASSES
             # TODO: maybe refactor this to pass configuration into the model instead.
             self.net = get_model('custom_faster_rcnn_fpn', classes=classes, transfer=None,
                                  dataset=self._cfg.dataset,
@@ -184,7 +206,7 @@ def __init__(self, config, logger=None):
                                  num_fpn_filters=self._cfg.num_fpn_filters,
                                  num_box_head_conv=self._cfg.num_box_head_conv,
                                  num_box_head_conv_filters=self._cfg.num_box_head_conv_filters,
-                                 num_box_head_dense_filters=self.cfg.num_box_head_dense_filters,
+                                 num_box_head_dense_filters=self._cfg.num_box_head_dense_filters,
                                  short=self._cfg.image_short, max_size=self._cfg.image_max_size,
                                  min_stage=2, max_stage=6, nms_thresh=self._cfg.nms_thresh,
                                  nms_topk=self._cfg.nms_topk, post_nms=self._cfg.post_nms,
@@ -433,16 +455,19 @@ def fit(self):
                     current_map = 0.
                 _save_params(self.net, self._logger, best_map, current_map, epoch,
                              self._cfg.save_interval, self._cfg.save_prefix)
+                if self._reporter:
+                    self._reporter(epoch=epoch, map_reward=current_map)
 
-    def evaluate(self):
+    def evaluate(self, dataset):
         """Evaluate the current model on dataset.
 
         Parameters
         ----------
         dataset : mxnet.gluon.data.DataLoader
             DataLoader containing dataset for evaluation.
         """
-        return self._validate(self.val_dataset, self.ctx, self.eval_metric)
+        dataloader = _get_testloader(self.net, dataset, len(self.ctx), self._cfg)
+        return self._validate(dataloader, self.ctx, self.eval_metric)
 
     def predict(self, x):
         """Predict an individual example.
@@ -456,3 +481,15 @@ def predict(self, x):
         x = x.as_in_context(self.ctx[0])
         ids, scores, bboxes = [xx[0].asnumpy() for xx in self.net(x)]
         return ids, scores, bboxes
+
+    def load_parameters(self, parameters, multi_precision=False):
+        """Load saved parameters into the model"""
+        param_dict = self.net._collect_params_with_prefix()
+        kwargs = {'ctx': None} if mx.__version__[:3] == '1.4' else {'cast_dtype': multi_precision,
+                                                                    'ctx': None}
+        for k, v in param_dict.items():
+            param_dict[k]._load_init(parameters[k], **kwargs)
+
+    def get_parameters(self):
+        """Return model parameters"""
+        return self.net._collect_params_with_prefix()
diff --git a/gluoncv/pipelines/tasks/__init__.py b/gluoncv/pipelines/tasks/__init__.py
diff --git a/gluoncv/pipelines/tasks/object_detection.py b/gluoncv/pipelines/tasks/object_detection.py
@@ -0,0 +1,123 @@
+import logging
+
+import autogluon as ag
+from autogluon.core.decorator import sample_config
+from autogluon.scheduler.resource import get_cpu_count, get_gpu_count
+from autogluon.task import BaseTask
+from autogluon.utils import collect_params
+
+from ..estimators.rcnn import FasterRCNNEstimator
+from ... import utils as gutils
+
+__all__ = ['ObjectDetection']
+
+
+@ag.args()
+def _train_object_detection(args, reporter):
+    # fix seed for mxnet, numpy and python builtin random generator.
+    gutils.random.seed(args.seed)
+
+    # training contexts
+    if args.meta_arch == 'yolo3':
+        net_name = '_'.join((args.meta_arch, args.net, 'custom'))
+    elif args.meta_arch == 'faster_rcnn':
+        net_name = '_'.join(('custom', args.meta_arch, 'fpn'))
+        kwargs = {'network': args.net, 'base_network_name': args.net,
+                  'image_short': args.data_shape, 'max_size': 1000, 'nms_thresh': 0.5,
+                  'nms_topk': -1, 'min_stage': 2, 'max_stage': 6, 'post_nms': -1,
+                  'roi_mode': 'align', 'roi_size': (7, 7), 'strides': (4, 8, 16, 32, 64),
+                  'clip': 4.14, 'rpn_channel': 256, 'anchor_scales': (2, 4, 8, 16, 32),
+                  'anchor_aspect_ratio': (0.5, 1, 2), 'anchor_alloc_size': (384, 384),
+                  'rpn_nms_thresh': 0.7, 'rpn_train_pre_nms': 12000, 'rpn_train_post_nms': 2000,
+                  'rpn_test_pre_nms': 6000, 'rpn_test_post_nms': 1000, 'rpn_min_size': 1,
+                  'per_device_batch_size': args.batch_size // args.num_gpus, 'num_sample': 512,
+                  'rcnn_pos_iou_thresh': 0.5, 'rcnn_pos_ratio': 0.25, 'max_num_gt': 100,
+                  'custom_model': True, 'no_pretrained_base': True, 'num_fpn_filters': 256,
+                  'num_box_head_conv': 4, 'num_box_head_conv_filters': 256, 'amp': False,
+                  'num_box_head_dense_filters': 1024, 'image_max_size': 1333, 'kv_store': 'nccl',
+                  'anchor_base_size': 16, 'rcnn_num_samples': 512, 'rpn_smoothl1_rho': 0.001,
+                  'rcnn_smoothl1_rho': 0.001, 'lr_warmup_factor': 1. / 3., 'lr_warmup': 500,
+                  'executor_threads': 4, 'disable_hybridization': False, 'static_alloc': False}
+        vars(args).update(kwargs)
+    else:
+        raise NotImplementedError(args.meta_arch, 'is not implemented.')
+    args.save_prefix += net_name
+
+    if args.meta_arch == 'faster_rcnn':
+        estimator = FasterRCNNEstimator(args, reporter=reporter)
+    else:
+        raise NotImplementedError('%s' % args.meta_arch)
+
+    # training
+    estimator.fit()
+
+    if args.final_fit:
+        return {'model_params': collect_params(estimator.net)}
+
+
+class ObjectDetection(BaseTask):
+    def __init__(self, config, logger=None):
+        super(ObjectDetection, self).__init__()
+        self._logger = logger if logger is not None else logging.getLogger(__name__)
+        self._config = config
+        nthreads_per_trial = get_cpu_count() if self._config.nthreads_per_trial > get_cpu_count() \
+            else self._config.nthreads_per_trial
+        if self._config.ngpus_per_trial > get_gpu_count():
+            self._logger.warning(
+                "The number of requested GPUs is greater than the number of available GPUs.")
+        ngpus_per_trial = get_gpu_count() if self._config.ngpus_per_trial > get_gpu_count() \
+            else self._config.ngpus_per_trial
+
+        _train_object_detection.register_args(
+            meta_arch=self._config.meta_arch, dataset=self._config.dataset, net=self._config.net,
+            lr=self._config.lr, loss=self._config.loss, num_gpus=self._config.ngpus_per_trial,
+            batch_size=self._config.batch_size, split_ratio=self._config.split_ratio,
+            epochs=self._config.epochs, num_workers=self._config.nthreads_per_trial,
+            hybridize=self._config.hybridize, verbose=self._config.verbose, final_fit=False,
+            seed=self._config.seed, data_shape=self._config.data_shape, start_epoch=0,
+            transfer=self._config.transfer, lr_mode=self._config.lr_mode,
+            lr_decay=self._config.lr_decay, lr_decay_period=self._config.lr_decay_period,
+            lr_decay_epoch=self._config.lr_decay_epoch, warmup_lr=self._config.warmup_lr,
+            warmup_epochs=self._config.warmup_epochs, warmup_iters=self._config.warmup_iters,
+            warmup_factor=self._config.warmup_factor, momentum=self._config.momentum,
+            wd=self._config.wd, log_interval=self._config.log_interval,
+            save_prefix=self._config.save_prefix, save_interval=self._config.save_interval,
+            val_interval=self._config.val_interval, num_samples=self._config.num_samples,
+            no_random_shape=self._config.no_random_shape, no_wd=self._config.no_wd,
+            mixup=self._config.mixup, no_mixup_epochs=self._config.no_mixup_epochs,
+            label_smooth=self._config.label_smooth, resume=self._config.resume,
+            syncbn=self._config.syncbn, reuse_pred_weights=self._config.reuse_pred_weights,
+            horovod=self._config.horovod, gpus='0,1,2,3,4,5,6,7', use_fpn=True,
+            norm_layer='syncbn' if self._config.syncbn else None,
+        )
+
+        self._config.scheduler_options = {
+            'resource': {'num_cpus': nthreads_per_trial, 'num_gpus': ngpus_per_trial},
+            'checkpoint': self._config.checkpoint,
+            'num_trials': self._config.num_trials,
+            'time_out': self._config.time_limits,
+            'resume': self._config.resume,
+            'visualizer': self._config.visualizer,
+            'time_attr': 'epoch',
+            'reward_attr': 'map_reward',
+            'dist_ip_addrs': self._config.dist_ip_addrs,
+            'searcher': self._config.search_strategy,
+            'search_options': self._config.search_options,
+        }
+        if self._config.search_strategy == 'hyperband':
+            self._config.scheduler_options.update({
+                'searcher': 'random',
+                'max_t': self._config.epochs,
+                'grace_period': self._config.grace_period if self._config.grace_period
+                else self._config.epochs // 4})
+
+    def fit(self):
+        results = self.run_fit(_train_object_detection, self._config.search_strategy,
+                               self._config.scheduler_options)
+        self._logger.info(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> finish model fitting")
+        best_config = sample_config(_train_object_detection.args, results['best_config'])
+        self._logger.info('The best config: {}'.format(results['best_config']))
+
+        estimator = FasterRCNNEstimator(best_config)
+        estimator.load_parameters(results.pop('model_params'))
+        return estimator