def main():
args = parse_args()
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset_type = 'OIDSegDataset'
data_root = 'gs://oid2019/data/'
img_norm_cfg = dict(
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
#assert args.img_scale in [1333,1600]
print('img_scale:', args.img_scale)
print('flip_ratio:', args.flip)
dataset = get_dataset(
dict(
type=dataset_type,
ann_file='/home/bo_liu/' + args.ann_file,
img_prefix=data_root + ('val/' if 'seg_val' in args.ann_file else 'OD_test/'),
img_scale=args.img_scale,
img_norm_cfg=img_norm_cfg,
size_divisor=32,
flip_ratio=int(args.flip),
with_mask=True,
with_label=False,
test_mode=True)
)
data_loader = build_dataloader(
dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
test_cfg = mmcv.ConfigDict(dict(
rpn=dict(
nms_across_levels=False,
nms_pre=1000,
nms_post=1000,
max_num=1000,
nms_thr=0.7,
min_bbox_size=0),
rcnn=dict(
score_thr=args.thres,
# score_thr=0.0,
nms=dict(type='nms', iou_thr=0.5),
max_per_img=args.max_per_img,
mask_thr_binary=0.5),
keep_all_stages=False))
# build the model and load checkpoint
models = []
for config_path, checkpoint_path in zip(args.cfg_list, args.checkpoint):
print(f"config: {config_path}\\n checkpoint: {checkpoint_path}")
tmp_cfg = mmcv.Config.fromfile(config_path)
tmp_cfg.model.pretrained = None
tmp_cfg.data.test.test_mode = True
model = build_detector(tmp_cfg.model, train_cfg=None, test_cfg=test_cfg)
checkpoint = load_checkpoint(model, checkpoint_path, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
models.append(model)
model = EnsembleModel(models)
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('writing results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = args.out
coco_eval(result_file, eval_types, dataset.coco)
else:
if not isinstance(outputs[0], dict):
result_file = args.out + '.json'
results2json(dataset, outputs, result_file)
coco_eval(result_file, eval_types, dataset.coco)
else:
for name in outputs[0]:
print('Evaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}.json'.format(name)
results2json(dataset, outputs_, result_file)
coco_eval(result_file, eval_types, dataset.coco)
def main():
args = parse_args()
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = get_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
load_checkpoint(model, args.checkpoint, map_location='cpu')
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = args.out
coco_eval(result_file, eval_types, dataset.coco)
else:
if not isinstance(outputs[0], dict):
result_file = args.out + '.json'
results2json(dataset, outputs, result_file)
coco_eval(result_file, eval_types, dataset.coco)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}.json'.format(name)
results2json(dataset, outputs_, result_file)
coco_eval(result_file, eval_types, dataset.coco)
def main():
args = parse_args()
assert args.out or args.show or args.json_out, \
('Please specify at least one operation (save or show the results) '
'with the argument "--out" or "--show" or "--json_out"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
if args.json_out is not None and args.json_out.endswith('.json'):
args.json_out = args.json_out[:-5]
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = args.out
coco_eval(result_file, eval_types, dataset.coco)
else:
if not isinstance(outputs[0], dict):
result_files = results2json(dataset, outputs, args.out)
coco_eval(result_files, eval_types, dataset.coco)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}'.format(name)
result_files = results2json(dataset, outputs_,
result_file)
coco_eval(result_files, eval_types, dataset.coco)
# Save predictions in the COCO json format
if args.json_out and rank == 0:
if not isinstance(outputs[0], dict):
results2json(dataset, outputs, args.json_out)
else:
for name in outputs[0]:
outputs_ = [out[name] for out in outputs]
result_file = args.json_out + '.{}'.format(name)
results2json(dataset, outputs_, result_file)
def train_detector(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
logger = get_root_logger(cfg.log_level)
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
if 'imgs_per_gpu' in cfg.data:
logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
'Please use "samples_per_gpu" instead')
if 'samples_per_gpu' in cfg.data:
logger.warning(
f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
f'={cfg.data.imgs_per_gpu} is used in this experiments')
else:
logger.warning(
'Automatically set "samples_per_gpu"="imgs_per_gpu"='
f'{cfg.data.imgs_per_gpu} in this experiments')
cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
# cfg.gpus will be ignored if distributed
len(cfg.gpu_ids),
dist=distributed,
seed=cfg.seed) for ds in dataset
]
# put model on gpus
if distributed:
find_unused_parameters = cfg.get('find_unused_parameters', False)
# Sets the `find_unused_parameters` parameter in
# torch.nn.parallel.DistributedDataParallel
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDataParallel(model.cuda(cfg.gpu_ids[0]),
device_ids=cfg.gpu_ids)
# build runner
optimizer = build_optimizer(model, cfg.optimizer)
if 'runner' not in cfg:
cfg.runner = {
'type': 'EpochBasedRunner',
'max_epochs': cfg.total_epochs
}
warnings.warn(
'config is now expected to have a `runner` section, '
'please set `runner` in your config.', UserWarning)
else:
if 'total_epochs' in cfg:
assert cfg.total_epochs == cfg.runner.max_epochs
runner = build_runner(cfg.runner,
default_args=dict(model=model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
meta=meta))
# an ugly workaround to make .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg,
distributed=distributed)
elif distributed and 'type' not in cfg.optimizer_config:
optimizer_config = OptimizerHook(**cfg.optimizer_config)
else:
optimizer_config = cfg.optimizer_config
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config,
cfg.get('momentum_config', None))
if distributed:
if isinstance(runner, EpochBasedRunner):
runner.register_hook(DistSamplerSeedHook())
# register eval hooks
if validate:
# Support batch_size > 1 in validation
val_samples_per_gpu = cfg.data.val.pop('samples_per_gpu', 1)
if val_samples_per_gpu > 1:
# Replace 'ImageToTensor' to 'DefaultFormatBundle'
cfg.data.val.pipeline = replace_ImageToTensor(
cfg.data.val.pipeline)
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
val_dataloader = build_dataloader(
val_dataset,
samples_per_gpu=val_samples_per_gpu,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
eval_cfg = cfg.get('evaluation', {})
eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
eval_hook = DistEvalHook if distributed else EvalHook
runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
# user-defined hooks
if cfg.get('custom_hooks', None):
custom_hooks = cfg.custom_hooks
assert isinstance(custom_hooks, list), \
f'custom_hooks expect list type, but got {type(custom_hooks)}'
for hook_cfg in cfg.custom_hooks:
assert isinstance(hook_cfg, dict), \
'Each item in custom_hooks expects dict type, but got ' \
f'{type(hook_cfg)}'
hook_cfg = hook_cfg.copy()
priority = hook_cfg.pop('priority', 'NORMAL')
hook = build_from_cfg(hook_cfg, HOOKS)
runner.register_hook(hook, priority=priority)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow)
def main():
args = parse_args()
assert args.out or args.eval or args.format_only or args.show \
or args.show_dir, \
('Please specify at least one operation (save/eval/format/show the '
'results / save the results) with the argument "--out", "--eval"'
', "--format-only", "--show" or "--show-dir"')
if args.eval and args.format_only:
raise ValueError('--eval and --format_only cannot be both specified')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = Config.fromfile(args.config)
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# import modules from string list.
if cfg.get('custom_imports', None):
from mmcv.utils import import_modules_from_strings
import_modules_from_strings(**cfg['custom_imports'])
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
if cfg.model.get('neck'):
if isinstance(cfg.model.neck, list):
for neck_cfg in cfg.model.neck:
if neck_cfg.get('rfp_backbone'):
if neck_cfg.rfp_backbone.get('pretrained'):
neck_cfg.rfp_backbone.pretrained = None
elif cfg.model.neck.get('rfp_backbone'):
if cfg.model.neck.rfp_backbone.get('pretrained'):
cfg.model.neck.rfp_backbone.pretrained = None
# in case the test dataset is concatenated
if isinstance(cfg.data.test, dict):
cfg.data.test.test_mode = True
elif isinstance(cfg.data.test, list):
for ds_cfg in cfg.data.test:
ds_cfg.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
samples_per_gpu = cfg.data.test.pop('samples_per_gpu', 1)
if samples_per_gpu > 1:
# Replace 'ImageToTensor' to 'DefaultFormatBundle'
cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
if args.fuse_conv_bn:
model = fuse_conv_bn(model)
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show, args.show_dir,
args.show_score_thr)
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = multi_gpu_test(model, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if rank == 0:
if args.out:
print(f'\nwriting results to {args.out}')
mmcv.dump(outputs, args.out)
kwargs = {} if args.eval_options is None else args.eval_options
if args.format_only:
dataset.format_results(outputs, **kwargs)
if args.eval:
eval_kwargs = cfg.get('evaluation', {}).copy()
# hard-code way to remove EvalHook args
for key in ['interval', 'tmpdir', 'start', 'gpu_collect']:
eval_kwargs.pop(key, None)
eval_kwargs.update(dict(metric=args.eval, **kwargs))
print(dataset.evaluate(outputs, **eval_kwargs))
def evaluate_model(model_name, paper_arxiv_id, file_id, weights_name,
paper_results, config):
print('---')
print('Now Evaluating %s' % model_name)
evaluator = COCOEvaluator(root='./.data/vision/coco',
model_name=model_name,
paper_arxiv_id=paper_arxiv_id,
paper_results=paper_results)
out = 'results.pkl'
launcher = 'none'
if out is not None and not out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(config)
cfg.data.test[
'ann_file'] = './.data/vision/coco/annotations/instances_val2017.json'
cfg.data.test['img_prefix'] = './.data/vision/coco/val2017/'
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if launcher == 'none':
distributed = False
else:
distributed = True
init_dist(launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
destination = '%s/.cache/torch/' % (str(Path.home()))
download_file_from_google_drive(file_id,
destination,
filename=weights_name)
local_checkpoint = os.path.join(destination, weights_name)
print(local_checkpoint)
# '/home/ubuntu/GCNet/mask_rcnn_r50_fpn_1x_20181010-069fa190.pth'
checkpoint = load_checkpoint(model, local_checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
evaluator.reset_time()
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs, cache_exists = single_gpu_test(model, data_loader, False,
evaluator)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
if cache_exists:
print('Cache exists: %s' % (evaluator.batch_hash))
evaluator.save()
else:
rank, _ = get_dist_info()
if out and rank == 0:
print('\nwriting results to {}'.format(out))
mmcv.dump(outputs, out)
eval_types = ['bbox']
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = out
else:
if not isinstance(outputs[0], dict):
result_files = dataset.results2json(outputs, out)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = out + '.{}'.format(name)
result_files = dataset.results2json(
outputs_, result_file)
anns = json.load(open(result_files['bbox']))
evaluator.detections = []
evaluator.add(anns)
evaluator.save()
def _dist_train(model,
dataset,
cfg,
validate=False,
logger=None,
timestamp=None,
runner_attr_dict=dict()):
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(ds,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
dist=True) for ds in dataset
]
# build runner
runner_attr_dict.update({
'imgs_per_gpu': cfg.data.imgs_per_gpu,
'initial_lr': cfg.optimizer['lr']
})
if hasattr(dataset, 'CLASSES'):
runner_attr_dict.update({'classes': dataset.CLASSES})
optimizer = build_optimizer(model, cfg.optimizer)
search_optimizer = getattr(cfg, 'search_config',
{}).pop('search_optimizer', None)
runner = Runner(model,
batch_processor,
optimizer,
search_optimizer,
cfg.work_dir,
logger=logger,
runner_attr_dict=runner_attr_dict)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg)
else:
optimizer_config = DistOptimizerHook(**cfg.optimizer_config)
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config)
runner.register_hook(DistSamplerSeedHook())
if search_optimizer is not None:
runner.register_hook(DistSearchOptimizerHook())
runner.register_hook(SearchHook(**cfg.search_config))
# register eval hooks
if validate:
val_dataset_cfg = cfg.data.val
eval_cfg = cfg.get('evaluation', {})
if isinstance(model.module, RPN):
# TODO: implement recall hooks for other datasets
runner.register_hook(
CocoDistEvalRecallHook(val_dataset_cfg, **eval_cfg))
else:
dataset_type = DATASETS.get(val_dataset_cfg.type)
if issubclass(dataset_type, datasets.CocoDataset):
runner.register_hook(
CocoDistEvalmAPHook(val_dataset_cfg, **eval_cfg))
else:
runner.register_hook(
DistEvalmAPHook(val_dataset_cfg, **eval_cfg))
register_hooks(runner, cfg)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def main():
args = parse_args()
assert args.out or args.show or args.show_dir, \
('Please specify at least one operation (save or show the results) '
'with the argument "--out", "--show" or "show-dir"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# import modules from string list.
if cfg.get('custom_imports', None):
from mmcv.utils import import_modules_from_strings
import_modules_from_strings(**cfg['custom_imports'])
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
if args.workers == 0:
args.workers = cfg.data.workers_per_gpu
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# set random seeds
if args.seed is not None:
set_random_seed(args.seed)
if 'all' in args.corruptions:
corruptions = [
'gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur',
'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog',
'brightness', 'contrast', 'elastic_transform', 'pixelate',
'jpeg_compression', 'speckle_noise', 'gaussian_blur', 'spatter',
'saturate'
]
elif 'benchmark' in args.corruptions:
corruptions = [
'gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur',
'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog',
'brightness', 'contrast', 'elastic_transform', 'pixelate',
'jpeg_compression'
]
elif 'noise' in args.corruptions:
corruptions = ['gaussian_noise', 'shot_noise', 'impulse_noise']
elif 'blur' in args.corruptions:
corruptions = [
'defocus_blur', 'glass_blur', 'motion_blur', 'zoom_blur'
]
elif 'weather' in args.corruptions:
corruptions = ['snow', 'frost', 'fog', 'brightness']
elif 'digital' in args.corruptions:
corruptions = [
'contrast', 'elastic_transform', 'pixelate', 'jpeg_compression'
]
elif 'holdout' in args.corruptions:
corruptions = ['speckle_noise', 'gaussian_blur', 'spatter', 'saturate']
elif 'None' in args.corruptions:
corruptions = ['None']
args.severities = [0]
else:
corruptions = args.corruptions
rank, _ = get_dist_info()
aggregated_results = {}
for corr_i, corruption in enumerate(corruptions):
aggregated_results[corruption] = {}
for sev_i, corruption_severity in enumerate(args.severities):
# evaluate severity 0 (= no corruption) only once
if corr_i > 0 and corruption_severity == 0:
aggregated_results[corruption][0] = \
aggregated_results[corruptions[0]][0]
continue
test_data_cfg = copy.deepcopy(cfg.data.test)
# assign corruption and severity
if corruption_severity > 0:
corruption_trans = dict(type='Corrupt',
corruption=corruption,
severity=corruption_severity)
# TODO: hard coded "1", we assume that the first step is
# loading images, which needs to be fixed in the future
test_data_cfg['pipeline'].insert(1, corruption_trans)
# print info
print(f'\nTesting {corruption} at severity {corruption_severity}')
# build the dataloader
# TODO: support multiple images per gpu
# (only minor changes are needed)
dataset = build_dataset(test_data_cfg)
data_loader = build_dataloader(dataset,
samples_per_gpu=1,
workers_per_gpu=args.workers,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
cfg.model.train_cfg = None
model = build_detector(cfg.model, test_cfg=cfg.get('test_cfg'))
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model,
args.checkpoint,
map_location='cpu')
# old versions did not save class info in checkpoints,
# this walkaround is for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
show_dir = args.show_dir
if show_dir is not None:
show_dir = osp.join(show_dir, corruption)
show_dir = osp.join(show_dir, str(corruption_severity))
if not osp.exists(show_dir):
osp.makedirs(show_dir)
outputs = single_gpu_test(model, data_loader, args.show,
show_dir, args.show_score_thr)
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
if args.out and rank == 0:
eval_results_filename = (osp.splitext(args.out)[0] +
'_results' +
osp.splitext(args.out)[1])
mmcv.dump(outputs, args.out)
eval_types = args.eval
if cfg.dataset_type == 'VOCDataset':
if eval_types:
for eval_type in eval_types:
if eval_type == 'bbox':
test_dataset = mmcv.runner.obj_from_dict(
cfg.data.test, datasets)
logger = 'print' if args.summaries else None
mean_ap, eval_results = \
voc_eval_with_return(
args.out, test_dataset,
args.iou_thr, logger)
aggregated_results[corruption][
corruption_severity] = eval_results
else:
print('\nOnly "bbox" evaluation \
is supported for pascal voc')
else:
if eval_types:
print(f'Starting evaluate {" and ".join(eval_types)}')
if eval_types == ['proposal_fast']:
result_file = args.out
else:
if not isinstance(outputs[0], dict):
result_files = dataset.results2json(
outputs, args.out)
else:
for name in outputs[0]:
print(f'\nEvaluating {name}')
outputs_ = [out[name] for out in outputs]
result_file = args.out
+ f'.{name}'
result_files = dataset.results2json(
outputs_, result_file)
eval_results = coco_eval_with_return(
result_files, eval_types, dataset.coco)
aggregated_results[corruption][
corruption_severity] = eval_results
else:
print('\nNo task was selected for evaluation;'
'\nUse --eval to select a task')
# save results after each evaluation
mmcv.dump(aggregated_results, eval_results_filename)
if rank == 0:
# print filan results
print('\nAggregated results:')
prints = args.final_prints
aggregate = args.final_prints_aggregate
if cfg.dataset_type == 'VOCDataset':
get_results(eval_results_filename,
dataset='voc',
prints=prints,
aggregate=aggregate)
else:
get_results(eval_results_filename,
dataset='coco',
prints=prints,
aggregate=aggregate)
def train_detector(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
logger = get_root_logger(cfg.log_level)
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
if 'imgs_per_gpu' in cfg.data:
logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
'Please use "samples_per_gpu" instead')
if 'samples_per_gpu' in cfg.data:
logger.warning(
f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
f'={cfg.data.imgs_per_gpu} is used in this experiments')
else:
logger.warning(
'Automatically set "samples_per_gpu"="imgs_per_gpu"='
f'{cfg.data.imgs_per_gpu} in this experiments')
cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
# cfg.gpus will be ignored if distributed
len(cfg.gpu_ids),
dist=distributed,
seed=cfg.seed) for ds in dataset
]
# put model on gpus
if distributed:
find_unused_parameters = cfg.get('find_unused_parameters', False)
# Sets the `find_unused_parameters` parameter in
# torch.nn.parallel.DistributedDataParallel
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDataParallel(
model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
# build runner
optimizer = build_optimizer(model, cfg.optimizer)
runner = Runner(
model,
batch_processor,
optimizer,
cfg.work_dir,
logger=logger,
meta=meta)
# an ugly workaround to make .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(
**cfg.optimizer_config, **fp16_cfg, distributed=distributed)
elif distributed and 'type' not in cfg.optimizer_config:
optimizer_config = OptimizerHook(**cfg.optimizer_config)
else:
optimizer_config = cfg.optimizer_config
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config,
cfg.get('momentum_config', None))
if distributed:
runner.register_hook(DistSamplerSeedHook())
# register eval hooks
if validate:
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
val_dataloader = build_dataloader(
val_dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
eval_cfg = cfg.get('evaluation', {})
eval_hook = DistEvalHook if distributed else EvalHook
runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def main():
args = parse_args()
if args.out is not None and not args.out.endswith(('.json', '.pickle')):
raise ValueError('The output file must be a pkl file.')
for i in range(args.checkpoint_start, args.checkpoint_end):
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(
dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model,
train_cfg=None,
test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
if not args.mean_teacher:
while not osp.exists(args.checkpoint + str(i) + '.pth'):
time.sleep(5)
while i + 1 != args.checkpoint_end and not osp.exists(
args.checkpoint + str(i + 1) + '.pth'):
time.sleep(5)
checkpoint = load_checkpoint(model,
args.checkpoint + str(i) + '.pth',
map_location='cpu')
else:
while not osp.exists(args.checkpoint + str(i) + '.pth.stu'):
time.sleep(5)
while i + 1 != args.checkpoint_end and not osp.exists(
args.checkpoint + str(i + 1) + '.pth.stu'):
time.sleep(5)
checkpoint = load_checkpoint(model,
args.checkpoint + str(i) + '.pth.stu',
map_location='cpu')
checkpoint['meta'] = dict()
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show,
args.save_img, args.save_img_dir)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
res = []
for id, boxes in enumerate(outputs):
boxes = boxes[0]
if type(boxes) == list:
boxes = boxes[0]
boxes[:, [2, 3]] -= boxes[:, [0, 1]]
if len(boxes) > 0:
for box in boxes:
# box[:4] = box[:4] / 0.6
temp = dict()
temp['image_id'] = id + 1
temp['category_id'] = 1
temp['bbox'] = box[:4].tolist()
temp['score'] = float(box[4])
res.append(temp)
with open(args.out, 'w') as f:
json.dump(res, f)
MRs = validate('/media/ljp/Data/data/crowdhuman/validation.json',
args.out)
print(MRs)
print(
'Checkpoint %d: [Reasonable: %.2f%%], [Bare: %.2f%%], [Partial: %.2f%%], [Heavy: %.2f%%]'
% (i, MRs[0] * 100, MRs[1] * 100, MRs[2] * 100, MRs[3] * 100))
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.options is not None:
cfg.merge_from_dict(args.options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# work_directory is determined in this priority: CLI > segment in file > filename
if args.work_directory is not None:
# update configs according to CLI args if args.work_directory is not None
cfg.work_directory = args.work_directory
elif cfg.get('work_directory', None) is None:
# use config filename as default work_directory if cfg.work_directory is None
cfg.work_directory = osp.join('./work_dirs', osp.splitext(osp.basename(args.config))[0])
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.gpu_ids is not None:
cfg.gpu_ids = args.gpu_ids
else:
cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# create work_directory
mmcv.mkdir_or_exist(osp.abspath(cfg.work_directory))
# dump config
cfg.dump(osp.join(cfg.work_directory, osp.basename(args.config)))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_directory, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' + dash_line)
meta['env_info'] = env_info
# log some basic info
logger.info(f'Distributed training: {distributed}')
logger.info(f'Config:\n{cfg.pretty_text}')
# ---------- MI-AOD Training and Test Start Here ---------- #
# set random seeds
if args.seed is not None:
logger.info(f'Set random seed to {args.seed}, deterministic: {args.deterministic}')
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
X_L, X_U, X_all, all_anns = get_X_L_0(cfg)
# # load set and model
# # Please change it to the timestamp directory which you want to load data from.
# last_timestamp = '/20201013_154728'
# # Please change it to the cycle which you want to load data from.
# load_cycle = 0
# X_L = np.load(cfg.work_directory + last_timestamp +'/X_L_' + str(load_cycle) + '.npy')
# X_U = np.load(cfg.work_directory + last_timestamp +'/X_U_' + str(load_cycle) + '.npy')
# cfg.cycles = list(range(load_cycle, 7))
cfg.work_directory = cfg.work_directory + '/' + timestamp
mmcv.mkdir_or_exist(osp.abspath(cfg.work_directory))
np.save(cfg.work_directory + '/X_L_' + '0' + '.npy', X_L)
np.save(cfg.work_directory + '/X_U_' + '0' + '.npy', X_U)
initial_step = cfg.lr_config.step
for cycle in cfg.cycles:
# set random seeds
if args.seed is not None:
logger.info(f'Set random seed to {args.seed}, deterministic: {args.deterministic}')
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
# get the config of the labeled dataset
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
# load model
model = build_detector(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
# # Please change it to the epoch which you want to load model at.
# model_file_name = '/latest.pth'
# model.load_state_dict(torch.load(cfg.work_directory[:16] + last_timestamp + model_file_name)['state_dict'])
# load dataset
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None and cycle == 0:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmdet_version=__version__ + get_git_hash()[:7],
config=cfg.pretty_text, CLASSES=datasets[0].CLASSES)
model.CLASSES = datasets[0].CLASSES
for epoch in range(cfg.epoch):
# Only in the last 3 epoch does the learning rate need to be reduced and the model needs to be evaluated.
if epoch == cfg.epoch - 1:
cfg.lr_config.step = initial_step
cfg.evaluation.interval = cfg.epoch_ratio[0]
else:
cfg.lr_config.step = [1000]
cfg.evaluation.interval = 100
# ---------- Label Set Training ----------
if epoch == 0:
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(0)
cfg.total_epochs = cfg.epoch_ratio[0]
cfg_bak = cfg.deepcopy()
time.sleep(2)
for name, value in model.named_parameters():
value.requires_grad = True
train_detector(model, datasets, cfg,
distributed=distributed, validate=(not args.no_validate), timestamp=timestamp, meta=meta)
cfg = cfg_bak
# ---------- Re-weighting and Minimizing Instance Uncertainty ----------
cfg_u = create_X_U_file(cfg.deepcopy(), X_U, all_anns, cycle)
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets_u = [build_dataset(cfg_u.data.train)]
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(1)
cfg_u.total_epochs = cfg_u.epoch_ratio[1]
cfg.total_epochs = cfg.epoch_ratio[1]
cfg_u_bak = cfg_u.deepcopy()
cfg_bak = cfg.deepcopy()
time.sleep(2)
for name, value in model.named_parameters():
if name in cfg.theta_f_1:
value.requires_grad = False
elif name in cfg.theta_f_2:
value.requires_grad = False
else:
value.requires_grad = True
train_detector(model, [datasets, datasets_u], [cfg, cfg_u],
distributed=distributed, validate=(not args.no_validate), timestamp=timestamp, meta=meta)
cfg_u = cfg_u_bak
cfg = cfg_bak
# ---------- Re-weighting and Maximizing Instance Uncertainty ----------
cfg_u = create_X_U_file(cfg.deepcopy(), X_U, all_anns, cycle)
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets_u = [build_dataset(cfg_u.data.train)]
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(2)
cfg_u.total_epochs = cfg_u.epoch_ratio[1]
cfg.total_epochs = cfg.epoch_ratio[1]
cfg_u_bak = cfg_u.deepcopy()
cfg_bak = cfg.deepcopy()
time.sleep(2)
for name, value in model.named_parameters():
if name in cfg.theta_f_1:
value.requires_grad = True
elif name in cfg.theta_f_2:
value.requires_grad = True
else:
value.requires_grad = False
train_detector(model, [datasets, datasets_u], [cfg, cfg_u],
distributed=distributed,validate=(not args.no_validate), timestamp=timestamp, meta=meta)
cfg_u = cfg_u_bak
cfg = cfg_bak
# ---------- Label Set Training ----------
cfg = create_X_L_file(cfg, X_L, all_anns, cycle)
datasets = [build_dataset(cfg.data.train)]
losstype.update_vars(0)
cfg.total_epochs = cfg.epoch_ratio[0]
cfg_bak = cfg.deepcopy()
for name, value in model.named_parameters():
value.requires_grad = True
time.sleep(2)
train_detector(model, datasets, cfg,
distributed=distributed, validate=args.no_validate, timestamp=timestamp, meta=meta)
cfg = cfg_bak
# ---------- Informative Image Selection ----------
if cycle != cfg.cycles[-1]:
cfg_u, subset = create_selection_unlabeled_set(cfg.deepcopy(), X_all, X_L, all_anns, cycle)
# get new labeled data
dataset_al = build_dataset(cfg_u.data.test)
data_loader = build_dataloader(dataset_al, samples_per_gpu=1, workers_per_gpu=cfg.data.workers_per_gpu,
dist=False, shuffle=False)
# set random seeds
if args.seed is not None:
logger.info(f'Set random seed to {args.seed}, deterministic: {args.deterministic}')
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
uncertainty = calculate_uncertainty(cfg, model, data_loader, return_box=False)
# uncertainty = torch.rand(cfg.num_samples//cfg.subset_p)
# uncertainty = uncertainty.data.cpu()
# update labeled set
X_L, X_U = update_X_L(uncertainty, X_all, subset, X_L, cfg.X_S_size)
# save set and model
np.save(cfg.work_directory + '/X_L_' + str(cycle+1) + '.npy', X_L)
np.save(cfg.work_directory + '/X_U_' + str(cycle+1) + '.npy', X_U)
def main():
args = parse_args()
assert args.out or args.show or args.json_out, \
('Please specify at least one operation (save or show the results) '
'with the argument "--out" or "--show" or "--json_out"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
if args.out is not None and not args.auto_dir:
assert osp.exists(osp.dirname(
args.out)), 'output file directory does not exist!!'
if args.json_out is not None and args.json_out.endswith('.json'):
args.json_out = args.json_out[:-5]
cfg = mmcv.Config.fromfile(args.config)
if args.auto_dir:
work_dir = cfg.work_dir
args.checkpoint = osp.join(work_dir, args.checkpoint)
if args.out:
args.out = osp.join(work_dir, args.out)
if args.json_out:
args.json_out = osp.join(work_dir, args.json_out)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# if cfg.dataset_type == 'LvisDataSet':
# cfg.test_cfg.rcnn.max_per_img = 300
# cfg.test_cfg.rcnn.score_thr = 0.01
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
# use val dataset for visualize ground truch bboxes
dataset = build_dataset(cfg.data.val if args.show_gt else cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=1,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if args.analyze:
result = dict()
result['fc_weight'] = model.bbox_head.fc_cls.weight.data
samples_per_cls_file = cfg.data.train.samples_per_cls_file
if osp.exists(samples_per_cls_file): # add samples_per_cls_file
with open(samples_per_cls_file, 'r') as f:
samples_per_cls = torch.Tensor(
[int(line.strip()) for line in f.readlines()])
analyze(result, samples_per_cls)
exit()
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model,
data_loader,
args.show,
show_gt=args.show_gt,
work_dir=cfg.work_dir)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if not isinstance(outputs[0], dict):
result_files = results2json(dataset,
outputs,
args.out,
dump_json=False)
lvis_eval(result_files, eval_types, dataset.lvis)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}'.format(name)
result_files = results2json(dataset,
outputs_,
result_file,
dump_json=False)
lvis_eval(result_files, eval_types, dataset.lvis)
# Save predictions in the COCO json format
if args.json_out and rank == 0:
if not isinstance(outputs[0], dict):
results2json(dataset, outputs, args.json_out, dump_json=True)
else:
for name in outputs[0]:
outputs_ = [out[name] for out in outputs]
result_file = args.json_out + '.{}'.format(name)
results2json(dataset, outputs_, result_file, dump_json=True)
def main():
parser = argparse.ArgumentParser(description='MMDet test detector')
parser.add_argument('config', help='test config file path')
parser.add_argument('checkpoint', help='checkpoint file')
parser.add_argument('--out', help='output result file')
args = parser.parse_args()
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
dataset = get_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=False,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
#单GPU检测
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader)
if args.out:
#输出pkl文件
mmcv.dump(outputs, args.out)
print('Starting evaluate {}'.format(' and '.join("bbox")))
result_file = args.out + '.json'
#输出json文件
results2json(dataset, outputs, result_file)
coco = dataset.coco
assert isinstance(coco, COCO)
cocoGt = coco
cocoDt = coco.loadRes(result_file)
img_ids = coco.getImgIds()
iou_type = 'bbox'
params = {}
params = Params(iouType=iou_type)
params.imgIds = sorted(cocoGt.getImgIds())
params.catIds = sorted(cocoGt.getCatIds())
evalImgs = defaultdict(list)
eval = {}
_gts = defaultdict(list)
_dts = defaultdict(list)
_paramsEval = {}
stats = []
ious = {}
#prepare
params.imgIds = list(np.unique(params.imgIds))
params.catIds = list(np.unique(params.catIds))
params.maxDets = sorted(params.maxDets)
gts = cocoGt.loadAnns(
cocoGt.getAnnIds(imgIds=params.imgIds, catIds=params.catIds))
dts = cocoDt.loadAnns(
cocoDt.getAnnIds(imgIds=params.imgIds, catIds=params.catIds))
_gts = defaultdict(list)
_dts = defaultdict(list)
catIds = params.catIds
for gt in gts:
_gts[gt['image_id'], gt['category_id']].append(gt)
for dt in dts:
_dts[dt['image_id'], dt['category_id']].append(dt)
print("图片总数为：")
print(len(params.imgIds))
pic_num = len(params.imgIds)
#未找到目标的图片id
dtnum, gtnum = [], []
for i in _gts:
gtnum.append(i[0])
s1 = set(gtnum)
for j in _dts:
dtnum.append(j[0])
s2 = set(dtnum)
miss = s1 - s2
print("未找到目标的图片id：")
print(miss)
#未找到目标id对应的图片名
missname = []
for i in miss:
missname.append(cocoGt.loadImgs(i)[0]['file_name'])
print(missname)
#含有多检测框图片id
multi = []
for k, v in _dts.items():
if len(v) > 1:
multi.append(k)
print("含有多个检测框的图片总数：")
print(len(multi))
#evaluate
ious = {(imgId, catId): computeIoU(params, _gts, _dts, imgId, catId) \
for imgId in params.imgIds
for catId in catIds}
maxDet = params.maxDets[-1] #100
evalImgs = defaultdict(
list) # per-image per-category evaluation results
evalImgs = [evaluateImg(params,_gts, _dts, ious, imgId, catId, areaRng, maxDet) \
for catId in catIds
for areaRng in params.areaRng
for imgId in params.imgIds
]
_paramsEval = copy.deepcopy(params)
#IOU不符合的图片
cnt_50 = []
cnt_55 = []
cnt_60 = []
cnt_65 = []
cnt_70 = []
cnt_75 = []
cnt_80 = []
cnt_85 = []
cnt_90 = []
cnt_95 = []
cnt_100 = []
for i in range(pic_num):
e = evalImgs[i]['dtMatches']
zero_50 = np.count_nonzero(e[0])
if zero_50 == 0:
cnt_50.append(i + 1)
continue
zero_55 = np.count_nonzero(e[1])
if zero_55 == 0:
cnt_55.append(i + 1)
continue
zero_60 = np.count_nonzero(e[2])
if zero_60 == 0:
cnt_60.append(i + 1)
continue
zero_65 = np.count_nonzero(e[3])
if zero_65 == 0:
cnt_65.append(i + 1)
continue
zero_70 = np.count_nonzero(e[4])
if zero_70 == 0:
cnt_70.append(i + 1)
continue
zero_75 = np.count_nonzero(e[5])
if zero_75 == 0:
cnt_75.append(i + 1)
continue
zero_80 = np.count_nonzero(e[6])
if zero_80 == 0:
cnt_80.append(i + 1)
continue
zero_85 = np.count_nonzero(e[7])
if zero_85 == 0:
cnt_85.append(i + 1)
continue
zero_90 = np.count_nonzero(e[8])
if zero_90 == 0:
cnt_90.append(i + 1)
continue
zero_95 = np.count_nonzero(e[9])
if zero_95 == 0:
cnt_95.append(i + 1)
continue
cnt_100.append(i + 1)
print("IoU不足相应百分比的图片数量：")
print("#50#")
print(len(cnt_50))
print("#55#")
print(len(cnt_55))
print("#60#")
print(len(cnt_60))
print("#65#")
print(len(cnt_65))
print("#70#")
print(len(cnt_70))
print("#75#")
print(len(cnt_75))
print("#80#")
print(len(cnt_80))
print("#85#")
print(len(cnt_85))
print("#90#")
print(len(cnt_90))
print("#95#")
print(len(cnt_95))
print("#100#")
print(len(cnt_100))
print(
len(cnt_50) + len(cnt_55) + len(cnt_60) + len(cnt_65) +
len(cnt_70) + len(cnt_75) + len(cnt_80) + len(cnt_85) +
len(cnt_90) + len(cnt_95) + len(cnt_100))
for i in range(pic_num):
e = evalImgs[i]['dtMatches']
#print(e[0])
#画图
#显示不足指定iou的图片
imgIds_f = cocoDt.getImgIds(list(range(1,
len(img_ids) +
1))) #未识别成功的图片(false)
imgIds_t = cocoDt.getImgIds(
list(set([x + 1 for x in range(pic_num)]) - set(cnt_75)))
imageFile = "/media/wl/000675B10007A33A/DatasetRepo/haier/JPEG/"
plt.figure()
for i in range(len(imgIds_f)):
imgId = imgIds_f[i]
Img_dt = cocoDt.loadImgs(imgId)[0]
Img_gt = cocoGt.loadImgs(imgId)[0]
imageUrl = imageFile + Img_dt['file_name']
#显示GT标签
annId_gt = cocoGt.getAnnIds(Img_gt['id'])
imgAnn_gt = cocoGt.loadAnns(ids=annId_gt)
#gt缺少segmentation字段，补上
for ann in imgAnn_gt:
ann['segmentation'] = [[
ann['bbox'][0], ann['bbox'][1], ann['bbox'][0],
ann['bbox'][1] + ann['bbox'][3],
ann['bbox'][0] + ann['bbox'][2],
ann['bbox'][1] + ann['bbox'][3],
ann['bbox'][0] + ann['bbox'][2], ann['bbox'][1]
]]
I = io.imread(imageUrl)
plt.subplot(1, 2, 1)
plt.imshow(I)
cocoGt.showAnns(imgAnn_gt)
plt.title('GT')
#显示DT标签
annId_dt = cocoDt.getAnnIds(Img_dt['id'])
imgAnn_dt = cocoDt.loadAnns(ids=annId_dt)
plt.subplot(1, 2, 2)
plt.imshow(I)
cocoDt.showAnns(imgAnn_dt)
plt.title('DT')
#plt.show()
#保存图片到指定文件夹
plt.rcParams['savefig.dpi'] = 300 #dpi为300,图片尺寸为1800*1200
plt.savefig(
'/home/wl/mmdetection/pic_result/haiermask/{picname}.svg'.
format(picname=Img_gt['file_name']))
plt.clf()
#显示未识别到的图片
'''
print("未识别到目标的图片")
for failpic in miss:
p = cocoGt.loadImgs(failpic)[0]
imageUrl = imageFile+p['file_name']
I=io.imread(imageUrl)
plt.subplot(1,2,1)
plt.imshow(I)
Id_gt = cocoGt.getAnnIds(p['id'])
Ann_gt = cocoGt.loadAnns(ids=Id_gt)
for ann in Ann_gt:
ann['segmentation'] = [[ann['bbox'][0], ann['bbox'][1],
ann['bbox'][0], ann['bbox'][1]+ann['bbox'][3],
ann['bbox'][0]+ann['bbox'][2], ann['bbox'][1]+ann['bbox'][3],
ann['bbox'][0]+ann['bbox'][2], ann['bbox'][1]]]
cocoGt.showAnns(Ann_gt)
#plt.rcParams['savefig.dpi']=300 #dpi为300,图片尺寸为1800*1200
#plt.savefig('/home/wl/mmdetection/pic_result/miss/{picname}'.format(picname=p['file_name']))
plt.subplot(1,2,2)
plt.imshow(I)
plt.show()
'''
evalsmall = [evaluateImg(params,_gts, _dts, ious, imgId, catId, [0 ** 2, 32 ** 2], maxDet) \
for catId in catIds
for imgId in params.imgIds
]
#print("evalImags:")
#print(evalsmall[0])
#accumulate
p = params
p.catIds = p.catIds if p.useCats == 1 else [-1]
T = len(p.iouThrs)
R = len(p.recThrs)
K = len(p.catIds) if p.useCats else 1
A = len(p.areaRng)
M = len(p.maxDets)
precision = -np.ones(
(T, R, K, A, M)) # -1 for the precision of absent categories
recall = -np.ones((T, K, A, M))
scores = -np.ones((T, R, K, A, M))
# create dictionary for future indexing
_pe = _paramsEval
catIds = _pe.catIds if _pe.useCats else [-1]
setK = set(catIds)
setA = set(map(tuple, _pe.areaRng))
setM = set(_pe.maxDets)
setI = set(_pe.imgIds)
# get inds to evaluate
k_list = [n for n, k in enumerate(p.catIds) if k in setK]
m_list = [m for n, m in enumerate(p.maxDets) if m in setM]
a_list = [
n for n, a in enumerate(map(lambda x: tuple(x), p.areaRng))
if a in setA
]
i_list = [n for n, i in enumerate(p.imgIds) if i in setI]
I0 = len(_pe.imgIds)
A0 = len(_pe.areaRng)
# retrieve E at each category, area range, and max number of detections
for k, k0 in enumerate(k_list):
Nk = k0 * A0 * I0
for a, a0 in enumerate(a_list):
Na = a0 * I0
for m, maxDet in enumerate(m_list):
E = [evalImgs[Nk + Na + i] for i in i_list]
E = [e for e in E if not e is None]
if len(E) == 0:
continue
dtScores = np.concatenate(
[e['dtScores'][0:maxDet] for e in E])
# different sorting method generates slightly different results.
# mergesort is used to be consistent as Matlab implementation.
inds = np.argsort(-dtScores, kind='mergesort')
dtScoresSorted = dtScores[inds]
dtm = np.concatenate(
[e['dtMatches'][:, 0:maxDet] for e in E], axis=1)[:,
inds]
dtIg = np.concatenate(
[e['dtIgnore'][:, 0:maxDet] for e in E], axis=1)[:,
inds]
gtIg = np.concatenate([e['gtIgnore'] for e in E])
npig = np.count_nonzero(gtIg == 0)
if npig == 0:
continue
tps = np.logical_and(dtm, np.logical_not(dtIg))
fps = np.logical_and(np.logical_not(dtm),
np.logical_not(dtIg))
tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float)
fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float)
for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
tp = np.array(tp)
fp = np.array(fp)
nd = len(tp)
rc = tp / npig
pr = tp / (fp + tp + np.spacing(1))
q = np.zeros((R, ))
ss = np.zeros((R, ))
if nd:
recall[t, k, a, m] = rc[-1]
else:
recall[t, k, a, m] = 0
# numpy is slow without cython optimization for accessing elements
# use python array gets significant speed improvement
pr = pr.tolist()
q = q.tolist()
for i in range(nd - 1, 0, -1):
if pr[i] > pr[i - 1]:
pr[i - 1] = pr[i]
inds = np.searchsorted(rc, p.recThrs, side='left')
try:
for ri, pi in enumerate(inds):
q[ri] = pr[pi]
ss[ri] = dtScoresSorted[pi]
except:
pass
precision[t, :, k, a, m] = np.array(q)
scores[t, :, k, a, m] = np.array(ss)
eval = {
'params': p,
'counts': [T, R, K, A, M],
'date': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
'precision': precision,
'recall': recall,
'scores': scores,
}
#summarize
def _summarize(ap=1, iouThr=None, areaRng='all', maxDets=100):
p = params
iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
typeStr = '(AP)' if ap == 1 else '(AR)'
iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
if iouThr is None else '{:0.2f}'.format(iouThr)
aind = [
i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng
]
mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
if ap == 1:
# dimension of precision: [TxRxKxAxM]
s = eval['precision']
# IoU
if iouThr is not None:
t = np.where(iouThr == p.iouThrs)[0]
s = s[t]
s = s[:, :, :, aind, mind]
else:
# dimension of recall: [TxKxAxM]
s = eval['recall']
if iouThr is not None:
t = np.where(iouThr == p.iouThrs)[0]
s = s[t]
s = s[:, :, aind, mind]
if len(s[s > -1]) == 0:
mean_s = -1
else:
mean_s = np.mean(s[s > -1])
print(
iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets,
mean_s))
return mean_s
def _summarizeDets():
stats = np.zeros((12, ))
stats[0] = _summarize(1)
stats[1] = _summarize(1, iouThr=.5, maxDets=params.maxDets[2])
stats[2] = _summarize(1, iouThr=.75, maxDets=params.maxDets[2])
stats[3] = _summarize(1,
areaRng='small',
maxDets=params.maxDets[2])
stats[4] = _summarize(1,
areaRng='medium',
maxDets=params.maxDets[2])
stats[5] = _summarize(1,
areaRng='large',
maxDets=params.maxDets[2])
stats[6] = _summarize(0, maxDets=params.maxDets[0])
stats[7] = _summarize(0, maxDets=params.maxDets[1])
stats[8] = _summarize(0, maxDets=params.maxDets[2])
stats[9] = _summarize(0,
areaRng='small',
maxDets=params.maxDets[2])
stats[10] = _summarize(0,
areaRng='medium',
maxDets=params.maxDets[2])
stats[11] = _summarize(0,
areaRng='large',
maxDets=params.maxDets[2])
return stats
def _summarizeKps():
stats = np.zeros((10, ))
stats[0] = _summarize(1, maxDets=20)
stats[1] = _summarize(1, maxDets=20, iouThr=.5)
stats[2] = _summarize(1, maxDets=20, iouThr=.75)
stats[3] = _summarize(1, maxDets=20, areaRng='medium')
stats[4] = _summarize(1, maxDets=20, areaRng='large')
stats[5] = _summarize(0, maxDets=20)
stats[6] = _summarize(0, maxDets=20, iouThr=.5)
stats[7] = _summarize(0, maxDets=20, iouThr=.75)
stats[8] = _summarize(0, maxDets=20, areaRng='medium')
stats[9] = _summarize(0, maxDets=20, areaRng='large')
return stats
iouType = params.iouType
summarize = _summarizeDets
stats = summarize()
def train_detector(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
logger = get_root_logger(cfg.log_level)
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
if 'imgs_per_gpu' in cfg.data:
logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
'Please use "samples_per_gpu" instead')
if 'samples_per_gpu' in cfg.data:
logger.warning(
f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
f'={cfg.data.imgs_per_gpu} is used in this experiments')
else:
logger.warning(
'Automatically set "samples_per_gpu"="imgs_per_gpu"='
f'{cfg.data.imgs_per_gpu} in this experiments')
cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
# cfg.gpus will be ignored if distributed
len(cfg.gpu_ids),
dist=distributed,
seed=cfg.seed) for ds in dataset
]
if cfg.load_from:
load_checkpoint(model=model, filename=cfg.load_from)
# put model on gpus
if torch.cuda.is_available():
model = model.cuda()
# nncf model wrapper
nncf_enable_compression = bool(cfg.get('nncf_config'))
if nncf_enable_compression:
compression_ctrl, model = wrap_nncf_model(model, cfg, data_loaders[0],
get_fake_input)
else:
compression_ctrl = None
map_location = 'default'
if torch.cuda.is_available():
if distributed:
# put model on gpus
find_unused_parameters = cfg.get('find_unused_parameters', False)
# Sets the `find_unused_parameters` parameter in
# torch.nn.parallel.DistributedDataParallel
model = MMDistributedDataParallel(
model,
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDataParallel(model.cuda(cfg.gpu_ids[0]),
device_ids=cfg.gpu_ids)
else:
model = MMDataCPU(model)
map_location = 'cpu'
if nncf_enable_compression and distributed:
compression_ctrl.distributed()
# build runner
optimizer = build_optimizer(model, cfg.optimizer)
runner = EpochBasedRunner(model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
meta=meta)
# an ugly workaround to make .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg,
distributed=distributed)
elif distributed and 'type' not in cfg.optimizer_config:
optimizer_config = OptimizerHook(**cfg.optimizer_config)
else:
optimizer_config = cfg.optimizer_config
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config,
cfg.get('momentum_config', None))
if distributed:
runner.register_hook(DistSamplerSeedHook())
add_logging_on_first_and_last_iter(runner)
# register eval hooks
if validate:
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
val_dataloader = build_dataloader(
val_dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
eval_cfg = cfg.get('evaluation', {})
eval_hook = DistEvalHook if distributed else EvalHook
if nncf_enable_compression:
eval_hook = DistEvalPlusBeforeRunHook if distributed else EvalPlusBeforeRunHook
runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
if nncf_enable_compression:
runner.register_hook(
CompressionHook(compression_ctrl=compression_ctrl))
if cfg.resume_from:
runner.resume(cfg.resume_from, map_location=map_location)
runner.run(data_loaders,
cfg.workflow,
cfg.total_epochs,
compression_ctrl=compression_ctrl)
def main():
args = parse_args()
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = get_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
load_checkpoint(model, args.checkpoint, map_location='cpu')
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
vggmodel = VGG()
vggmodel.load_state_dict(torch.load('~/.torch/models/vgg16-397923af.pth'))
vggmodel.cuda()
_, vgg2deepfea = vggmodel(outputs)
c1model = rankmodel()
c1model.load_state_dict(torch.load('model/c1.pth'))
c1gmodel.cuda()
_, c1deep = rankmodel(vgg2deepfea)
outputs_c1x = [out[x] for out in outputs]
outputs_c1y = [out[y] for out in outputs]
fov = 75 * M_PI / 180;
invRot=calcRotation(outputs_c1x,outputs_c1y,fov)
vggmodel = VGG()
vggmodel.load_state_dict(torch.load('~/.torch/models/vgg16-397923af.pth'))
vggmodel.cuda()
_, c2deepfea = vggmodel(invRot)
c2model = rankmodel()
c2model.load_state_dict(torch.load('model/c2.pth'))
c2gmodel.cuda()
_, c2deep = rankmodel(c2deepfea)
model = RF_Model(c2deep)
def main():
args = parse_args()
assert args.out or args.eval or args.format_only or args.show \
or args.show_dir, \
('Please specify at least one operation (save/eval/format/show the '
'results / save the results) with the argument "--out", "--eval"'
', "--format-only", "--show" or "--show-dir"')
if args.eval and args.format_only:
raise ValueError('--eval and --format_only cannot be both specified')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = Config.fromfile(args.config)
if cfg.get('USE_MMDET', False):
from mmdet.apis import multi_gpu_test, single_gpu_test
from mmdet.datasets import build_dataloader
from mmdet.models import build_detector as build_model
else:
from mmtrack.apis import multi_gpu_test, single_gpu_test
from mmtrack.datasets import build_dataloader
from mmtrack.models import build_model
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# cfg.model.pretrains = None
if hasattr(cfg.model, 'detector'):
cfg.model.detector.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
if cfg.get('test_cfg', False):
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
else:
model = build_model(cfg.model)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
if args.checkpoint is not None:
checkpoint = load_checkpoint(model,
args.checkpoint,
map_location='cpu')
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
if not hasattr(model, 'CLASSES'):
model.CLASSES = dataset.CLASSES
if args.fuse_conv_bn:
model = fuse_conv_bn(model)
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show, args.show_dir,
args.show_score_thr)
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = multi_gpu_test(model, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if rank == 0:
if args.out:
print(f'\nwriting results to {args.out}')
mmcv.dump(outputs, args.out)
kwargs = {} if args.eval_options is None else args.eval_options
if args.format_only:
dataset.format_results(outputs, **kwargs)
if args.eval:
eval_kwargs = cfg.get('evaluation', {}).copy()
# hard-code way to remove EvalHook args
for key in ['interval', 'tmpdir', 'start', 'gpu_collect']:
eval_kwargs.pop(key, None)
eval_kwargs.update(dict(metric=args.eval, **kwargs))
print(dataset.evaluate(outputs, **eval_kwargs))
def _dist_train(model,
dataset,
cfg,
validate=False,
logger=None,
timestamp=None,
meta=None):
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(ds,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
dist=True,
seed=cfg.seed) for ds in dataset
]
# put model on gpus
find_unused_parameters = cfg.get('find_unused_parameters', True)
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
# build runner
optimizer = build_optimizer(model, cfg.optimizer)
runner = Runner(model,
batch_processor,
optimizer,
cfg.work_dir,
logger=logger,
meta=meta)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg)
else:
optimizer_config = DistOptimizerHook(**cfg.optimizer_config)
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config)
runner.register_hook(DistSamplerSeedHook())
# register eval hooks
if validate:
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
val_dataloader = build_dataloader(
val_dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=True,
shuffle=False)
eval_cfg = cfg.get('evaluation', {})
runner.register_hook(DistEvalHook(val_dataloader, **eval_cfg))
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def main():
args = parse_args()
assert args.eval or args.show \
or args.show_dir, \
('Please specify at least one operation (eval/show the '
'results) with the argument "--eval"'
', "--show" or "--show-dir"')
cfg = Config.fromfile(args.config)
if cfg.get('USE_MMDET', False):
from mmdet.apis import multi_gpu_test, single_gpu_test
from mmdet.datasets import build_dataloader
from mmdet.models import build_detector as build_model
if 'detector' in cfg.model:
cfg.model = cfg.model.detector
elif cfg.get('USE_MMCLS', False):
from mmtrack.apis import multi_gpu_test, single_gpu_test
from mmtrack.datasets import build_dataloader
from mmtrack.models import build_reid as build_model
if 'reid' in cfg.model:
cfg.model = cfg.model.reid
else:
from mmtrack.apis import multi_gpu_test, single_gpu_test
from mmtrack.datasets import build_dataloader
from mmtrack.models import build_model
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
logger = get_logger('SOTParamsSearcher', log_file=args.log)
# build the model and load checkpoint
if cfg.get('test_cfg', False):
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
else:
model = build_model(cfg.model)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
if args.checkpoint is not None:
checkpoint = load_checkpoint(model,
args.checkpoint,
map_location='cpu')
if 'meta' in checkpoint and 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
if not hasattr(model, 'CLASSES'):
model.CLASSES = dataset.CLASSES
if args.fuse_conv_bn:
model = fuse_conv_bn(model)
if not distributed:
model = MMDataParallel(model, device_ids=[0])
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
# init best_score, best_results and best parames
if 'meta' in checkpoint and 'hook_msgs' in checkpoint[
'meta'] and 'best_score' in checkpoint['meta']['hook_msgs']:
best_score = checkpoint['meta']['hook_msgs']['best_score']
else:
best_score = 0
key_metric = cfg.evaluation.save_best
best_result = {f'{key_metric}': best_score}
best_params = dict(penalty_k=cfg.model.test_cfg.rpn.penalty_k,
lr=cfg.model.test_cfg.rpn.lr,
win_influ=cfg.model.test_cfg.rpn.window_influence)
print_log(f'init best score as: {best_score}', logger)
print_log(f'init best params as: {best_params}', logger)
num_cases = len(args.penalty_k_range) * len(args.lr_range) * len(
args.win_influ_range)
case_count = 0
# compare function setting in parameter search
rule_map = {'greater': lambda x, y: x > y, 'less': lambda x, y: x < y}
compare_func = rule_map[cfg.evaluation.rule]
for penalty_k in args.penalty_k_range:
for lr in args.lr_range:
for win_influ in args.win_influ_range:
case_count += 1
cfg.model.test_cfg.rpn.penalty_k = penalty_k
cfg.model.test_cfg.rpn.lr = lr
cfg.model.test_cfg.rpn.window_influence = win_influ
print_log(f'-----------[{case_count}/{num_cases}]-----------',
logger)
print_log(
f'penalty_k={penalty_k} lr={lr} win_influence={win_influ}',
logger)
if not distributed:
outputs = single_gpu_test(
model,
data_loader,
args.show,
args.show_dir,
show_score_thr=args.show_score_thr)
else:
outputs = multi_gpu_test(model, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if rank == 0:
kwargs = args.eval_options if args.eval_options else {}
if args.eval:
eval_kwargs = cfg.get('evaluation', {}).copy()
# hard-code way to remove EvalHook args
eval_hook_args = [
'interval', 'tmpdir', 'start', 'gpu_collect',
'save_best', 'rule', 'by_epoch'
]
for key in eval_hook_args:
eval_kwargs.pop(key, None)
eval_kwargs.update(dict(metric=args.eval, **kwargs))
eval_results = dataset.evaluate(outputs, **eval_kwargs)
print_log(f'evaluation results: {eval_results}',
logger)
print_log('------------------------------------------',
logger)
if compare_func(eval_results[key_metric],
best_result[key_metric]):
best_result = eval_results
best_params['penalty_k'] = penalty_k,
best_params['lr'] = lr,
best_params['win_influ'] = win_influ
print_log(
f'The current best evaluation results: \
{best_result}', logger)
print_log(f'The current best params: {best_params}',
logger)
print_log(
f'After parameter searching, the best evaluation results: \
{best_result}', logger)
print_log(f'After parameter searching, the best params: {best_params}',
logger)
def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpuid
img_dir = args.img_dir
out_dir = args.out_dir
batch_size = args.batch_size
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
if args.img_dir != '':
file_list = common.load_filepaths(args.img_dir,
suffix=('.jpg', '.png', '.jpeg'),
recursive=True)
elif args.img_list != '':
file_list = parse_testfile(args.img_list)
else:
raise "Both img_dir and img_list is empty."
dataset = FilesDataset(file_list, cfg.test_pipeline)
data_loader = build_dataloader(dataset,
imgs_per_gpu=batch_size,
workers_per_gpu=batch_size,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
model = reweight_cls(model, args.tau).cuda()
model = MMDataParallel(model, device_ids=[0])
model.eval()
count = 0
for i, data in enumerate(data_loader):
with torch.no_grad():
# bbox_results, segm_results
results = model(return_loss=False, rescale=True, **data)
# batch
#for result in results:
# file_path = file_list[count]
# save_name = file_path.replace('/home/songbai.xb/workspace/projects/TAO/data/TAO/frames/val/', '')
# save_path = os.path.join(out_dir, save_name)
# common.makedirs(os.path.dirname(save_path))
# save_in_tao_format(result, save_path)
# count += 1
file_path = file_list[i]
save_name = file_path.replace(
'/home/songbai.xb/workspace/projects/TAO/data/TAO/frames/val/', '')
save_name = save_name.replace('.jpg', '.pkl').replace('.jpeg', )
save_path = os.path.join(out_dir, save_name)
common.makedirs(os.path.dirname(save_path))
save_in_tao_format(results[0], save_path)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.work_dir is not None:
cfg.work_dir = args.work_dir
pathlib.Path(cfg.work_dir).mkdir(parents=True, exist_ok=True)
cfg.gpus = args.gpus
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# init logger before other steps
logger = get_root_logger(cfg.work_dir)
logger.info('Distributed training: {}'.format(distributed))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}'.format(args.seed))
set_random_seed(args.seed)
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
if distributed:
model = MMDistributedDataParallel(model.cuda())
else:
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
train_dataset = get_dataset(cfg.data.train)
optimizer = build_optimizer(model, cfg.optimizer)
train_loader = build_dataloader(train_dataset,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
dist=distributed)
start_epoch = it = 0
last_epoch = -1
lr_scheduler, lr_warmup_scheduler = build_scheduler(
optimizer,
total_iters_each_epoch=len(train_loader),
total_epochs=cfg.total_epochs,
last_epoch=last_epoch,
optim_cfg=cfg.optimizer,
lr_cfg=cfg.lr_config)
# -----------------------start training---------------------------
logger.info('**********************Start training**********************')
train_model(model,
optimizer,
train_loader,
lr_scheduler=lr_scheduler,
optim_cfg=cfg.optimizer,
start_epoch=start_epoch,
total_epochs=cfg.total_epochs,
start_iter=it,
rank=args.local_rank,
logger=logger,
ckpt_save_dir=cfg.work_dir,
lr_warmup_scheduler=lr_warmup_scheduler,
ckpt_save_interval=cfg.checkpoint_config.interval,
max_ckpt_save_num=args.max_ckpt_save_num,
log_interval=cfg.log_config.interval)
logger.info('**********************End training**********************')
def _non_dist_train(model,
dataset,
cfg,
validate=False,
logger=None,
timestamp=None,
runner_attr_dict=dict()):
if validate:
raise NotImplementedError('Built-in validation is not implemented '
'yet in not-distributed training. Use '
'distributed training or test.py and '
'*eval.py scripts instead.')
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(ds,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False) for ds in dataset
]
# build runner
runner_attr_dict.update({
'imgs_per_gpu': cfg.data.imgs_per_gpu,
'initial_lr': cfg.optimizer['lr']
})
if hasattr(dataset, 'CLASSES'):
runner_attr_dict.update({'classes': dataset.CLASSES})
optimizer = build_optimizer(model, cfg.optimizer)
search_optimizer = getattr(getattr(cfg, 'search_config', {}),
'search_optimizer', None)
assert search_optimizer is None, "Not support"
runner = Runner(model,
batch_processor,
optimizer,
search_optimizer,
cfg.work_dir,
logger=logger,
runner_attr_dict=runner_attr_dict)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg,
distributed=False)
else:
optimizer_config = cfg.optimizer_config
runner.register_training_hooks(cfg.lr_config, optimizer_config)
# register eval hooks
if validate:
if isinstance(model.module, RPN):
runner.register_hook(CocoDistEvalRecallHook(cfg.data.val))
else:
if cfg.data.val.type == 'CocoDataset':
runner.register_hook(CocoDistEvalmAPHook(cfg.data.val))
else:
runner.register_hook(DistEvalmAPHook(cfg.data.val))
register_hooks(runner, cfg)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def train_detector(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
logger = get_root_logger(cfg.log_level)
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
if 'imgs_per_gpu' in cfg.data:
logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
'Please use "samples_per_gpu" instead')
if 'samples_per_gpu' in cfg.data:
logger.warning(
f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
f'={cfg.data.imgs_per_gpu} is used in this experiments')
else:
logger.warning(
'Automatically set "samples_per_gpu"="imgs_per_gpu"='
f'{cfg.data.imgs_per_gpu} in this experiments')
cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
# cfg.gpus will be ignored if distributed
len(cfg.gpu_ids),
dist=distributed,
seed=cfg.seed) for ds in dataset
]
# put model on gpus
if distributed:
find_unused_parameters = cfg.get('find_unused_parameters', False)
# Sets the `find_unused_parameters` parameter in
# torch.nn.parallel.DistributedDataParallel
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDataParallel(model.cuda(cfg.gpu_ids[0]),
device_ids=cfg.gpu_ids)
# model.module.backbone._init_weights('/home/pengyi/freq_attention/log/20201010_042828freq_resnet50_fp16_freq_sel8in1_vanillaSE_4layers_0.1lr_100epoch_cosinedecay_LSR_16freq/model_best.pth.tar')
# build runner
optimizer = build_optimizer(model, cfg.optimizer)
runner = EpochBasedRunner(model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
meta=meta)
# an ugly workaround to make .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
**fp16_cfg,
distributed=distributed)
elif distributed and 'type' not in cfg.optimizer_config:
optimizer_config = OptimizerHook(**cfg.optimizer_config)
else:
optimizer_config = cfg.optimizer_config
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config,
cfg.get('momentum_config', None))
if distributed:
runner.register_hook(DistSamplerSeedHook())
# register eval hooks
if validate:
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
val_dataloader = build_dataloader(
val_dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
eval_cfg = cfg.get('evaluation', {})
eval_hook = DistEvalHook if distributed else EvalHook
runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
# user-defined hooks
if cfg.get('custom_hooks', None):
custom_hooks = cfg.custom_hooks
assert isinstance(custom_hooks, list), \
f'custom_hooks expect list type, but got {type(custom_hooks)}'
for hook_cfg in cfg.custom_hooks:
assert isinstance(hook_cfg, dict), \
'Each item in custom_hooks expects dict type, but got ' \
f'{type(hook_cfg)}'
hook_cfg = hook_cfg.copy()
priority = hook_cfg.pop('priority', 'NORMAL')
hook = build_from_cfg(hook_cfg, HOOKS)
runner.register_hook(hook, priority=priority)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def main():
args = parse_args()
assert args.out or args.eval or args.format_only or args.show, \
('Please specify at least one operation (save/eval/format/show the '
'results) with the argument "--out", "--eval", "--format_only" '
'or "--show"')
if args.eval and args.format_only:
raise ValueError('--eval and --format_only cannot be both specified')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
# 计算需要剪枝的变量个数total
model.cuda()
total = 0
for m in model.backbone.modules():
if isinstance(m, nn.BatchNorm2d):
total += m.weight.data.shape[0]
# 确定剪枝的全局阈值
bn = torch.zeros(total)
index = 0
for m in model.backbone.modules():
if isinstance(m, nn.BatchNorm2d):
size = m.weight.data.shape[0]
bn[index:(index + size)] = m.weight.data.abs().clone()
index += size
# 按照权值大小排序
y, i = torch.sort(bn)
thre_index = int(total * args.percent)
# 确定要剪枝的阈值
thre = y[thre_index].cuda()
# ********************************预剪枝*********************************#
pruned = 0
cfg_ori = []
cfg = []
cfg_mask = []
model_backbone = list(model.backbone.modules())
for layer_id, m in enumerate(model_backbone):
if isinstance(m, nn.BatchNorm2d):
weight_copy = m.weight.data.abs().clone()
if isinstance(model_backbone[layer_id + 1], channel_selection):
mask = torch.ones(weight_copy.shape[0]).cuda()
else:
# 要保留的通道标记Mask图
mask = weight_copy.gt(thre).float().cuda()
# 要保留的通道标记Mask图
pruned = pruned + mask.shape[0] - torch.sum(mask)
# m.weight.data.mul_(mask)
# m.bias.data.mul_(mask)
cfg.append(int(torch.sum(mask)))
cfg_ori.append(mask.shape[0])
cfg_mask.append(mask.clone())
print(
'layer index: {:d} \t total channel: {:d} \t remaining channel: {:d}'
.format(layer_id, mask.shape[0], int(torch.sum(mask))))
pruned_ratio = pruned / total
print("剪枝比例：")
print(pruned_ratio)
print('Pre-processing Successful!')
print('cfg:')
print(cfg)
# ******************************* 正式剪枝 ********************************#
# 每个阶的最一层不剪枝
newmodel = copy.deepcopy(model)
newmodel.backbone = PResNet(depth=101,
num_stages=4,
out_indices=(0, 1, 2, 3),
frozen_stages=1,
norm_cfg=dict(type='BN', requires_grad=True),
style='pytorch',
cfg=cfg)
newmodel.cuda()
# print(newmodel.backbone)
num_parameters = sum([param.nelement() for param in newmodel.parameters()])
savepath = os.path.join(args.save_path, "prune.txt")
# with open(savepath, "w") as fp:
# fp.write("Configuration: \n" + str(cfg) + "\n")
# fp.write("Number of parameters: \n" + str(num_parameters) + "\n")
# fp.write("Test accuracy: \n" + str(acc))
old_modules = list(model.backbone.modules())
new_modules = list(newmodel.backbone.modules())
layer_id_in_cfg = 0
start_mask = torch.ones(3)
end_mask = cfg_mask[layer_id_in_cfg]
conv_count = 0
# downsample_conv_list = [17, 48, 88, 299]
downsample_conv_list = [17, 49, 90, 302]
for layer_id, m0 in enumerate(old_modules):
# m0 = old_modules[layer_id]
# print('m0:')
# print(m0)
m1 = new_modules[layer_id]
# print('m1:')
# print(m1)
if isinstance(m0, nn.BatchNorm2d):
idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
if idx1.size == 1:
idx1 = np.resize(idx1, (1, ))
if isinstance(old_modules[layer_id + 1], channel_selection):
# If the next layer is the channel selection layer, then the current batchnorm 2d layer won't be pruned.
m1.weight.data = m0.weight.data.clone()
m1.bias.data = m0.bias.data.clone()
m1.running_mean = m0.running_mean.clone()
m1.running_var = m0.running_var.clone()
# We need to set the channel selection layer.
m2 = new_modules[layer_id + 1]
m2.indexes.data.zero_()
m2.indexes.data[idx1.tolist()] = 1.0
layer_id_in_cfg += 1
start_mask = end_mask.clone()
if layer_id_in_cfg < len(cfg_mask):
end_mask = cfg_mask[layer_id_in_cfg]
else:
m1.weight.data = m0.weight.data[idx1.tolist()].clone()
m1.bias.data = m0.bias.data[idx1.tolist()].clone()
m1.running_mean = m0.running_mean[idx1.tolist()].clone()
m1.running_var = m0.running_var[idx1.tolist()].clone()
layer_id_in_cfg += 1
start_mask = end_mask.clone()
if layer_id_in_cfg < len(
cfg_mask): # do not change in Final FC
end_mask = cfg_mask[layer_id_in_cfg]
elif isinstance(m0, nn.Conv2d):
if conv_count == 0:
m1.weight.data = m0.weight.data.clone()
conv_count += 1
continue
if layer_id in downsample_conv_list:
# We need to consider the case where there are downsampling convolutions.
# For these convolutions, we just copy the weights.
m1.weight.data = m0.weight.data.clone()
continue
if isinstance(old_modules[layer_id + 1], nn.BatchNorm2d):
# This convers the convolutions in the residual block.
# The convolutions are either after the channel selection layer or after the batch normalization layer.
idx0 = np.squeeze(
np.argwhere(np.asarray(start_mask.cpu().numpy())))
idx1 = np.squeeze(
np.argwhere(np.asarray(end_mask.cpu().numpy())))
print('In shape: {:d}, Out shape {:d}.'.format(
idx0.size, idx1.size))
if idx0.size == 1:
idx0 = np.resize(idx0, (1, ))
if idx1.size == 1:
idx1 = np.resize(idx1, (1, ))
w1 = m0.weight.data[:, idx0.tolist(), :, :].clone()
# If the current convolution is not the last convolution in the residual block, then we can change the
# number of output channels. Currently we use `conv_count` to detect whether it is such convolution.
# if conv_count % 3 != 0:
w1 = w1[idx1.tolist(), :, :, :].clone()
m1.weight.data = w1.clone()
conv_count += 1
continue
elif isinstance(m0, nn.Linear):
idx0 = np.squeeze(np.argwhere(np.asarray(
start_mask.cpu().numpy())))
if idx0.size == 1:
idx0 = np.resize(idx0, (1, ))
m1.weight.data = m0.weight.data[:, idx0].clone()
m1.bias.data = m0.bias.data.clone()
# torch.save({'cfg': cfg, 'state_dict': newmodel.state_dict()}, os.path.join(args.save_path, 'pruned.pth.tar'))
# torch.save(newmodel.state_dict(), os.path.join(args.save_path, 'pruned.pth'))
print(newmodel)
torch.save(newmodel, os.path.join(args.save_path, 'pruned.pth'))
# print(newmodel)
if not distributed:
newmodel = MMDataParallel(newmodel, device_ids=[0])
newmodel = single_gpu_test(newmodel, data_loader, args.show)
else:
newmodel = MMDistributedDataParallel(
newmodel.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = multi_gpu_test(newmodel, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if rank == 0:
if args.out:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
kwargs = {} if args.options is None else args.options
if args.format_only:
dataset.format_results(outputs, **kwargs)
if args.eval:
dataset.evaluate(outputs, args.eval, **kwargs)
def main():
# os.environ["CUDA_VISIBLE_DEVICES"] = "1"
args = parse_args()
assert args.out or args.show, \
('Please specify at least one operation (save or show the results) '
'with the argument "--out" or "--show"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
result_file = args.out
# args = parser.parse_args()
# cfg = mmcv.Config.fromfile(args.config)
# test_dataset = mmcv.runner.obj_from_dict(cfg.data.test, datasets)
# txt_eval(args.result, test_dataset, args.iou_thr)
txt_eval(result_file, dataset, iou_thr=args.iou_thr)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# update configs according to CLI args
if args.dir is not None:
if args.dir.startswith('//'):
cfg.work_dir = args.dir[2:]
else:
localhost = get_localhost().split('.')[0]
# results from server saved to /private
if 'gpu' in localhost:
output_dir = '/private/huangchenxi/mmdet/outputs'
else:
output_dir = 'work_dirs'
if args.dir.endswith('-c'):
args.dir = args.dir[:-2]
args.resume_from = search_and_delete(os.path.join(
output_dir, args.dir),
prefix=cfg.work_dir,
suffix=localhost)
cfg.work_dir += time.strftime("_%m%d_%H%M") + '_' + localhost
cfg.work_dir = os.path.join(output_dir, args.dir, cfg.work_dir)
if args.workers_per_gpu != -1:
cfg.data['workers_per_gpu'] = args.workers_per_gpu
if args.resume_from is not None:
cfg.resume_from = args.resume_from
cfg.gpus = args.gpus
if args.profiler or args.speed:
cfg.data.imgs_per_gpu = 1
if cfg.resume_from or cfg.load_from:
cfg.model['pretrained'] = None
if args.test:
cfg.data.train['ann_file'] = cfg.data.val['ann_file']
cfg.data.train['img_prefix'] = cfg.data.val['img_prefix']
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
num_gpus = args.gpus
rank = 0
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
num_gpus = torch.cuda.device_count()
rank, _ = get_dist_info()
if cfg.optimizer['type'] == 'SGD':
cfg.optimizer['lr'] *= num_gpus * cfg.data.imgs_per_gpu / 256
else:
cfg.optimizer['lr'] *= ((num_gpus / 8) * (cfg.data.imgs_per_gpu / 2))
# init logger before other steps
logger = get_root_logger(nlogger, cfg.log_level)
if rank == 0:
logger.set_logger_dir(cfg.work_dir, 'd')
logger.info("Config: ------------------------------------------\n" +
cfg.text)
logger.info('Distributed training: {}'.format(distributed))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}'.format(args.seed))
set_random_seed(args.seed)
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
if rank == 0:
# describe_vars(model)
writer = set_writer(cfg.work_dir)
# try:
# # describe_features(model.backbone)
# writer.add_graph(model, torch.zeros((1, 3, 800, 800)))
# except (NotImplementedError, TypeError):
# logger.warn("Add graph failed.")
# except Exception as e:
# logger.warn("Add graph failed:", e)
if not args.graph and not args.profiler and not args.speed:
if distributed:
model = MMDistributedDataParallel(model.cuda())
else:
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
if isinstance(cfg.data.train, list):
for t in cfg.data.train:
logger.info("loading training set: " + str(t.ann_file))
train_dataset = [build_dataset(t) for t in cfg.data.train]
CLASSES = train_dataset[0].CLASSES
else:
logger.info("loading training set: " +
str(cfg.data.train.ann_file))
train_dataset = build_dataset(cfg.data.train)
logger.info("{} images loaded!".format(len(train_dataset)))
CLASSES = train_dataset.CLASSES
if cfg.checkpoint_config is not None:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmdet_version=__version__,
config=cfg.text,
CLASSES=CLASSES)
# add an attribute for visualization convenience
if hasattr(model, 'module'):
model.module.CLASSES = CLASSES
else:
model.CLASSES = CLASSES
train_detector(model,
train_dataset,
cfg,
distributed=distributed,
validate=args.validate,
logger=logger,
runner_attr_dict={'task_name': args.dir})
else:
from mmcv.runner.checkpoint import load_checkpoint
from mmdet.datasets import build_dataloader
from mmdet.core.utils.model_utils import register_hooks
from mmdet.apis.train import parse_losses
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
if args.profiler == 'test' or args.speed == 'test':
model.eval()
dataset = build_dataset(cfg.data.test)
else:
model.train()
dataset = build_dataset(cfg.data.train)
if cfg.load_from and (args.profiler or args.speed):
logger.info('load checkpoint from %s', cfg.load_from)
load_checkpoint(model,
cfg.load_from,
map_location='cpu',
strict=True)
data_loader = build_dataloader(dataset,
cfg.data.imgs_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False,
shuffle=False)
if args.graph:
id_dict = {}
for name, parameter in model.named_parameters():
id_dict[id(parameter)] = name
for i, data_batch in enumerate(data_loader):
if args.graph:
outputs = model(**data_batch)
loss, log_vars = parse_losses(outputs)
get_dot = register_hooks(loss, id_dict)
loss.backward()
dot = get_dot()
dot.save('graph.dot')
break
elif args.profiler:
with torch.autograd.profiler.profile(use_cuda=True) as prof:
if args.profiler == 'train':
outputs = model(**data_batch)
loss, log_vars = parse_losses(outputs)
loss.backward()
else:
with torch.no_grad():
model(**data_batch, return_loss=False)
if i == 20:
prof.export_chrome_trace('./trace.json')
logger.info(prof)
break
elif args.speed:
if args.speed == 'train':
start = time.perf_counter()
outputs = model(**data_batch)
loss, log_vars = parse_losses(outputs)
loss.backward()
torch.cuda.synchronize()
end = time.perf_counter()
else:
start = time.perf_counter()
with torch.no_grad():
model(**data_batch, return_loss=False)
end = time.perf_counter()
logger.info("{:.3f} s/iter, {:.1f} iters/s".format(
end - start, 1. / (end - start)))
def main(**kwargs):
args = parse_args()
for k, v in kwargs.items():
args.__setattr__(k, v)
assert args.out or args.show or args.json_out, \
('Please specify at least one operation (save or show the results) '
'with the argument "--out" or "--show" or "--json_out"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
if args.json_out is not None and args.json_out.endswith('.json'):
args.json_out = args.json_out[:-5]
if isinstance(args.config, str):
cfg = mmcv.Config.fromfile(args.config)
else:
cfg = args.config
# have first model?
first_model = None
if 'first_model_cfg' in cfg._cfg_dict and cfg.first_model_cfg is not None:
first_code_py = import_module(cfg.first_code_py)
first_model = first_code_py.Inference(cfg.first_model_cfg, cfg.first_model_path)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
if 'val' in cfg.data:
cfg.data.val.test_mode = True
if 'test' in cfg.data:
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
if args.mode == 'val':
dataset = build_dataset(cfg.data.val)
else:
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(
dataset,
imgs_per_gpu=args.imgs_per_gpu,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs, result_times = single_gpu_test(model, data_loader, args.show, first_model=first_model)
else:
model = MMDistributedDataParallel(model.cuda())
outputs, result_times = multi_gpu_test(model, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = args.out
coco_eval(result_file, eval_types, dataset.coco, classwise=True)
else:
if not isinstance(outputs[0], dict):
result_files = results2json(dataset, outputs, args.out)
coco_eval(result_files, eval_types, dataset.coco, classwise=True)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}'.format(name)
result_files = results2json(dataset, outputs_, result_file)
coco_eval(result_files, eval_types, dataset.coco, classwise=True)
# Save predictions in the COCO json format
defect_test_results = {}
if args.json_out and rank == 0:
if not isinstance(outputs[0], dict):
result_files = results2json(dataset, outputs, args.json_out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if 'ignore_ids' not in cfg.data[args.mode]:
cfg.data[args.mode]['ignore_ids'] = None
ignore_ids = cfg.data[args.mode]['ignore_ids']
coco_result = coco_eval(result_files, eval_types, dataset.coco, classwise=True, ignore_ids=ignore_ids)
threshold = cfg.test_cfg['rcnn']['score_thr']
defect_rst = defect_eval(result_files['bbox'], dataset.coco.dataset, result_times, threshold=threshold)
defect_test_results.update(log='\n'.join([coco_result['bbox']['log'], defect_rst['log']]))
defect_test_results.update(coco_result=coco_result['bbox']['data'])
defect_test_results.update(defect_result=defect_rst['data'])
else:
for name in outputs[0]:
outputs_ = [out[name] for out in outputs]
result_file = args.json_out + '.{}'.format(name)
results2json(dataset, outputs_, result_file)
return defect_test_results
def main():
args = parse_args()
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.val.test_mode = True
dataset = obj_from_dict(cfg.data.val, datasets, dict(test_mode=True))
if args.gpus == 1:
model = build_detector(cfg.model,
train_cfg=None,
test_cfg=cfg.test_cfg)
load_checkpoint(model, args.checkpoint)
model = MMDataParallel(model, device_ids=[0])
data_loader = build_dataloader(
dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
num_gpus=1,
dist=False,
shuffle=False)
# outputs is an array, each element is an image result, it may be resize to original image size
outputs = single_test(model, data_loader, args.show)
#print(outputs)
else:
model_args = cfg.model.copy()
model_args.update(train_cfg=None, test_cfg=cfg.test_cfg)
model_type = getattr(detectors, model_args.pop('type'))
outputs = parallel_test(model_type,
model_args,
args.checkpoint,
dataset,
_data_func,
range(args.gpus),
workers_per_gpu=args.proc_per_gpu)
if args.out:
print('writing results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = args.out
coco_eval(result_file, eval_types, dataset.coco)
else:
if not isinstance(outputs[0], dict):
result_file = args.out + '.json'
results2json(dataset, outputs, result_file)
coco_eval(result_file, eval_types, dataset.coco)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}.json'.format(name)
results2json(dataset, outputs_, result_file)
coco_eval(result_file, eval_types, dataset.coco)
cfg = Config.fromfile(config)
cfg.data.test.test_mode = True
distributed = False
#set device cpu or gpu
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
#build dataset
dataset = build_dataset(cfg.data.test)
print(cfg.data.test)
#build dataloader
data_loader = build_dataloader(dataset,
samples_per_gpu=1,
workers_per_gpu=1,
dist=distributed,
shuffle=False)
#build detector
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
#load weights
checkpoint = load_checkpoint(model, checkpoint,
map_location='cpu') # 'cuda:0'
model.CLASSES = dataset.CLASSES
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, False)
results = []
def main():
args = parse_args()
assert args.out or args.eval or args.format_only or args.show, \
('Please specify at least one operation (save/eval/format/show the '
'results) with the argument "--out", "--eval", "--format_only" '
'or "--show"')
if args.eval and args.format_only:
raise ValueError('--eval and --format_only cannot be both specified')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
imgs_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
if args.fuse_conv_bn:
model = fuse_module(model)
# old versions did not save class info in checkpoints, this walkaround is
# for backward compatibility
if ('meta' in checkpoint) and ('CLASSES' in checkpoint['meta']):
model.CLASSES = checkpoint['meta']['CLASSES']
else:
model.CLASSES = dataset.CLASSES
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show)
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
outputs = multi_gpu_test(model, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if rank == 0:
if args.out:
print('\nwriting results to {}'.format(args.out))
mmcv.dump(outputs, args.out)
kwargs = {} if args.options is None else args.options
if args.format_only:
dataset.format_results(outputs, **kwargs)
if args.eval:
dataset.evaluate(outputs, args.eval, **kwargs)
def main():
args = parse_args()
config_file1 = './swa/swa_cascade_rcnn_r50_rfp_sac_iou_alldata-v3_e15/swa_cascade_rcnn_r50_rfp_sac_iou_alldata-v3_e15.py' #
checkpoint_file1 = './swa/swa_cascade_rcnn_r50_rfp_sac_iou_alldata-v3_e15/swa_model_12.pth'
config_file2 = './swa/cascade_rcnn_r50_rfp_sac_iou_ls_alldata-v3_e15.py'
checkpoint_file2 = './swa/epoch_15.pth'
device = 'cuda:0'
cfg1 = Config.fromfile(config_file1)
cfg2 = Config.fromfile(config_file2)
# build model
# model1
model1 = build_detector(cfg1.model, test_cfg=cfg1.get('test_cfg'))
load_checkpoint(model1, checkpoint_file1, map_location=device)
# model2
model2 = build_detector(cfg2.model, test_cfg=cfg2.get('test_cfg'))
load_checkpoint(model2, checkpoint_file2, map_location=device)
test_json_raw = json.load(open(cfg1.data.test.ann_file))
imgid2name = {}
for imageinfo in test_json_raw['images']:
imgid = imageinfo['id']
imgid2name[imageinfo['file_name']] = imgid
wrap_fp16_model(model1) # 采用fp16加速预测
wrap_fp16_model(model2)
# build the dataloader
samples_per_gpu = cfg1.data.test.pop('samples_per_gpu',
1) # aug_test不支持batch_size>1
dataset = build_dataset(cfg1.data.test)
data_loader = build_dataloader(dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=4,
dist=False,
shuffle=False)
model1 = MMDataParallel(model1, device_ids=[0]) # 为啥加？(不加就错了)
model2 = MMDataParallel(model2, device_ids=[0])
model1.eval()
model2.eval()
json_results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result1 = model1(return_loss=False, rescale=True, **data)
result2 = model2(return_loss=False, rescale=True, **data)
batch_size = len(result1)
assert len(result1) == len(result2)
result1 = result1[0] # 每次只输入一张
result2 = result2[0]
img_metas = data['img_metas'][0].data[0]
img_shape = img_metas[0]['ori_shape']
bboxes, scores, labels = post_predictions(result1, img_shape)
e_bboxes, e_scores, e_labels = post_predictions(result2, img_shape)
bboxes_list = [bboxes, e_bboxes]
scores_list = [scores, e_scores]
labels_list = [labels, e_labels]
bboxes, scores, labels = weighted_boxes_fusion(bboxes_list,
scores_list,
labels_list,
weights=[1, 1],
iou_thr=0.6,
skip_box_thr=0.0001,
conf_type='max')
# basename = img_metas[0]['ori_filename']
# image = cv2.imread(os.path.join(cfg.data.test.img_prefix, basename))
for (box, score, label) in zip(bboxes, scores, labels):
xmin, ymin, xmax, ymax = box.tolist()
xmin, ymin, xmax, ymax = round(
float(xmin) * img_shape[1],
2), round(float(ymin) * img_shape[0],
2), round(float(xmax) * img_shape[1],
2), round(float(ymax) * img_shape[0], 2)
data = dict()
data['image_id'] = imgid2name[img_metas[0]['ori_filename']]
data['bbox'] = [xmin, ymin, xmax - xmin, ymax - ymin]
data['score'] = float(score)
data['category_id'] = label + 1
json_results.append(data)
for _ in range(batch_size):
prog_bar.update()
mmcv.dump(json_results, args.jsonfile)
