def yolo3_darknet53_coco(pretrained_base=True, pretrained=False, num_sync_bn_devices=-1, **kwargs): """YOLO3 multi-scale with darknet53 base network on COCO dataset. Parameters ---------- pretrained_base : boolean Whether fetch and load pretrained weights for base network. pretrained : boolean Whether fetch and load pretrained weights for the entire network. num_sync_bn_devices : int, default is -1 Number of devices for training. If `num_sync_bn_devices < 2`, SyncBatchNorm is disabled. Returns ------- mxnet.gluon.HybridBlock Fully hybrid yolo3 network. """ from gluoncv.data import COCODetection pretrained_base = False if pretrained else pretrained_base base_net = darknet53(pretrained=pretrained_base, num_sync_bn_devices=num_sync_bn_devices) stages = [base_net.features[:15], base_net.features[15:24], base_net.features[24:]] anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] strides = [8, 16, 32] classes = COCODetection.CLASSES return get_yolov3( 'darknet53', stages, [512, 256, 128], anchors, strides, classes, 'coco', pretrained=pretrained, num_sync_bn_devices=num_sync_bn_devices, **kwargs)
def main(_argv): # get dataset dataset = get_dataset(FLAGS.dataset) # fix for tiny datasets of 1 or few elements batch_size = FLAGS.batch_size if len(dataset) < batch_size: batch_size = len(dataset) # handle gpu usage gpus = FLAGS.gpus if batch_size < len(gpus): gpus = [int(gpus[0])] # contexts ctx = [mx.gpu(int(i)) for i in gpus] ctx = ctx if ctx else [mx.cpu()] # dataloader loader = get_dataloader(dataset, batch_size) # setup network if FLAGS.network == 'darknet53': net = darknet53(root='models', pretrained=True, norm_layer=BatchNorm, norm_kwargs=None) elif FLAGS.network == 'mobilenet1.0': net = get_mobilenet(root='models', multiplier=1, pretrained=True, norm_layer=BatchNorm, norm_kwargs=None) else: raise NotImplementedError( 'Backbone CNN model {} not implemented.'.format(FLAGS.network)) # organise the save directories for the results if FLAGS.dataset in ['voc', 'coco', 'det', 'vid']: save_dir = os.path.join('models', FLAGS.network, FLAGS.save_dir, FLAGS.dataset) else: save_dir = os.path.join('models', FLAGS.network, FLAGS.save_dir) os.makedirs(save_dir, exist_ok=True) # attempt to load predictions extract(save_dir, net, dataset, loader, ctx, net_name=FLAGS.network)
def yolo3_darknet53_custom(classes, transfer=None, pretrained_base=True, pretrained=False, num_sync_bn_devices=-1, **kwargs): """YOLO3 multi-scale with darknet53 base network on custom dataset. Parameters ---------- classes : iterable of str Names of custom foreground classes. `len(classes)` is the number of foreground classes. transfer : str or None If not `None`, will try to reuse pre-trained weights from SSD networks trained on other datasets. pretrained_base : boolean Whether fetch and load pretrained weights for base network. num_sync_bn_devices : int, default is -1 Number of devices for training. If `num_sync_bn_devices < 2`, SyncBatchNorm is disabled. Returns ------- mxnet.gluon.HybridBlock Fully hybrid yolo3 network. """ if transfer is None: base_net = darknet53(pretrained=pretrained_base, num_sync_bn_devices=num_sync_bn_devices) stages = [base_net.features[:15], base_net.features[15:24], base_net.features[24:]] anchors = [ [10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] strides = [8, 16, 32] net = get_yolov3( 'darknet53', stages, [512, 256, 128], anchors, strides, classes, 'coco', pretrained=pretrained, num_sync_bn_devices=num_sync_bn_devices, **kwargs) else: from gluoncv.model_zoo import get_model net = get_model('yolo3_darknet53_' + str(transfer), pretrained=True, **kwargs) net.reset_class(classes) return net
def yolo3_darknet53(classes, dataset_name, transfer=None, pretrained_base=True, pretrained=False, norm_layer=BatchNorm, norm_kwargs=None, **kwargs): """YOLO3 multi-scale with darknet53 base network on any dataset. Modified from: https://github.com/dmlc/gluon-cv/blob/0dbd05c5eb8537c25b64f0e87c09be979303abf2/gluoncv/model_zoo/yolo/yolo3.py Parameters ---------- classes : iterable of str Names of custom foreground classes. `len(classes)` is the number of foreground classes. dataset_name : str The name of the dataset, used for model save name transfer : str or None If not `None`, will try to reuse pre-trained weights from yolo networks trained on other datasets. pretrained_base : boolean Whether fetch and load pretrained weights for base network. norm_layer : object Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`) Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`. norm_kwargs : dict Additional `norm_layer` arguments, for example `num_devices=4` for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`. Returns ------- mxnet.gluon.HybridBlock Fully hybrid yolo3 network. """ if pretrained: warnings.warn( "Custom models don't provide `pretrained` weights, ignored.") if transfer is None: base_net = darknet53(pretrained=pretrained_base, norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) stages = [ base_net.features[:15], base_net.features[15:24], base_net.features[24:] ] anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] strides = [8, 16, 32] net = get_yolov3('darknet53', stages, [512, 256, 128], anchors, strides, classes, dataset_name, norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) else: net = get_model('yolo3_darknet53_' + str(transfer), pretrained=True, **kwargs) reuse_classes = [x for x in classes if x in net.classes] net.reset_class(classes, reuse_weights=reuse_classes) return net
def yolo3_darknet53_custom(classes, transfer=None, pretrained_base=True, pretrained=False, norm_layer=BatchNorm, norm_kwargs=None, **kwargs): """YOLO3 multi-scale with darknet53 base network on custom dataset. Parameters ---------- classes : iterable of str Names of custom foreground classes. `len(classes)` is the number of foreground classes. transfer : str or None If not `None`, will try to reuse pre-trained weights from yolo networks trained on other datasets. pretrained_base : boolean Whether fetch and load pretrained weights for base network. norm_layer : object Normalization layer used (default: :class:`mxnet.gluon.nn.BatchNorm`) Can be :class:`mxnet.gluon.nn.BatchNorm` or :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`. norm_kwargs : dict Additional `norm_layer` arguments, for example `num_devices=4` for :class:`mxnet.gluon.contrib.nn.SyncBatchNorm`. Returns ------- mxnet.gluon.HybridBlock Fully hybrid yolo3 network. """ if pretrained: warnings.warn( "Custom models don't provide `pretrained` weights, ignored.") if transfer is None: base_net = darknet53(pretrained=pretrained_base, norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) stages = [ base_net.features[:15], base_net.features[15:24], base_net.features[24:] ] #256, #512 #1024 anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] anchors_half = [[5, 6, 8, 15, 16, 11], [15, 30, 31, 22, 29, 59], [58, 45, 78, 99, 186, 163]] anchors_fourth = [[2, 3, 4, 7, 8, 5], [7, 15, 15, 11, 14, 29], [29, 22, 39, 49, 93, 81]] strides = [8, 16, 32] net = get_yolov3('darknet53', stages, [512, 256, 128], anchors, strides, classes, '', norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs) else: from gluoncv.model_zoo import get_model net = get_model('yolo3_darknet53_' + str(transfer), pretrained=True, **kwargs) reuse_classes = [x for x in classes if x in net.classes] net.reset_class(classes, reuse_weights=reuse_classes) return net