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
'ONNX file is saved to {}'.format(output_onnx_file))
# Convert to IR.
logging.info('Converting to IR...')
output_ir_dir = osp.join(args.output_dir, 'ir',
target['dst_file'])
mkdir_for_file(output_ir_dir)
output_ir_dir = osp.dirname(output_ir_dir)
status = call([
args.model_optimizer, '--framework', 'onnx',
'--input_model', output_onnx_file, '--output_dir',
output_ir_dir, '--input', 'im_data,im_info',
'--output', 'boxes,scores,classes,batch_ids,raw_masks',
'--mean_values', 'im_data{},im_info[0,0,0]'.format(
str(target['mean_pixel']).replace(' ', ''))
])
if status:
logging.warning('Failed to convert model to IR.')
else:
logging.info(
'IR files saved to {}'.format(output_ir_dir))
except Exception as ex:
logging.warning(repr(ex))
if __name__ == '__main__':
setup_logging()
args = parse_args()
main(args)
def __init__(self, work_dir, config):
super().__init__()
self.identifier = config['identifier']
self.description = config['description']
self.root_directory = work_dir if work_dir else osp.join(
osp.dirname(osp.abspath(__file__)), '..')
self.run_directory = self.create_run_directory(
osp.join(self.root_directory, 'models'))
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 = config['training_details']['batch_size']
self.virtual_iter_size = config['training_details'][
'virtual_iter_size']
model_class = make_text_detector(**config['model'])
alphabet_decoder = AlphabetDecoder()
# Training dataset.
training_transforms = Compose([
getattr(sys.modules[__name__], k)(**v)
for k, v in config['training_transforms'].items()
] + [AlphabetDecodeTransform(alphabet_decoder)])
training_dataset_name = config['training_dataset_name']
logger.info('Training dataset {}'.format(training_dataset_name))
training_dataset = get_dataset(training_dataset_name,
True,
True,
training_transforms,
alphabet_decoder=alphabet_decoder,
remove_images_without_text=True)
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([
getattr(sys.modules[__name__], k)(**v)
for k, v in config['validation_transforms'].items()
])
self.confidence_threshold = config['validation_confidence_threshold']
validation_datasets = []
validation_dataset_name = config['validation_dataset_name']
logger.info('Validation dataset #{}: {}'.format(
len(validation_datasets) + 1, validation_dataset_name))
validation_datasets.append(
get_dataset(validation_dataset_name,
False,
False,
validation_transforms,
alphabet_decoder=alphabet_decoder))
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 = config['training_details']['validate_every']
for validation_dataset in validation_datasets:
assert training_dataset.classes_num == validation_dataset.classes_num
# Model and optimizer.
logger.info('Model:')
self.model = model_class(cls_num=training_dataset.classes_num,
shape=config['shape'],
num_chars=len(alphabet_decoder.alphabet))
logger.info(self.model)
self.training_iterations_num = config['training_details'][
'training_iterations_num']
lr_scheduler_milestones = config['training_details'][
'lr_scheduler_milestones']
base_lr = config['training_details']['base_lr']
weight_decay = config['training_details']['weight_decay']
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
if 'backbone_checkpoint' in config and config['backbone_checkpoint']:
checkpoint_file_path = osp.join(self.root_directory,
config['backbone_checkpoint'])
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,
verbose=True,
skip_prefix='text_recogn')
if 'checkpoint' in config and config['checkpoint']:
checkpoint_file_path = osp.join(self.root_directory,
config['checkpoint'])
if not osp.exists(checkpoint_file_path):
raise IOError('Checkpoint file "{}" does not exist. '.format(
checkpoint_file_path))
logger.info(
'Loading weights from "{}"'.format(checkpoint_file_path))
load_checkpoint(self.model, checkpoint_file_path, verbose=True)
# Loggers and misc. stuff.
self.loggers = [
TextLogger(logger),
TensorboardLogger(self.run_directory)
]
self.log_every = 50
self.checkpoint_every = config['training_details']['checkpoint_every']