def get_dataloader(net, train_dataset, val_dataset, short, max_size,
batch_size, num_workers):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
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)
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 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(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 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.Append() for _ in range(5)])
train_loader = mx.gluon.data.DataLoader(train_dataset.transform(
train_transform(net.short, net.max_size)),
batch_size,
True,
batchify_fn=train_bfn,
last_batch='rollover',
num_workers=num_workers)
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
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_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 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)])
val_loader = mx.gluon.data.DataLoader(val_dataset.transform(
val_transform(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 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_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 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=True)),
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 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))
