def main(args):
if args.dataset:
dataset = get_dataset(args.dataset, False, False, None)
classes_num = dataset.classes_num
else:
classes_num = args.classes_num
net = locate(args.model)(classes_num, fc_detection_head=False)
load_checkpoint(net, args.ckpt)
if args.show_flops:
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
printable_graph = onnx_export(net, args.input_size, args.output_file, check=args.check, verbose=args.verbose)
if args.verbose:
logging.info(printable_graph)
if args.show_flops:
net.stop_flops_count()
logging.info('Computational complexity: {}'.format(flops_to_string(net.compute_average_flops_cost())))
if args.verbose:
print_model_with_flops(net)
def main(args):
transforms = Compose([
Resize(max_size=args.fit_max_image_size,
window_size=args.fit_window_size,
size=args.size),
ToTensor(),
Normalize(mean=args.mean_pixel, std=[1., 1., 1.], rgb=args.rgb),
])
dataset = get_dataset(args.dataset, False, False, transforms)
logging.info(dataset)
batch_size = 1
logging.info('Using {} backend'.format(args.backend))
logging.info('Loading network...')
if args.backend == 'pytorch':
net = locate(args.pytorch_model_class)(dataset.classes_num)
net.eval()
load_checkpoint(net, args.checkpoint_file_path)
if torch.cuda.is_available():
net = net.cuda()
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
elif args.backend == 'openvino':
net = MaskRCNNOpenVINO(
args.openvino_model_path,
args.checkpoint_file_path,
device=args.device,
plugin_dir=args.plugin_dir,
cpu_extension_lib_path=args.cpu_extension,
collect_perf_counters=args.show_performance_counters)
else:
raise ValueError('Unknown backend "{}"'.format(args.backend))
viz = Visualizer(dataset.classes,
confidence_threshold=args.prob_threshold,
show_boxes=args.show_boxes,
show_scores=args.show_scores)
inference_timer = Timer(cuda_sync=True, warmup=1)
timer = Timer(cuda_sync=False, warmup=1)
timer.tic()
logging.info('Configuring data source...')
if args.video:
try:
args.video = int(args.video)
except ValueError:
pass
demo_dataset = VideoDataset(args.video,
labels=dataset.classes,
transforms=transforms)
num_workers = 0
tracker = StaticIOUTracker()
else:
demo_dataset = ImagesDataset(args.images,
labels=dataset.classes,
transforms=transforms)
num_workers = 1
tracker = None
data_loader = torch.utils.data.DataLoader(demo_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False,
collate_fn=collate)
logging.info('Processing data...')
frames_num = len(demo_dataset)
for data_batch in tqdm(
iter(data_loader),
total=frames_num if frames_num != sys.maxsize else 0):
im_data = data_batch['im_data']
im_info = data_batch['im_info']
if torch.cuda.is_available():
im_data = [i.cuda() for i in im_data]
im_info = [i.cuda() for i in im_info]
with torch.no_grad(), inference_timer:
boxes, classes, scores, _, masks = net(im_data, im_info)
meta = data_batch['meta'][0]
scores, classes, boxes, masks = postprocess(
scores,
classes,
boxes,
masks,
im_h=meta['original_size'][0],
im_w=meta['original_size'][1],
im_scale_y=meta['processed_size'][0] / meta['original_size'][0],
im_scale_x=meta['processed_size'][1] / meta['original_size'][1],
full_image_masks=True,
encode_masks=False,
confidence_threshold=args.prob_threshold)
masks_ids = tracker(masks, classes) if tracker is not None else None
image = data_batch['original_image'][0]
visualization = viz(image,
boxes,
classes,
scores,
segms=masks,
ids=masks_ids)
fps = 1 / timer.toc()
if args.show_fps:
visualization = cv2.putText(visualization,
'FPS: {:>2.2f}'.format(fps), (30, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 0, 255), 2)
cv2.imshow('result', visualization)
key = cv2.waitKey(args.delay)
if key == 27:
break
timer.tic()
if inference_timer.average_time > 0:
logging.info('Average inference FPS: {:3.2f}'.format(
1 / inference_timer.average_time))
if args.backend == 'pytorch':
net.stop_flops_count()
if args.show_flops:
logging.info('Average FLOPs: {}'.format(
flops_to_string(net.compute_average_flops_cost())))
if args.show_layers_flops:
logging.info('Thorough computational complexity statistics:')
print_model_with_flops(net)
if torch.cuda.is_available():
logging.info('GPU memory footprint:')
logging.info('\tMax allocated: {:.2f} MiB'.format(
torch.cuda.max_memory_allocated() / 1024**2))
logging.info('\tMax cached: {:.2f} MiB'.format(
torch.cuda.max_memory_cached() / 1024**2))
else:
if args.show_performance_counters:
net.print_performance_counters()
cv2.destroyAllWindows()
del net
def __init__(self):
super().__init__()
self.identifier = 'instance-segmentation-security-0050'
self.description = 'Training of instance-segmentation-security-0050'
self.root_directory = osp.join(osp.dirname(osp.abspath(__file__)), '..')
self.run_directory = self.create_run_directory(osp.join(self.root_directory, 'outputs'))
setup_logging(file_path=osp.join(self.run_directory, 'log.txt'))
logger.info('Running {}'.format(self.identifier))
logger.info(self.description)
logger.info('Working directory "{}"'.format(self.run_directory))
self.batch_size = 32
self.virtual_iter_size = 1
# Training dataset.
training_transforms = Compose(
[
RandomResize(mode='size', heights=(416, 448, 480, 512, 544), widths=(416, 448, 480, 512, 544)),
RandomHorizontalFlip(prob=0.5),
ToTensor(),
Normalize(mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], rgb=False),
],
)
training_dataset_name = 'coco_2017_train'
logger.info('Training dataset {}'.format(training_dataset_name))
training_dataset = get_dataset(training_dataset_name, True, True, training_transforms)
logger.info(training_dataset)
self.training_data_loader = torch.utils.data.DataLoader(
training_dataset, batch_size=self.batch_size, num_workers=0,
shuffle=True, drop_last=True, collate_fn=collate
)
# Validation datasets.
validation_transforms = Compose(
[
Resize(size=[480, 480]),
ToTensor(),
Normalize(mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], rgb=False),
]
)
validation_datasets = []
validation_dataset_name = 'coco_2017_val'
logger.info('Validation dataset #{}: {}'.format(len(validation_datasets) + 1, validation_dataset_name))
validation_datasets.append(get_dataset(validation_dataset_name, False, False, validation_transforms))
logger.info(validation_datasets[-1])
self.validation_data_loaders = []
for validation_dataset in validation_datasets:
self.validation_data_loaders.append(torch.utils.data.DataLoader(
validation_dataset,
batch_size=1, num_workers=8,
shuffle=False, drop_last=False, collate_fn=collate)
)
self.validate_every = 10000
for validation_dataset in validation_datasets:
assert training_dataset.classes_num == validation_dataset.classes_num
# Model and optimizer.
logger.info('Model:')
self.model = Model(training_dataset.classes_num)
logger.info(self.model)
self.training_iterations_num = 270000
lr_scheduler_milestones = [220000, 250000]
base_lr = 0.02
weight_decay = 0.0001
logger.info('Optimizer:')
self.optimizer = torch.optim.SGD(self.setup_optimizer(self.model, base_lr, weight_decay),
lr=base_lr, weight_decay=weight_decay, momentum=0.9)
logger.info(self.optimizer)
logger.info('Learning Rate scheduler:')
self.lr_scheduler = MultiStepLRWithWarmUp(
self.optimizer,
milestones=lr_scheduler_milestones,
warmup_iters=1000,
warmup_method='linear',
warmup_factor_base=0.333,
gamma=0.1,
last_epoch=0
)
logger.info(self.lr_scheduler)
self.start_step = 0
checkpoint_file_path = osp.join(self.root_directory, 'data', 'pretrained_models',
'converted', 'imagenet', 'detectron', 'resnet50.pth')
if not osp.exists(checkpoint_file_path):
raise IOError('Initial checkpoint file "{}" does not exist. '
'Please fetch pre-trained backbone networks using '
'tools/download_pretrained_weights.py script first.'.format(checkpoint_file_path))
logger.info('Loading weights from "{}"'.format(checkpoint_file_path))
load_checkpoint(self.model.backbone, checkpoint_file_path)
# Loggers and misc. stuff.
self.loggers = [TextLogger(logger),
TensorboardLogger(self.run_directory)]
self.log_every = 50
self.checkpoint_every = 10000
def main(args):
transforms = Compose([
Resize(max_size=args.fit_max_image_size,
window_size=args.fit_window_size,
size=args.size),
ToTensor(),
Normalize(mean=args.mean_pixel, std=[1., 1., 1.], rgb=args.rgb),
])
dataset = get_dataset(args.dataset, False, False, transforms)
logging.info(dataset)
num_workers = args.num_workers
inference_timer = Timer()
logging.info('Using {} backend'.format(args.backend))
logging.info('Loading network...')
if args.backend == 'pytorch':
net = locate(args.pytorch_model_class)(dataset.classes_num,
force_max_output_size=False)
net.eval()
load_checkpoint(net, args.checkpoint_file_path, verbose=True)
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
if torch.cuda.is_available():
torch.backends.cudnn.deterministic = True
net = net.cuda()
net = ShallowDataParallel(net)
batch_size = args.batch_size
elif args.backend == 'openvino':
net = MaskRCNNOpenVINO(
args.openvino_model_path,
args.checkpoint_file_path,
device=args.device,
plugin_dir=args.plugin_dir,
cpu_extension_lib_path=args.cpu_extension,
collect_perf_counters=args.show_performance_counters)
batch_size = 1
else:
raise ValueError('Unknown backend "{}"'.format(args.backend))
logging.info('Using batch size {}'.format(batch_size))
logging.info('Number of prefetching processes {}'.format(num_workers))
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False,
collate_fn=collate)
logging.info('Processing dataset...')
boxes_all = []
masks_all = []
classes_all = []
scores_all = []
for data_batch in tqdm(iter(data_loader)):
batch_meta = data_batch['meta']
actual_batch_size = len(batch_meta)
with torch.no_grad(), inference_timer:
boxes, classes, scores, batch_ids, masks = net(**data_batch)
im_heights = [meta['original_size'][0] for meta in batch_meta]
im_widths = [meta['original_size'][1] for meta in batch_meta]
im_scale_y = [
meta['processed_size'][0] / meta['original_size'][0]
for meta in batch_meta
]
im_scale_x = [
meta['processed_size'][1] / meta['original_size'][1]
for meta in batch_meta
]
scores, classes, boxes, masks = postprocess_batch(
batch_ids,
scores,
classes,
boxes,
masks,
actual_batch_size,
im_h=im_heights,
im_w=im_widths,
im_scale_y=im_scale_y,
im_scale_x=im_scale_x,
full_image_masks=True,
encode_masks=True)
boxes_all.extend(boxes)
masks_all.extend(masks)
classes_all.extend(classes)
scores_all.extend(scores)
try:
del data_loader
except ConnectionResetError:
pass
logging.info('Evaluating results...')
evaluation_results = dataset.evaluate(scores_all, classes_all, boxes_all,
masks_all)
logging.info(evaluation_results)
logging.info('Average inference time {}'.format(
inference_timer.average_time))
if args.backend == 'pytorch':
if torch.cuda.is_available():
net = net.module
net.stop_flops_count()
if args.show_flops:
logging.info('Average FLOPs: {}'.format(
flops_to_string(net.compute_average_flops_cost())))
if args.show_layers_flops:
logging.info('Thorough computational complexity statistics:')
print_model_with_flops(net)
else:
if args.show_performance_counters:
net.print_performance_counters()
del net
def main(args):
num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
train_tool = DefaultMaskRCNNTrainingEngine()
train_tool.identifier = 'MaskRCNN'
train_tool.description = ''
train_tool.set_random_seeds()
train_tool.root_directory = osp.join(args.output_dir, train_tool.identifier)
train_tool.run_directory = train_tool.create_run_directory(train_tool.root_directory)
train_tool.batch_size = args.batch_size
train_tool.virtual_iter_size = \
train_tool.adjust_virtual_iteration_size(num_gpus, args.batch_size, args.batch_size_per_gpu)
transforms = Compose(
[
Resize(max_size=args.max_image_size),
ToTensor(),
Normalize(mean=args.mean_pixel, std=[1., 1., 1.], rgb=False),
]
)
if args.dataset == 'coco2017':
train_dataset_name = 'coco_2017_train'
val_dataset_name = 'coco_2017_val'
else:
raise ValueError('Invalid dataset name "{}"'.format(args.dataset))
train_dataset = get_dataset(train_dataset_name, True, True, transforms)
val_dataset = get_dataset(val_dataset_name, False, False, transforms)
assert train_dataset.classes_num == val_dataset.classes_num
train_tool.training_data_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size // train_tool.virtual_iter_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=collate
)
train_tool.validation_data_loaders = [
torch.utils.data.DataLoader(
val_dataset,
batch_size=1,
num_workers=args.num_workers,
shuffle=False,
collate_fn=collate
)
]
train_tool.validate_every = args.test_interval
train_tool.model = locate(args.model)(train_dataset.classes_num)
train_tool.training_iterations_num = args.max_iter
train_tool.lr_scheduler_milestones = args.drop_lr
params = train_tool.setup_optimizer(train_tool.model, args.lr, args.weight_decay)
train_tool.optimizer = torch.optim.SGD(params, momentum=args.momentum)
start_step = 0
if args.load_ckpt or args.load_backbone:
start_step = train_tool.load_checkpoint(train_tool.model, train_tool.optimizer,
args.load_ckpt, args.load_backbone, args.resume)
train_tool.lr_scheduler = MultiStepLRWithWarmUp(
train_tool.optimizer,
milestones=args.drop_lr,
warmup_iters=args.warmup_iters,
warmup_method=args.warmup_method,
warmup_factor_base=args.warmup_factor,
gamma=args.gamma,
last_epoch=start_step
)
text_log = TextLogger(logging.getLogger(__name__))
tensorboard_log = TensorboardLogger(train_tool.run_directory)
train_tool.loggers = [text_log, tensorboard_log]
train_tool.log_every = args.display_interval
train_tool.checkpoint_every = args.checkpoint_interval
# Begin training
train_tool.run(start_step)