def get_dataloader(net, train_dataset, val_dataset, batch_size, num_workers):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(FasterRCNNDefaultTrainTransform(net.short, net.max_size, net)),
batch_size, True, batchify_fn=train_bfn, last_batch='rollover', num_workers=num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(FasterRCNNDefaultValTransform(net.short, net.max_size)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
return train_loader, val_loader
def test_transforms_presets_rcnn():
im_fname = gcv.utils.download(
'https://github.com/dmlc/web-data/blob/master/' +
'gluoncv/detection/biking.jpg?raw=true',
path='biking.jpg')
x, orig_img = rcnn.load_test(im_fname, short=600, max_size=1000)
x1, orig_img1 = rcnn.transform_test(mx.image.imread(im_fname),
short=600,
max_size=1000)
np.testing.assert_allclose(x.asnumpy(), x1.asnumpy())
np.testing.assert_allclose(orig_img, orig_img1)
if not osp.isdir(osp.expanduser('~/.mxnet/datasets/voc')):
return
train_dataset = gcv.data.VOCDetection(splits=((2007, 'trainval'),
(2012, 'trainval')))
val_dataset = gcv.data.VOCDetection(splits=[(2007, 'test')])
width, height = (512, 512)
net = gcv.model_zoo.get_model('faster_rcnn_resnet50_v1b_voc',
pretrained=False,
pretrained_base=False)
net.initialize()
num_workers = 0
short, max_size = 600, 1000
batch_size = 4
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
train_loader = mx.gluon.data.DataLoader(train_dataset.transform(
rcnn.FasterRCNNDefaultTrainTransform(short, max_size, net)),
batch_size,
True,
batchify_fn=train_bfn,
last_batch='rollover',
num_workers=num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_loader = mx.gluon.data.DataLoader(val_dataset.transform(
rcnn.FasterRCNNDefaultValTransform(short, max_size)),
batch_size,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=num_workers)
train_loader2 = gluon.data.DataLoader(
train_dataset.transform(
rcnn.FasterRCNNDefaultTrainTransform(short, max_size)),
batch_size,
True,
batchify_fn=batchify.Tuple(*[batchify.Append() for _ in range(2)]),
last_batch='rollover',
num_workers=num_workers)
for loader in [train_loader, val_loader, train_loader2]:
for i, batch in enumerate(loader):
if i > 1:
break
pass
def get_dataloader(net, train_dataset, val_dataset, train_transform, val_transform, batch_size,
num_workers, multi_stage):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(7)])
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(train_transform(net, ashape=net.ashape, multi_stage=multi_stage)),
batch_size, True, batchify_fn=train_bfn, last_batch='rollover', num_workers=num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
val_loader = mx.gluon.data.DataLoader(val_dataset, batch_size, False, batchify_fn=val_bfn,
last_batch='keep', num_workers=num_workers)
return train_loader, val_loader
def get_dataloader(net, train_dataset, val_dataset, train_transform, val_transform, batch_size,
num_workers):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Stack() for _ in range(6)])
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(train_transform(
net.short, net.base_stride, net.valid_range)),
batch_size,shuffle=True, batchify_fn=train_bfn, last_batch='rollover',
num_workers=num_workers)
val_bfn = batchify.Tuple(*[batchify.Stack() for _ in range(2)])
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(val_transform(net.short, net.base_stride)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
return train_loader, val_loader
def get_dataloader(net, train_dataset, val_dataset, train_transform, val_transform, batch_size,
num_workers, multi_stage):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(6)])
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(train_transform(net.short, net.max_size, net, ashape=net.ashape,
multi_stage=multi_stage)),
batch_size, True, batchify_fn=train_bfn, last_batch='rollover', num_workers=num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
short = net.short[-1] if isinstance(net.short, (tuple, list)) else net.short
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(val_transform(short, net.max_size)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
return train_loader, val_loader
def get_loader(net, train_dataset, val_dataset, batch_size, args):
train_transform = FasterRCNNDefaultTrainTransform(net.short, net.max_size, net,
ashape=net.ashape, multi_stage=args.use_fpn)
# return images, labels, rpn_cls_targets, rpn_box_targets, rpn_box_masks
train_batchify_fn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
train_loader = mx.gluon.data.DataLoader(train_dataset.transform(train_transform),
batch_size, shuffle=True, batchify_fn=train_batchify_fn, last_batch='rollover')
short = net.short[-1] if isinstance(net.short, (tuple, list)) else net.short
val_transform = FasterRCNNDefaultValTransform(short, net.max_size)
# return to x, y, im_scale
val_batchify_fn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_loader = mx.gluon.data.DataLoader(val_dataset.transform(val_transform), batch_size, shuffle=False,
batchify_fn=val_batchify_fn, last_batch='keep')
"""
train_loader：
每个batch为（[data1, data2,...], [label, label2,...], [rpn_cls_targets1, rpn_cls_targets2, ...],
[rpn_box_targets1, rpn_box_targets2, ...], [rpn_box_masks1, rpn_box_masks2, ...]）
for data, label, rpn_cls_targets, rpn_box_targets, rpn_box_masks in zip(*train_loader):
rpn_cls_targets: (1, N) rpn_box_targets: (1, N, 4) rpn_box_masks: (1, N, 4)
val_loader:
每个batch为（[data1, data2,...], [label, label2,...], [im_scale1, im_scale2, ...]）
for data, label, im_scale in zip(*val_loader):
im_scale: (1, 1)
但是还没搞清楚im_scale是什么东西 T_T
cls_preds, box_preds, roi, samples, matches, rpn_score, rpn_box, anchors = net(data, gt_box)
cls_targets, box_targets, box_masks = net.target_generator(roi, samples, matches, gt_label, gt_box)
data: (1, 3, h, w)
label: (1, num_obj, 6)
# rpn
roi: (1, 128, 4)
samples: (1, 128)
matches: (1, 128)
rpn_cls_targets: (1, N) rpn_score: (1, N, 1)
rpn_box_targets: (1, N, 4) rpn_box: (1, N, 4)
rpn_box_masks: (1, N, 4)
# rcnn
cls_targets: (1, 128) cls_preds: (1, 128, num_cls + 1)
box_targets: (1, 128, num_cls, 4) box_preds: (1, 128, num_cls, 4)
rcnn box mask: (1, 128, num_cls, 4)
"""
return train_loader, val_loader
def get_dataloader(net, train_dataset, val_dataset, train_transform, val_transform, batch_size,
args):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(6)])
train_sampler = gcv.nn.sampler.SplitSampler(len(train_dataset), hvd.size(), hvd.rank()) if args.horovod else None
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(train_transform(net.short, net.max_size, net, ashape=net.ashape,
multi_stage=args.use_fpn)),
batch_size, train_sampler is None, sampler=train_sampler, batchify_fn=train_bfn,
last_batch='rollover', num_workers=args.num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
short = net.short[-1] if isinstance(net.short, (tuple, list)) else net.short
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(val_transform(short, net.max_size)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=args.num_workers)
return train_loader, val_loader
def get_data_loader(net, train_dataset, val_dataset, batch_size=1):
"""
get PASCAL-VOC train&val dataloader
:return:
"""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(FasterRCNNDefaultTrainTransform(net.short, net.max_size, net)),
batch_size=1, shuffle=True, batchify_fn=train_bfn, last_batch='rollover')
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(FasterRCNNDefaultValTransform(net.short, net.max_size)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep')
return (train_loader, val_loader)
def get_dataloader(net, train_dataset, val_dataset, train_transform,
val_transform, batch_size, num_shards_per_process, args):
"""Get dataloader."""
train_bfn = batchify.MaskRCNNTrainBatchify(net, num_shards_per_process)
train_sampler = \
gcv.nn.sampler.SplitSortedBucketSampler(train_dataset.get_im_aspect_ratio(),
batch_size,
num_parts=hvd.size() if args.horovod else 1,
part_index=hvd.rank() if args.horovod else 0,
shuffle=True)
train_loader = mx.gluon.data.DataLoader(train_dataset.transform(
train_transform(net.short,
net.max_size,
net,
ashape=net.ashape,
multi_stage=args.use_fpn)),
batch_sampler=train_sampler,
batchify_fn=train_bfn,
num_workers=args.num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
short = net.short[-1] if isinstance(net.short,
(tuple, list)) else net.short
# validation use 1 sample per device
val_loader = mx.gluon.data.DataLoader(val_dataset.transform(
val_transform(short, net.max_size)),
num_shards_per_process,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=args.num_workers)
return train_loader, val_loader
def get_dataloader(net, val_dataset, batch_size, num_workers):
"""Get dataloader."""
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(MaskRCNNDefaultValTransform(net.short, net.max_size)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
return val_loader
def get_dataloader(net, val_dataset, val_transform, batch_size, num_workers):
"""Get dataloader."""
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(4)])
short = net.short[-1] if isinstance(net.short, (tuple, list)) else net.short
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(val_transform(short, net.max_size, net.base_stride)),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
return val_loader
def get_dataloader(net, val_dataset, batch_size, num_workers):
"""Get dataloader."""
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
val_loader = mx.gluon.data.DataLoader(val_dataset,
batch_size,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=num_workers)
return val_loader
def get_dataloader(net, val_dataset, data_shape, batch_size, num_workers, args):
"""Get dataloader."""
width, height = data_shape, data_shape
# val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
# Copied from eval_mask_rcnn.py
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
val_loader = gluon.data.DataLoader(
val_dataset.transform(YOLO3UsdSegCocoValTransform(width, height, 50, 'coco')),
batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers)
return val_loader
def get_frcnn_data_loader(val_dataset, net):
""" load data in batches for frcnn model """
batch_size = BATCH_SIZE
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_loader = mx.gluon.data.DataLoader(val_dataset.transform(
FasterRCNNDefaultValTransform(net.short, net.max_size)),
batch_size,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=0)
return val_loader
def DogDataLoader(net,
root='./stanford_dog_dataset',
preload_label=True,
batch_size=1,
shuffle=True,
num_workers=0):
# dataset
train_dataset = DogDetection(root=root,
splits='train',
preload_label=preload_label)
val_dataset = DogDetection(root=root,
splits='test',
preload_label=preload_label)
val_metric = VOC07MApMetric(iou_thresh=0.5,
class_names=val_dataset.classes)
train_batchify = batchify.Tuple(*[batchify.Append() for _ in range(5)])
val_batchify = batchify.Tuple(*[batchify.Append() for _ in range(3)])
train_dataloader = gluon.data.DataLoader(train_dataset.transform(
FasterRCNNDefaultTrainTransform(net.short, net.max_size, net)),
batch_size=batch_size,
shuffle=True,
batchify_fn=train_batchify,
last_batch='rollover',
num_workers=num_workers)
val_dataloader = gluon.data.DataLoader(val_dataset.transform(
FasterRCNNDefaultValTransform(net.short, net.max_size)),
batch_size=batch_size,
shuffle=False,
batchify_fn=val_batchify,
last_batch='keep',
num_workers=num_workers)
return train_dataloader, val_dataloader, val_metric
def get_val_frcnn_dataloader(val_dataset,
batch_size,
num_workers,
short=600,
max_size=800):
"""Get dataloader."""
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_loader = mx.gluon.data.DataLoader(val_dataset.transform(
FasterRCNNDefaultValTransform(short, max_size)),
batch_size,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=num_workers)
return val_loader
def get_dataloader(net,
train_dataset,
batch_size,
num_workers,
short=600,
max_size=1000):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
train_loader = mx.gluon.data.DataLoader(train_dataset.transform(
FasterRCNNDefaultTrainTransform(short, max_size, net)),
batch_size,
True,
batchify_fn=train_bfn,
last_batch='rollover',
num_workers=num_workers)
return train_loader
def train(net, async_net, ctx, args):
"""Training pipeline"""
net.collect_params().reset_ctx(ctx)
if args.no_wd:
for k, v in net.collect_params(".*beta|.*gamma|.*bias").items():
v.wd_mult = 0.0
if args.label_smooth:
net._target_generator._label_smooth = True
if args.lr_decay_period > 0:
lr_decay_epoch = list(
range(args.lr_decay_period, args.epochs, args.lr_decay_period))
else:
lr_decay_epoch = [int(i) for i in args.lr_decay_epoch.split(',')]
lr_scheduler = LRSequential([
LRScheduler("linear",
base_lr=0,
target_lr=args.lr,
nepochs=args.warmup_epochs,
iters_per_epoch=args.batch_size),
LRScheduler(args.lr_mode,
base_lr=args.lr,
nepochs=args.epochs - args.warmup_epochs,
iters_per_epoch=args.batch_size,
step_epoch=lr_decay_epoch,
step_factor=args.lr_decay,
power=2),
])
if (args.optimizer == "sgd"):
trainer = gluon.Trainer(net.collect_params(),
args.optimizer, {
"wd": args.wd,
"momentum": args.momentum,
"lr_scheduler": lr_scheduler
},
kvstore="local")
elif (args.optimizer == "adam"):
trainer = gluon.Trainer(net.collect_params(),
args.optimizer, {"lr_scheduler": lr_scheduler},
kvstore="local")
else:
trainer = gluon.Trainer(net.collect_params(),
args.optimizer,
kvstore="local")
# targets
#sigmoid_ce = gluon.loss.SigmoidBinaryCrossEntropyLoss(from_sigmoid=False)
#l1_loss = gluon.loss.L1Loss()
# Intermediate Metrics:
train_metrics = (
mx.metric.Loss("ObjLoss"),
mx.metric.Loss("BoxCenterLoss"),
mx.metric.Loss("BoxScaleLoss"),
mx.metric.Loss("ClassLoss"),
mx.metric.Loss("TotalLoss"),
)
train_metric_ixs = range(len(train_metrics))
target_metric_ix = -1 # Train towards TotalLoss (the last one)
# Evaluation Metrics:
val_metric = VOC07MApMetric(iou_thresh=0.5)
# Data transformations:
train_dataset = gluon_pipe_mode.AugmentedManifestDetection(
args.train,
length=args.num_samples_train,
)
train_batchify_fn = batchify.Tuple(
*([batchify.Stack() for _ in range(6)] +
[batchify.Pad(axis=0, pad_val=-1) for _ in range(1)]))
if args.no_random_shape:
logger.debug("Creating train DataLoader without random transform")
train_transforms = YOLO3DefaultTrainTransform(args.data_shape,
args.data_shape,
net=async_net,
mixup=args.mixup)
train_dataloader = gluon.data.DataLoader(
train_dataset.transform(train_transforms),
batch_size=args.batch_size,
batchify_fn=train_batchify_fn,
last_batch="discard",
num_workers=args.num_workers,
shuffle=
False, # Note that shuffle *cannot* be used with AugmentedManifestDetection
)
else:
logger.debug("Creating train DataLoader with random transform")
train_transforms = [
YOLO3DefaultTrainTransform(x * 32,
x * 32,
net=async_net,
mixup=args.mixup)
for x in range(10, 20)
]
train_dataloader = RandomTransformDataLoader(
train_transforms,
train_dataset,
interval=10,
batch_size=args.batch_size,
batchify_fn=train_batchify_fn,
last_batch="discard",
num_workers=args.num_workers,
shuffle=
False, # Note that shuffle *cannot* be used with AugmentedManifestDetection
)
validation_dataset = None
validation_dataloader = None
if args.validation:
validation_dataset = gluon_pipe_mode.AugmentedManifestDetection(
args.validation,
length=args.num_samples_validation,
)
validation_dataloader = gluon.data.DataLoader(
validation_dataset.transform(
YOLO3DefaultValTransform(args.data_shape, args.data_shape), ),
args.batch_size,
shuffle=False,
batchify_fn=batchify.Tuple(batchify.Stack(),
batchify.Pad(pad_val=-1)),
last_batch="keep",
num_workers=args.num_workers,
)
# Prepare the inference-time configuration for our model's setup:
# (This will be saved alongside our network structure/params)
inference_config = config.InferenceConfig(image_size=args.data_shape)
logger.info(args)
logger.info(f"Start training from [Epoch {args.start_epoch}]")
prev_best_score = float("-inf")
best_epoch = args.start_epoch
logger.info("Sleeping for 3s in case training data file not yet ready")
time.sleep(3)
for epoch in range(args.start_epoch, args.start_epoch + args.epochs):
# if args.mixup:
# # TODO(zhreshold): more elegant way to control mixup during runtime
# try:
# train_data._dataset.set_mixup(np.random.beta, 1.5, 1.5)
# except AttributeError:
# train_data._dataset._data.set_mixup(np.random.beta, 1.5, 1.5)
# if epoch >= args.epochs - args.no_mixup_epochs:
# try:
# train_data._dataset.set_mixup(None)
# except AttributeError:
# train_data._dataset._data.set_mixup(None)
tic = time.time()
btic = time.time()
mx.nd.waitall()
net.hybridize()
logger.debug(
f"Input data dir contents: {os.listdir('/opt/ml/input/data/')}")
for i, batch in enumerate(train_dataloader):
logger.debug(f"Epoch {epoch}, minibatch {i}")
batch_size = batch[0].shape[0]
data = gluon.utils.split_and_load(batch[0],
ctx_list=ctx,
batch_axis=0,
even_split=False)
# objectness, center_targets, scale_targets, weights, class_targets
fixed_targets = [
gluon.utils.split_and_load(batch[it],
ctx_list=ctx,
batch_axis=0,
even_split=False)
for it in range(1, 6)
]
gt_boxes = gluon.utils.split_and_load(batch[6],
ctx_list=ctx,
batch_axis=0,
even_split=False)
loss_trackers = tuple([] for metric in train_metrics)
with autograd.record():
for ix, x in enumerate(data):
losses_raw = net(x, gt_boxes[ix],
*[ft[ix] for ft in fixed_targets])
# net outputs: [obj_loss, center_loss, scale_loss, cls_loss]
# Each a mx.ndarray 1xbatch_size. This is the same order as our
# train_metrics, so we just need to add a total vector:
total_loss = sum(losses_raw)
losses = losses_raw + [total_loss]
# If any sample's total loss is non-finite, sum will be:
if not isfinite(sum(total_loss)):
logger.error(
f"[Epoch {epoch}][Minibatch {i}] got non-finite losses: {losses_raw}"
)
# TODO: Terminate training if losses or gradient go infinite?
for ix in train_metric_ixs:
loss_trackers[ix].append(losses[ix])
autograd.backward(loss_trackers[target_metric_ix])
trainer.step(batch_size)
for ix in train_metric_ixs:
train_metrics[ix].update(0, loss_trackers[ix])
if args.log_interval and not (i + 1) % args.log_interval:
train_metrics_current = map(lambda metric: metric.get(),
train_metrics)
metrics_msg = "; ".join([
f"{name}={val:.3f}" for name, val in train_metrics_current
])
logger.info(
f"[Epoch {epoch}][Minibatch {i}] LR={trainer.learning_rate:.2E}; "
f"Speed={batch_size/(time.time()-btic):.3f} samples/sec; {metrics_msg};"
)
btic = time.time()
train_metrics_current = map(lambda metric: metric.get(), train_metrics)
metrics_msg = "; ".join(
[f"{name}={val:.3f}" for name, val in train_metrics_current])
logger.info(
f"[Epoch {epoch}] TrainingCost={time.time()-tic:.3f}; {metrics_msg};"
)
if not (epoch + 1) % args.val_interval:
logger.info(f"Validating [Epoch {epoch}]")
metric_names, metric_values = validate(
net, validation_dataloader, epoch, ctx,
VOC07MApMetric(iou_thresh=0.5), args)
if isinstance(metric_names, list):
val_msg = "; ".join(
[f"{k}={v}" for k, v in zip(metric_names, metric_values)])
current_score = float(metric_values[-1])
else:
val_msg = f"{metric_names}={metric_values}"
current_score = metric_values
logger.info(f"[Epoch {epoch}] Validation: {val_msg};")
else:
current_score = float("-inf")
save_progress(
net,
inference_config,
current_score,
prev_best_score,
args.model_dir,
epoch,
args.checkpoint_interval,
args.checkpoint_dir,
)
if current_score > prev_best_score:
prev_best_score = current_score
best_epoch = epoch
if (args.early_stopping and epoch >= args.early_stopping_min_epochs
and (epoch - best_epoch) >= args.early_stopping_patience):
logger.info(
f"[Epoch {epoch}] No improvement since epoch {best_epoch}: Stopping early"
)
break
def speed_test():
args = parse_args()
# context list
ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()]
ctx = [mx.cpu()] if not ctx else ctx
# Get net
net = gcv.model_zoo.get_model(args.network,
pretrained=False,
pretrained_base=False)
net.load_parameters(args.pretrained)
net.set_nms(0.45, 200)
net.collect_params().reset_ctx(ctx=ctx)
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
# Dataset
val_dataset = COCOInstance(root='/home/tutian/dataset', skip_empty=False)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(2)])
# val_bfn = Tuple(Stack(), Pad(pad_val=-1))
val_loader = gluon.data.DataLoader(val_dataset.transform(
YOLO3UsdSegCocoValTransform(416, 416, 50, 'coco')),
1,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=1)
# Some preparation
total_time_net = 0
total_time_post = 0
total_time_cpu = 0
total_time_predot = 0
total_time_dot = 0
total_time_genmask = 0
total_time_calculate = 0
colors = {
i: plt.get_cmap('hsv')(i / len(CLASSES))
for i in range(len(CLASSES))
}
img_ids = sorted(val_dataset.coco.getImgIds())
mx.nd.waitall()
net.hybridize()
save_images = False
# print(image_list_batch)
with tqdm(total=5000) as pbar:
for ibt, batch in enumerate(val_loader):
batch = split_and_load(batch, ctx_list=ctx)
for x, im_info in zip(*batch):
# get prediction results
t1 = time.time()
ids, scores, bboxes, coefs = net(x)
t_c0 = time.time()
mx.nd.waitall()
t_c1 = time.time()
t2 = time.time()
# Post process
t_cpu0 = time.time()
bboxes = bboxes.asnumpy()[0]
ids = ids.asnumpy()[0]
scores = scores.asnumpy()[0]
coefs = coefs.asnumpy()[0]
im_info = im_info.asnumpy()[0]
t_cpu1 = time.time()
total_time_cpu += (t_cpu1 - t_cpu0)
t_cpu0 = time.time()
im_height, im_width = [int(i) for i in im_info]
valid = np.where(((ids >= 0) & (scores >= 0.45)))[0]
ids = ids[valid]
scores = scores[valid]
bboxes = bboxes[valid]
coefs = coefs[valid]
coefs = coefs * sqrt_var + x_mean
t_cpu1 = time.time()
total_time_predot += (t_cpu1 - t_cpu0)
t_cpu0 = time.time()
masks = np.dot(coefs, bases)
t_cpu1 = time.time()
total_time_dot += (t_cpu1 - t_cpu0)
t_cpu0 = time.time()
bboxes, masks = generate_bbox_mask(coefs, bboxes, im_height,
im_width)
t_cpu1 = time.time()
total_time_genmask += (t_cpu1 - t_cpu0)
t3 = time.time()
if ibt >= 500:
total_time_net += (t2 - t1)
total_time_calculate += (t_c1 - t_c0)
total_time_post += (t3 - t2)
# Save the masks
if save_images:
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(1, 1, 1)
ax.imshow(
Image.open('/disk1/data/coco/val2017/' +
str(img_ids[ibt]).zfill(12) + '.jpg'))
for bbox, mask, idt in zip(bboxes, masks, ids):
idt = int(idt)
# bbox
rect = plt.Rectangle((bbox[0], bbox[1]),
bbox[2],
bbox[3],
fill=False,
edgecolor=colors[idt],
linewidth=2)
ax.add_patch(rect)
# mask
mask_channel = np.zeros((im_height, im_width, 4))
for i in range(4):
mask_channel[:, :, i] = mask * colors[idt][i]
ax.imshow(mask_channel, alpha=0.5)
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.subplots_adjust(top=0.995,
bottom=0.005,
right=0.995,
left=0.005,
hspace=0,
wspace=0)
plt.savefig(
os.path.join(args.save_dir,
str(img_ids[ibt]).zfill(12) + '.jpg'))
plt.close()
pbar.update(1)
print("network speed ", 4500 / total_time_net, "fps")
print("Sync speed ", 4500 / total_time_calculate, "fps")
print("post process speed ", 5000 / total_time_post, "fps")
print("GPU to CPU speed ", 5000 / total_time_cpu, "fps")
print("Pre-dot speed ", 5000 / total_time_predot, "fps")
print("np.dot speed ", 5000 / total_time_dot, "fps")
print("Gen-Mask speed ", 5000 / total_time_genmask, "fps")
print("total speed ", 5000 / (total_time_net + total_time_post),
"fps")
num_workers = 0
# init model
net = get_model("faster_rcnn_resnet50_v1b_voc",
pretrained=False,
pretrained_base=False)
net.load_parameters(model_path)
net.collect_params().reset_ctx(ctx)
# load val dataset
val_dataset = gdata.VOCDetection(root="/Users/rensike/Files/temp/voc_mini",
splits=[(0, 'val')])
eval_metric = VOCMApMetric(iou_thresh=0.5, class_names=val_dataset.classes)
# load val dataloader
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
val_data_loader = mx.gluon.data.DataLoader(val_dataset.transform(
FasterRCNNDefaultValTransform(net.short, net.max_size)),
1,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=num_workers)
# do evaluate
eval_metric.reset()
net.hybridize(static_alloc=True)
map_name, mean_ap = validate(net, val_data_loader, ctx, eval_metric)
val_msg = '\n'.join(['{}={}'.format(k, v) for k, v in zip(map_name, mean_ap)])
print('Validation: \n{}'.format(val_msg))
