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))