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