def _get_dataloader(net, test_dataset, data_shape, batch_size,
num_workers, num_devices, args):
"""Get dataloader."""
if args.meta_arch == 'yolo3':
width, height = data_shape, data_shape
val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
test_loader = gluon.data.DataLoader(
test_dataset.transform(
YOLO3DefaultValTransform(width, height)),
batch_size,
False,
batchify_fn=val_batchify_fn,
last_batch='keep',
num_workers=num_workers)
return test_loader
elif args.meta_arch == 'faster_rcnn':
"""Get faster rcnn dataloader."""
test_bfn = Tuple(*[Append() for _ in range(3)])
short = net.short[-1] if isinstance(net.short,
(tuple,
list)) else net.short
# validation use 1 sample per device
test_loader = gluon.data.DataLoader(
test_dataset.transform(
FasterRCNNDefaultValTransform(short, net.max_size)),
num_devices,
False,
batchify_fn=test_bfn,
last_batch='keep',
num_workers=args.num_workers)
return test_loader
else:
raise NotImplementedError('%s not implemented.' %
args.meta_arch)
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, batch_size, num_workers):
"""Get dataloader."""
train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)])
val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)])
dtype = 'float16' if cfg.GENERAL.FP16 else 'float32'
train_dataset = train_dataset.transform(
FasterRCNNDefaultTrainTransform(net.short,
net.max_size,
net,
dtype=dtype))
val_dataset = val_dataset.transform(
FasterRCNNDefaultValTransform(net.short, net.max_size, dtype=dtype))
train_loader = mx.gluon.data.DataLoader(train_dataset,
batch_size,
True,
batchify_fn=train_bfn,
last_batch='rollover',
num_workers=num_workers,
thread_pool=True)
val_loader = mx.gluon.data.DataLoader(val_dataset,
batch_size,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=num_workers,
thread_pool=True)
return train_loader, val_loader
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(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)
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_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_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_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 __init__(self,
root=os.path.join('~', '.mxnet', 'datasets', 'visualgenome'),
split='train'):
super(VGRelation, self).__init__(root)
self._root = os.path.expanduser(root)
self._img_path = os.path.join(self._root, 'VG_100K', '{}')
if split == 'train':
self._dict_path = os.path.join(self._root,
'rel_annotations_train.json')
elif split == 'val':
self._dict_path = os.path.join(self._root,
'rel_annotations_val.json')
else:
raise NotImplementedError
with open(self._dict_path) as f:
tmp = f.read()
self._dict = json.loads(tmp)
self._predicates_path = os.path.join(self._root, 'predicates.json')
with open(self._predicates_path, 'r') as f:
tmp = f.read()
self.rel_classes = json.loads(tmp)
self.num_rel_classes = len(self.rel_classes) + 1
self._objects_path = os.path.join(self._root, 'objects.json')
with open(self._objects_path, 'r') as f:
tmp = f.read()
self.obj_classes = json.loads(tmp)
self.num_obj_classes = len(self.obj_classes)
if split == 'val':
self.img_transform = FasterRCNNDefaultValTransform(short=600,
max_size=1000)
else:
self.img_transform = FasterRCNNDefaultTrainTransform(short=600,
max_size=1000)
self.split = split
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_voc_iterator(rank, num_workers, net, num_shards):
data_dir = "data-%d" % rank
try:
s3_client = boto3.client('s3')
for file in [
'VOCtrainval_06-Nov-2007.tar', 'VOCtest_06-Nov-2007.tar',
'VOCtrainval_11-May-2012.tar'
]:
s3_client.download_file(args.s3bucket, f'voc_tars/{file}',
f'/opt/ml/code/{file}')
with tarfile.open(filename) as tar:
tar.extractall(path=path)
except:
print('downloading from source')
download_voc(data_dir)
input_shape = (1, 256, 256, 3)
batch_size = args.batch_size
# might want to replace with mx.io.ImageDetRecordIter, this means you need data in RecordIO format
# train_iter = mx.io.MNISTIter(
# image="%s/train-images-idx3-ubyte" % data_dir,
# label="%s/train-labels-idx1-ubyte" % data_dir,
# input_shape=input_shape,
# batch_size=batch_size,
# shuffle=True,
# flat=False,
# num_parts=hvd.size(),
# part_index=hvd.rank()
# )
train_dataset = gdata.VOCDetection(
root=f'/opt/ml/code/data-{rank}/VOCdevkit/',
splits=[(2007, 'trainval'), (2012, 'trainval')])
val_dataset = gdata.VOCDetection(
root=f'/opt/ml/code/data-{rank}/VOCdevkit/',
splits=[(2007, 'test')])
val_metric = VOC07MApMetric(iou_thresh=0.5,
class_names=val_dataset.classes)
im_aspect_ratio = [1.] * len(train_dataset)
train_bfn = FasterRCNNTrainBatchify(net)
train_sampler = gluoncv.nn.sampler.SplitSortedBucketSampler(
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)
# had issue with multi_stage=True
train_iter = mx.gluon.data.DataLoader(train_dataset.transform(
FasterRCNNDefaultTrainTransform(net.short,
net.max_size,
net,
ashape=net.ashape,
multi_stage=False)),
batch_sampler=train_sampler,
batchify_fn=train_bfn,
num_workers=num_workers)
val_bfn = Tuple(*[Append() for _ in range(3)])
short = net.short[-1] if isinstance(net.short,
(tuple, list)) else net.short
# validation use 1 sample per device
val_iter = mx.gluon.data.DataLoader(val_dataset.transform(
FasterRCNNDefaultValTransform(short, net.max_size)),
num_shards,
False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=num_workers)
return train_iter, val_iter
num_parts=1,
part_index=0,
shuffle=True)
train_loader = mx.gluon.data.DataLoader(
train_dataset.transform(
FasterRCNNDefaultTrainTransform(
net.short,
net.max_size,
net,
ashape=net.ashape,
multi_stage=cfg["model"]["use_fpn"])),
batch_sampler=train_sampler,
batchify_fn=train_bfn,
num_workers=cfg["dataset"]["num_workers"],
)
val_bfn = Tuple(*[Append() for _ in range(3)])
# short是列表或元组，是为了在训练过程中增加图片大小抖动这一数据增强方式
short = net.short[-1] if isinstance(net.short,
(tuple, list)) else net.short
val_loader = mx.gluon.data.DataLoader(
val_dataset.transform(
FasterRCNNDefaultValTransform(short, net.max_size)),
batch_size=len(ctx), # 每张卡一张图片
shuffle=False,
batchify_fn=val_bfn,
last_batch='keep',
num_workers=cfg["dataset"]["num_workers"])
train(net, train_loader, val_loader, eval_metric, ctx, cfg)