def save_air(args):
'''save air'''
print('============= yolov3 start save air ==================')
devid = int(os.getenv('DEVICE_ID', '0')) if args.run_platform != 'CPU' else 0
context.set_context(mode=context.GRAPH_MODE, device_target=args.run_platform, save_graphs=False, device_id=devid)
num_classes = config.num_classes
anchors_mask = config.anchors_mask
num_anchors_list = [len(x) for x in anchors_mask]
network = backbone_HwYolov3(num_classes, num_anchors_list, args)
if os.path.isfile(args.pretrained):
param_dict = load_checkpoint(args.pretrained)
param_dict_new = {}
for key, values in param_dict.items():
if key.startswith('moments.'):
continue
elif key.startswith('network.'):
param_dict_new[key[8:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
print('load model {} success'.format(args.pretrained))
input_data = np.random.uniform(low=0, high=1.0, size=(args.batch_size, 3, 448, 768)).astype(np.float32)
tensor_input_data = Tensor(input_data)
export(network, tensor_input_data,
file_name=args.pretrained.replace('.ckpt', '_' + str(args.batch_size) + 'b.air'), file_format='AIR')
print("export model success.")
def define_network(args):
"""Define train network with TrainOneStepCell."""
# backbone and loss
num_classes = args.num_classes
num_anchors_list = args.num_anchors_list
anchors = args.anchors
anchors_mask = args.anchors_mask
momentum = args.momentum
args.logger.info('train opt momentum:{}'.format(momentum))
weight_decay = args.weight_decay * float(args.batch_size)
args.logger.info('real weight_decay:{}'.format(weight_decay))
lr_scale = args.world_size / 8
args.logger.info('lr_scale:{}'.format(lr_scale))
args.lr = warmup_step_new(args, lr_scale=lr_scale)
network = backbone_HwYolov3(num_classes, num_anchors_list, args)
criterion0 = YoloLoss(num_classes, anchors, anchors_mask[0], 64, 0, head_idx=0.0)
criterion1 = YoloLoss(num_classes, anchors, anchors_mask[1], 32, 0, head_idx=1.0)
criterion2 = YoloLoss(num_classes, anchors, anchors_mask[2], 16, 0, head_idx=2.0)
# load pretrain model
if os.path.isfile(args.pretrained):
param_dict = load_checkpoint(args.pretrained)
param_dict_new = {}
for key, values in param_dict.items():
if key.startswith('moments.'):
continue
elif key.startswith('network.'):
param_dict_new[key[8:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
args.logger.info('load model {} success'.format(args.pretrained))
train_net = BuildTrainNetworkV2(network, criterion0, criterion1, criterion2, args)
# optimizer
opt = nn.Momentum(params=train_net.trainable_params(), learning_rate=Tensor(args.lr), momentum=momentum,
weight_decay=weight_decay)
# package training process
if args.use_loss_scale:
train_net = TrainOneStepWithLossScaleCell(train_net, opt)
else:
train_net = nn.TrainOneStepCell(train_net, opt)
if args.world_size != 1:
train_net.set_broadcast_flag()
return train_net
def val(args):
'''eval'''
print('=============yolov3 start evaluating==================')
# logger
args.batch_size = config.batch_size
args.input_shape = config.input_shape
args.result_path = config.result_path
args.conf_thresh = config.conf_thresh
args.nms_thresh = config.nms_thresh
context.set_auto_parallel_context(parallel_mode=ParallelMode.STAND_ALONE,
device_num=args.world_size,
gradients_mean=True)
mindrecord_path = args.mindrecord_path
print('Loading data from {}'.format(mindrecord_path))
num_classes = config.num_classes
if num_classes > 1:
raise NotImplementedError(
'num_classes > 1: Yolov3 postprocess not implemented!')
anchors = config.anchors
anchors_mask = config.anchors_mask
num_anchors_list = [len(x) for x in anchors_mask]
reduction_0 = 64.0
reduction_1 = 32.0
reduction_2 = 16.0
labels = ['face']
classes = {0: 'face'}
# dataloader
ds = de.MindDataset(
mindrecord_path + "0",
columns_list=["image", "annotation", "image_name", "image_size"])
single_scale_trans = SingleScaleTrans(resize=args.input_shape)
ds = ds.batch(
args.batch_size,
per_batch_map=single_scale_trans,
input_columns=["image", "annotation", "image_name", "image_size"],
num_parallel_workers=8)
args.steps_per_epoch = ds.get_dataset_size()
# backbone
network = backbone_HwYolov3(num_classes, num_anchors_list, args)
# load pretrain model
if os.path.isfile(args.pretrained):
param_dict = load_checkpoint(args.pretrained)
param_dict_new = {}
for key, values in param_dict.items():
if key.startswith('moments.'):
continue
elif key.startswith('network.'):
param_dict_new[key[8:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
print('load model {} success'.format(args.pretrained))
else:
print(
'load model {} failed, please check the path of model, evaluating end'
.format(args.pretrained))
exit(0)
ds = ds.repeat(1)
det = {}
img_size = {}
img_anno = {}
model_name = args.pretrained.split('/')[-1].replace('.ckpt', '')
result_path = os.path.join(args.result_path, model_name)
if os.path.exists(result_path):
pass
if not os.path.isdir(result_path):
os.makedirs(result_path, exist_ok=True)
# result file
ret_files_set = {
'face': os.path.join(result_path, 'comp4_det_test_face_rm5050.txt'),
}
test_net = BuildTestNetwork(network, reduction_0, reduction_1, reduction_2,
anchors, anchors_mask, num_classes, args)
print('conf_thresh:', args.conf_thresh)
eval_times = 0
for data in ds.create_tuple_iterator(output_numpy=True):
batch_images = data[0]
batch_labels = data[1]
batch_image_name = data[2]
batch_image_size = data[3]
eval_times += 1
img_tensor = Tensor(batch_images, mstype.float32)
dets = []
tdets = []
coords_0, cls_scores_0, coords_1, cls_scores_1, coords_2, cls_scores_2 = test_net(
img_tensor)
boxes_0, boxes_1, boxes_2 = get_bounding_boxes(
coords_0, cls_scores_0, coords_1, cls_scores_1, coords_2,
cls_scores_2, args.conf_thresh, args.input_shape, num_classes)
converted_boxes_0, converted_boxes_1, converted_boxes_2 = tensor_to_brambox(
boxes_0, boxes_1, boxes_2, args.input_shape, labels)
tdets.append(converted_boxes_0)
tdets.append(converted_boxes_1)
tdets.append(converted_boxes_2)
batch = len(tdets[0])
for b in range(batch):
single_dets = []
for op in range(3):
single_dets.extend(tdets[op][b])
dets.append(single_dets)
det.update({
batch_image_name[k].decode('UTF-8'): v
for k, v in enumerate(dets)
})
img_size.update({
batch_image_name[k].decode('UTF-8'): v
for k, v in enumerate(batch_image_size)
})
img_anno.update({
batch_image_name[k].decode('UTF-8'): v
for k, v in enumerate(batch_labels)
})
print('eval times:', eval_times)
print('batch size: ', args.batch_size)
netw, neth = args.input_shape
reorg_dets = voc_wrapper.reorg_detection(det, netw, neth, img_size)
voc_wrapper.gen_results(reorg_dets, result_path, img_size, args.nms_thresh)
# compute mAP
ground_truth = parse_gt_from_anno(img_anno, classes)
ret_list = parse_rets(ret_files_set)
iou_thr = 0.5
evaluate = calc_recall_presicion_ap(ground_truth, ret_list, iou_thr)
aps_str = ''
for cls in evaluate:
per_line, = plt.plot(evaluate[cls]['recall'],
evaluate[cls]['presicion'], 'b-')
per_line.set_label('%s:AP=%.3f' % (cls, evaluate[cls]['ap']))
aps_str += '_%s_AP_%.3f' % (cls, evaluate[cls]['ap'])
plt.plot([i / 1000.0 for i in range(1, 1001)],
[i / 1000.0 for i in range(1, 1001)], 'y--')
plt.axis([0, 1.2, 0, 1.2])
plt.xlabel('recall')
plt.ylabel('precision')
plt.grid()
plt.legend()
plt.title('PR')
# save mAP
ap_save_path = os.path.join(
result_path,
result_path.replace('/', '_') + aps_str + '.png')
print('Saving {}'.format(ap_save_path))
plt.savefig(ap_save_path)
print('=============yolov3 evaluating finished==================')
def train(args):
'''train'''
print('=============yolov3 start trainging==================')
# init distributed
if args.world_size != 1:
init()
args.local_rank = get_rank()
args.world_size = get_group_size()
args.batch_size = config.batch_size
args.warmup_lr = config.warmup_lr
args.lr_rates = config.lr_rates
args.lr_steps = config.lr_steps
args.gamma = config.gamma
args.weight_decay = config.weight_decay
args.momentum = config.momentum
args.max_epoch = config.max_epoch
args.log_interval = config.log_interval
args.ckpt_path = config.ckpt_path
args.ckpt_interval = config.ckpt_interval
args.outputs_dir = os.path.join(args.ckpt_path, datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
print('args.outputs_dir', args.outputs_dir)
args.logger = get_logger(args.outputs_dir, args.local_rank)
if args.world_size != 8:
args.lr_steps = [i * 8 // args.world_size for i in args.lr_steps]
if args.world_size == 1:
args.weight_decay = 0.
if args.world_size != 1:
parallel_mode = ParallelMode.DATA_PARALLEL
else:
parallel_mode = ParallelMode.STAND_ALONE
context.set_auto_parallel_context(parallel_mode=parallel_mode, device_num=args.world_size, gradients_mean=True)
mindrecord_path = args.mindrecord_path
num_classes = config.num_classes
anchors = config.anchors
anchors_mask = config.anchors_mask
num_anchors_list = [len(x) for x in anchors_mask]
momentum = args.momentum
args.logger.info('train opt momentum:{}'.format(momentum))
weight_decay = args.weight_decay * float(args.batch_size)
args.logger.info('real weight_decay:{}'.format(weight_decay))
lr_scale = args.world_size / 8
args.logger.info('lr_scale:{}'.format(lr_scale))
# dataloader
args.logger.info('start create dataloader')
epoch = args.max_epoch
ds = de.MindDataset(mindrecord_path + "0", columns_list=["image", "annotation"], num_shards=args.world_size,
shard_id=args.local_rank)
ds = ds.map(input_columns=["image", "annotation"],
output_columns=["image", "annotation", 'coord_mask_0', 'conf_pos_mask_0', 'conf_neg_mask_0',
'cls_mask_0', 't_coord_0', 't_conf_0', 't_cls_0', 'gt_list_0', 'coord_mask_1',
'conf_pos_mask_1', 'conf_neg_mask_1', 'cls_mask_1', 't_coord_1', 't_conf_1',
't_cls_1', 'gt_list_1', 'coord_mask_2', 'conf_pos_mask_2', 'conf_neg_mask_2',
'cls_mask_2', 't_coord_2', 't_conf_2', 't_cls_2', 'gt_list_2'],
column_order=["image", "annotation", 'coord_mask_0', 'conf_pos_mask_0', 'conf_neg_mask_0',
'cls_mask_0', 't_coord_0', 't_conf_0', 't_cls_0', 'gt_list_0', 'coord_mask_1',
'conf_pos_mask_1', 'conf_neg_mask_1', 'cls_mask_1', 't_coord_1', 't_conf_1',
't_cls_1', 'gt_list_1', 'coord_mask_2', 'conf_pos_mask_2', 'conf_neg_mask_2',
'cls_mask_2', 't_coord_2', 't_conf_2', 't_cls_2', 'gt_list_2'],
operations=compose_map_func, num_parallel_workers=16, python_multiprocessing=True)
ds = ds.batch(args.batch_size, drop_remainder=True, num_parallel_workers=8)
args.steps_per_epoch = ds.get_dataset_size()
lr = warmup_step_new(args, lr_scale=lr_scale)
ds = ds.repeat(epoch)
args.logger.info('args.steps_per_epoch:{}'.format(args.steps_per_epoch))
args.logger.info('args.world_size:{}'.format(args.world_size))
args.logger.info('args.local_rank:{}'.format(args.local_rank))
args.logger.info('end create dataloader')
args.logger.save_args(args)
args.logger.important_info('start create network')
create_network_start = time.time()
# backbone and loss
network = backbone_HwYolov3(num_classes, num_anchors_list, args)
criterion0 = YoloLoss(num_classes, anchors, anchors_mask[0], 64, 0, head_idx=0.0)
criterion1 = YoloLoss(num_classes, anchors, anchors_mask[1], 32, 0, head_idx=1.0)
criterion2 = YoloLoss(num_classes, anchors, anchors_mask[2], 16, 0, head_idx=2.0)
# load pretrain model
if os.path.isfile(args.pretrained):
param_dict = load_checkpoint(args.pretrained)
param_dict_new = {}
for key, values in param_dict.items():
if key.startswith('moments.'):
continue
elif key.startswith('network.'):
param_dict_new[key[8:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
args.logger.info('load model {} success'.format(args.pretrained))
train_net = BuildTrainNetworkV2(network, criterion0, criterion1, criterion2, args)
# optimizer
opt = Momentum(params=train_net.trainable_params(), learning_rate=Tensor(lr), momentum=momentum,
weight_decay=weight_decay)
# package training process
train_net = TrainOneStepWithLossScaleCell(train_net, opt)
train_net.set_broadcast_flag()
# checkpoint
ckpt_max_num = args.max_epoch * args.steps_per_epoch // args.ckpt_interval
train_config = CheckpointConfig(save_checkpoint_steps=args.ckpt_interval, keep_checkpoint_max=ckpt_max_num)
ckpt_cb = ModelCheckpoint(config=train_config, directory=args.outputs_dir, prefix='{}'.format(args.local_rank))
cb_params = _InternalCallbackParam()
cb_params.train_network = train_net
cb_params.epoch_num = ckpt_max_num
cb_params.cur_epoch_num = 1
run_context = RunContext(cb_params)
ckpt_cb.begin(run_context)
train_net.set_train()
t_end = time.time()
t_epoch = time.time()
old_progress = -1
i = 0
scale_manager = DynamicLossScaleManager(init_loss_scale=2 ** 10, scale_factor=2, scale_window=2000)
for data in ds.create_tuple_iterator(output_numpy=True):
batch_images = data[0]
batch_labels = data[1]
coord_mask_0 = data[2]
conf_pos_mask_0 = data[3]
conf_neg_mask_0 = data[4]
cls_mask_0 = data[5]
t_coord_0 = data[6]
t_conf_0 = data[7]
t_cls_0 = data[8]
gt_list_0 = data[9]
coord_mask_1 = data[10]
conf_pos_mask_1 = data[11]
conf_neg_mask_1 = data[12]
cls_mask_1 = data[13]
t_coord_1 = data[14]
t_conf_1 = data[15]
t_cls_1 = data[16]
gt_list_1 = data[17]
coord_mask_2 = data[18]
conf_pos_mask_2 = data[19]
conf_neg_mask_2 = data[20]
cls_mask_2 = data[21]
t_coord_2 = data[22]
t_conf_2 = data[23]
t_cls_2 = data[24]
gt_list_2 = data[25]
img_tensor = Tensor(batch_images, mstype.float32)
coord_mask_tensor_0 = Tensor(coord_mask_0.astype(np.float32))
conf_pos_mask_tensor_0 = Tensor(conf_pos_mask_0.astype(np.float32))
conf_neg_mask_tensor_0 = Tensor(conf_neg_mask_0.astype(np.float32))
cls_mask_tensor_0 = Tensor(cls_mask_0.astype(np.float32))
t_coord_tensor_0 = Tensor(t_coord_0.astype(np.float32))
t_conf_tensor_0 = Tensor(t_conf_0.astype(np.float32))
t_cls_tensor_0 = Tensor(t_cls_0.astype(np.float32))
gt_list_tensor_0 = Tensor(gt_list_0.astype(np.float32))
coord_mask_tensor_1 = Tensor(coord_mask_1.astype(np.float32))
conf_pos_mask_tensor_1 = Tensor(conf_pos_mask_1.astype(np.float32))
conf_neg_mask_tensor_1 = Tensor(conf_neg_mask_1.astype(np.float32))
cls_mask_tensor_1 = Tensor(cls_mask_1.astype(np.float32))
t_coord_tensor_1 = Tensor(t_coord_1.astype(np.float32))
t_conf_tensor_1 = Tensor(t_conf_1.astype(np.float32))
t_cls_tensor_1 = Tensor(t_cls_1.astype(np.float32))
gt_list_tensor_1 = Tensor(gt_list_1.astype(np.float32))
coord_mask_tensor_2 = Tensor(coord_mask_2.astype(np.float32))
conf_pos_mask_tensor_2 = Tensor(conf_pos_mask_2.astype(np.float32))
conf_neg_mask_tensor_2 = Tensor(conf_neg_mask_2.astype(np.float32))
cls_mask_tensor_2 = Tensor(cls_mask_2.astype(np.float32))
t_coord_tensor_2 = Tensor(t_coord_2.astype(np.float32))
t_conf_tensor_2 = Tensor(t_conf_2.astype(np.float32))
t_cls_tensor_2 = Tensor(t_cls_2.astype(np.float32))
gt_list_tensor_2 = Tensor(gt_list_2.astype(np.float32))
scaling_sens = Tensor(scale_manager.get_loss_scale(), dtype=mstype.float32)
loss0, overflow, _ = train_net(img_tensor, coord_mask_tensor_0, conf_pos_mask_tensor_0,
conf_neg_mask_tensor_0, cls_mask_tensor_0, t_coord_tensor_0,
t_conf_tensor_0, t_cls_tensor_0, gt_list_tensor_0,
coord_mask_tensor_1, conf_pos_mask_tensor_1, conf_neg_mask_tensor_1,
cls_mask_tensor_1, t_coord_tensor_1, t_conf_tensor_1,
t_cls_tensor_1, gt_list_tensor_1, coord_mask_tensor_2,
conf_pos_mask_tensor_2, conf_neg_mask_tensor_2,
cls_mask_tensor_2, t_coord_tensor_2, t_conf_tensor_2,
t_cls_tensor_2, gt_list_tensor_2, scaling_sens)
overflow = np.all(overflow.asnumpy())
if overflow:
scale_manager.update_loss_scale(overflow)
else:
scale_manager.update_loss_scale(False)
args.logger.info('rank[{}], iter[{}], loss[{}], overflow:{}, loss_scale:{}, lr:{}, batch_images:{}, '
'batch_labels:{}'.format(args.local_rank, i, loss0, overflow, scaling_sens, lr[i],
batch_images.shape, batch_labels.shape))
# save ckpt
cb_params.cur_step_num = i + 1 # current step number
cb_params.batch_num = i + 2
if args.local_rank == 0:
ckpt_cb.step_end(run_context)
# save Log
if i == 0:
time_for_graph_compile = time.time() - create_network_start
args.logger.important_info('Yolov3, graph compile time={:.2f}s'.format(time_for_graph_compile))
if i % args.steps_per_epoch == 0:
cb_params.cur_epoch_num += 1
if i % args.log_interval == 0 and args.local_rank == 0:
time_used = time.time() - t_end
epoch = int(i / args.steps_per_epoch)
fps = args.batch_size * (i - old_progress) * args.world_size / time_used
args.logger.info('epoch[{}], iter[{}], loss:[{}], {:.2f} imgs/sec'.format(epoch, i, loss0, fps))
t_end = time.time()
old_progress = i
if i % args.steps_per_epoch == 0 and args.local_rank == 0:
epoch_time_used = time.time() - t_epoch
epoch = int(i / args.steps_per_epoch)
fps = args.batch_size * args.world_size * args.steps_per_epoch / epoch_time_used
args.logger.info('=================================================')
args.logger.info('epoch time: epoch[{}], iter[{}], {:.2f} imgs/sec'.format(epoch, i, fps))
args.logger.info('=================================================')
t_epoch = time.time()
i = i + 1
args.logger.info('=============yolov3 training finished==================')
ds = de.MindDataset(
mindrecord_path + "0",
columns_list=["image", "annotation", "image_name", "image_size"])
single_scale_trans = SingleScaleTrans(resize=args.input_shape)
ds = ds.batch(
args.batch_size,
per_batch_map=single_scale_trans,
input_columns=["image", "annotation", "image_name", "image_size"],
num_parallel_workers=8)
args.steps_per_epoch = ds.get_dataset_size()
# backbone
network = backbone_HwYolov3(num_classes, num_anchors_list, args)
# load pretrain model
if os.path.isfile(args.pretrained):
param_dict = load_checkpoint(args.pretrained)
param_dict_new = {}
for key, values in param_dict.items():
if key.startswith('moments.'):
continue
elif key.startswith('network.'):
param_dict_new[key[8:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
print('load model {} success'.format(args.pretrained))
else: