def train_detector(model,
dataset,
cfg,
validate=False,
timestamp=None,
meta=None):
logger = Logging.getLogger()
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [build_dataloader(ds, data=cfg.data) for ds in dataset]
device = select_device(cfg.device)
model = model.cuda(device)
model.device = device
if device.type != 'cpu' and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# build runner
optimizer = cfg.optimizer
if 'ema' in cfg:
ema = cfg.ema
else:
ema = None
runner = Runner(model,
batch_processor,
optimizer,
cfg.work_dir,
logger=logger,
meta=meta,
ema=ema)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# register eval hooks 需要放在日志前面，不然打印不出日志。
if validate:
cfg.data.val.train = False
val_dataset = build_from_dict(cfg.data.val, DATASET)
val_dataloader = build_dataloader(val_dataset,
shuffle=False,
data=cfg.data)
eval_cfg = cfg.get('evaluation', {})
from yolodet.models.hooks.eval_hook import EvalHook
runner.register_hook(EvalHook(val_dataloader, **eval_cfg))
# register hooks
# runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,cfg.checkpoint_config)
runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,
cfg.checkpoint_config, cfg.log_config)
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 inference(config,
checkpoint,
img_path,
device='cuda:0',
half=False,
augment=False,
scores_thr=0.3,
merge=False,
save_json=False,
save_file=False,
save_path='',
show=False):
device = select_device(device)
model = init_detector(config, checkpoint=checkpoint, device=device)
t0 = time.time()
if os.path.isdir(img_path):
imgs_paths = get_file_realpath(
img_path, *[
".jpg", ".JPG", ".jpeg", ".JPEG", ".png", ".PNG", ".bmp",
".BMP"
])
for i, img in enumerate(imgs_paths):
basename = os.path.basename(img)
result, t = inference_detector(model,
img=img,
scores_thr=scores_thr,
augment=augment,
half=half,
merge=merge)
#print(result)
# Print time (inference + NMS)
print('%s Done , object num : %s .time:(%.3fs)' %
(basename, len(result), t))
show_result(img,
result,
model.CLASSES,
show=show,
save_json=save_json,
save_file=save_file,
out_path=save_path,
file_name=basename)
else:
basename = os.path.basename(img_path)
t1 = torch_utils.time_synchronized()
result = inference_detector(model,
img=img_path,
scores_thr=scores_thr,
augment=augment,
half=half,
merge=merge)
t2 = torch_utils.time_synchronized()
# print(result)
# Print time (inference + NMS)
print('%s Done , object num : %s .time:(%.3fs)' %
(basename, len(result), t2 - t1))
show_result(img_path,
result,
model.CLASSES,
show=show,
save_json=save_json,
save_file=save_file,
out_path=save_path,
file_name=basename)
print('Done. (%.3fs)' % (time.time() - t0))
parser.add_argument('--half',
action='store_true',
help='fp16 half precision')
opt = parser.parse_args()
print(opt)
# config = '/disk2/project/mmdetection/mount/pytorch-YOLOv4/cfg/yolov5_coco_gpu.py'
# checkpoint = '/disk2/project/pytorch-YOLOv4/work_dirs/yolov5-l_epoch_24.pth'
cfg = Config.fromfile(opt.config)
cfg.data.val.train = False
val_dataset = build_from_dict(cfg.data.val, DATASET)
val_dataloader = build_dataloader(val_dataset,
data=cfg.data,
shuffle=False)
device = select_device(opt.device)
# model = init_detector(opt.config, checkpoint=opt.checkpoint, device=device)
model = init_detector(opt.config, checkpoint=opt.checkpoint, device=device)
result = single_gpu_test(model,
val_dataloader,
half=opt.half,
conf_thres=opt.conf_thres,
iou_thres=opt.iou_thres,
merge=opt.merge,
save_json=opt.save_json,
augment=opt.augment,
verbose=opt.verbose,
coco_val_path=opt.coco_val_path)
print(result)
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.work_dir is not None:
cfg.work_dir = args.work_dir
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.device is not None:
cfg.device = args.device
else:
cfg.device = None
device = select_device(cfg.device)
if args.autoscale_lr:
# apply the linear scaling rule (https://arxiv.org/abs/1706.02677)
cfg.optimizer['lr'] = cfg.optimizer['lr'] * len(cfg.gpu_ids) / 8
# create work_dir
file_utils.mkdir_or_exist(osp.abspath(cfg.work_dir))
# init the logger before other steps
# timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
# log_file = osp.join(cfg.work_dir, '{}.log'.format(timestamp))
logger = Logging.getLogger()
# 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([('{}: {}'.format(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
meta['batch_size'] = cfg.data.batch_size
meta['subdivisions'] = cfg.data.subdivisions
meta['multi_scale'] = args.multi_scale
# log some basic info
logger.info('Config:\n{}'.format(cfg.text))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}, deterministic: {}'.format(
args.seed, args.deterministic))
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
model = build_from_dict(cfg.model, DETECTORS)
model = model.cuda(device)
# model.device = device
if device.type != 'cpu' and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.device = device
datasets = [build_from_dict(cfg.data.train, DATASET)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_from_dict(val_dataset, DATASET))
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(config=cfg.text,
CLASSES=datasets[0].CLASSES)
# add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
train_detector(model,
datasets,
cfg,
validate=args.validate,
timestamp=timestamp,
meta=meta)