def create_datasets_and_loaders(args, model_config):
input_config = resolve_input_config(args, model_config=model_config)
dataset_train, dataset_eval = create_dataset(args.dataset, args.root)
# setup labeler in loader/collate_fn if not enabled in the model bench
labeler = None
if not args.bench_labeler:
labeler = AnchorLabeler(Anchors.from_config(model_config),
model_config.num_classes,
match_threshold=0.5)
loader_train = create_loader(
dataset_train,
input_size=input_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
# color_jitter=args.color_jitter,
# auto_augment=args.aa,
interpolation=args.train_interpolation
or input_config['interpolation'],
fill_color=input_config['fill_color'],
mean=input_config['mean'],
std=input_config['std'],
num_workers=args.workers,
distributed=args.distributed,
pin_mem=args.pin_mem,
anchor_labeler=labeler,
)
if args.val_skip > 1:
dataset_eval = SkipSubset(dataset_eval, args.val_skip)
loader_eval = create_loader(
dataset_eval,
input_size=input_config['input_size'],
batch_size=args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=input_config['interpolation'],
fill_color=input_config['fill_color'],
mean=input_config['mean'],
std=input_config['std'],
num_workers=args.workers,
distributed=args.distributed,
pin_mem=args.pin_mem,
anchor_labeler=labeler,
)
evaluator = create_evaluator(args.dataset,
loader_eval.dataset,
distributed=args.distributed,
pred_yxyx=False)
return loader_train, loader_eval, evaluator
def validate(args):
setup_default_logging()
if args.amp:
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
args.pretrained = args.pretrained or not args.checkpoint # might as well try to validate something
args.prefetcher = not args.no_prefetcher
# create model
with set_layer_config(scriptable=args.torchscript):
bench = create_model(
args.model,
bench_task='predict',
num_classes=args.num_classes,
pretrained=args.pretrained,
redundant_bias=args.redundant_bias,
soft_nms=args.soft_nms,
checkpoint_path=args.checkpoint,
checkpoint_ema=args.use_ema,
)
model_config = bench.config
param_count = sum([m.numel() for m in bench.parameters()])
print('Model %s created, param count: %d' % (args.model, param_count))
bench = bench.cuda()
amp_autocast = suppress
if args.apex_amp:
bench = amp.initialize(bench, opt_level='O1')
print('Using NVIDIA APEX AMP. Validating in mixed precision.')
elif args.native_amp:
amp_autocast = torch.cuda.amp.autocast
print('Using native Torch AMP. Validating in mixed precision.')
else:
print('AMP not enabled. Validating in float32.')
if args.num_gpu > 1:
bench = torch.nn.DataParallel(bench,
device_ids=list(range(args.num_gpu)))
dataset = create_dataset(args.dataset, args.root, args.split)
input_config = resolve_input_config(args, model_config)
loader = create_loader(dataset,
input_size=input_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=input_config['interpolation'],
fill_color=input_config['fill_color'],
mean=input_config['mean'],
std=input_config['std'],
num_workers=args.workers,
pin_mem=args.pin_mem)
evaluator = create_evaluator(args.dataset, dataset, pred_yxyx=False)
bench.eval()
batch_time = AverageMeter()
end = time.time()
last_idx = len(loader) - 1
with torch.no_grad():
for i, (input, target) in enumerate(loader):
with amp_autocast():
output = bench(input, img_info=target)
evaluator.add_predictions(output, target)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.log_freq == 0 or i == last_idx:
print(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
.format(i,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg))
mean_ap = 0.
if dataset.parser.has_labels:
mean_ap = evaluator.evaluate()
else:
evaluator.save(args.results)
return mean_ap
def create_datasets_and_loaders(
args,
model_config,
transform_train_fn=None,
transform_eval_fn=None,
collate_fn=None,
):
""" Setup datasets, transforms, loaders, evaluator.
Args:
args: Command line args / config for training
model_config: Model specific configuration dict / struct
transform_train_fn: Override default image + annotation transforms (see note in loaders.py)
transform_eval_fn: Override default image + annotation transforms (see note in loaders.py)
collate_fn: Override default fast collate function
Returns:
Train loader, validation loader, evaluator
"""
input_config = resolve_input_config(args, model_config=model_config)
dataset_train, dataset_eval = create_dataset(args.dataset, args.root)
# setup labeler in loader/collate_fn if not enabled in the model bench
labeler = None
if not args.bench_labeler:
labeler = AnchorLabeler(Anchors.from_config(model_config),
model_config.num_classes,
match_threshold=0.5)
loader_train = create_loader(
dataset_train,
input_size=input_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
# color_jitter=args.color_jitter,
# auto_augment=args.aa,
interpolation=args.train_interpolation
or input_config['interpolation'],
fill_color=input_config['fill_color'],
mean=input_config['mean'],
std=input_config['std'],
num_workers=args.workers,
distributed=args.distributed,
pin_mem=args.pin_mem,
anchor_labeler=labeler,
transform_fn=transform_train_fn,
collate_fn=collate_fn,
)
if args.val_skip > 1:
dataset_eval = SkipSubset(dataset_eval, args.val_skip)
loader_eval = create_loader(
dataset_eval,
input_size=input_config['input_size'],
batch_size=args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=input_config['interpolation'],
fill_color=input_config['fill_color'],
mean=input_config['mean'],
std=input_config['std'],
num_workers=args.workers,
distributed=args.distributed,
pin_mem=args.pin_mem,
anchor_labeler=labeler,
transform_fn=transform_eval_fn,
collate_fn=collate_fn,
)
evaluator = create_evaluator(args.dataset,
loader_eval.dataset,
distributed=args.distributed,
pred_yxyx=False)
return loader_train, loader_eval, evaluator
def validate(args):
setup_default_logging()
if args.amp:
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
args.pretrained = args.pretrained or not args.checkpoint # might as well try to validate something
args.prefetcher = not args.no_prefetcher
# create model
with set_layer_config(scriptable=args.torchscript):
bench = create_model(
args.model,
bench_task='predict',
num_classes=args.num_classes,
pretrained=args.pretrained,
redundant_bias=args.redundant_bias,
soft_nms=args.soft_nms,
checkpoint_path=args.checkpoint,
checkpoint_ema=args.use_ema,
)
model_config = bench.config
param_count = sum([m.numel() for m in bench.parameters()])
print('Model %s created, param count: %d' % (args.model, param_count))
bench = bench.cuda()
amp_autocast = suppress
if args.apex_amp:
bench = amp.initialize(bench, opt_level='O1')
print('Using NVIDIA APEX AMP. Validating in mixed precision.')
elif args.native_amp:
amp_autocast = torch.cuda.amp.autocast
print('Using native Torch AMP. Validating in mixed precision.')
else:
print('AMP not enabled. Validating in float32.')
if args.num_gpu > 1:
bench = torch.nn.DataParallel(bench,
device_ids=list(range(args.num_gpu)))
dataset = create_dataset(args.dataset, args.root, args.split)
input_config = resolve_input_config(args, model_config)
loader = create_loader(dataset,
input_size=input_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=input_config['interpolation'],
fill_color=input_config['fill_color'],
mean=input_config['mean'],
std=input_config['std'],
num_workers=args.workers,
pin_mem=args.pin_mem)
evaluator = create_evaluator(args.dataset, dataset, pred_yxyx=False)
bench.eval()
batch_time = AverageMeter()
end = time.time()
last_idx = len(loader) - 1
imgs = []
with torch.no_grad():
for i, (input, target) in enumerate(loader):
for b in range(input.shape[0]):
imgs.append(input[b].cpu().numpy())
# targets.append(target[b].cpu().numpy())
with amp_autocast():
output = bench(input, img_info=target)
evaluator.add_predictions(output, target)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.log_freq == 0 or i == last_idx:
print(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
.format(i,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg))
mean_ap = 0.
if dataset.parser.has_labels:
preds = [p[:2, :] for p in evaluator.predictions]
anns = evaluator.coco_api.imgToAnns
targets = [
np.asarray((anns[k][0]['bbox'], anns[k][1]['bbox']))
for k in range(len(imgs))
]
mean_ap = evaluator.evaluate()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
for i, img in enumerate(imgs):
img = imgs[i]
img_m = np.mean(img, axis=0)
for c in range(3):
img[c] = img_m
img_ = img.transpose(1, 2, 0)
m = img_.min()
M = img_.max()
img_ = ((img_ - m) / (M - m) * 255).astype('uint8').copy()
img_ = draw_bbox(img_, preds[i], targets[i])
cv2.imwrite(os.path.join(args.out_dir, '%d.jpg' % i), img_)
else:
evaluator.save(args.results)
return mean_ap
