context.reset_auto_parallel_context()
parallel_mode = ParallelMode.STAND_ALONE
degree = 1
if args.is_distributed:
parallel_mode = ParallelMode.DATA_PARALLEL
degree = get_group_size()
context.set_auto_parallel_context(parallel_mode=parallel_mode,
gradients_mean=True,
device_num=degree)
network = YOLOV4CspDarkNet53(is_training=True)
# default is kaiming-normal
config = ConfigYOLOV4CspDarkNet53()
args.checkpoint_filter_list = config.checkpoint_filter_list
default_recurisive_init(network)
load_yolov4_params(args, network)
network = YoloWithLossCell(network)
args.logger.info('finish get network')
config.label_smooth = args.label_smooth
config.label_smooth_factor = args.label_smooth_factor
if args.training_shape:
config.multi_scale = [convert_training_shape(args.training_shape)]
if args.resize_rate:
config.resize_rate = args.resize_rate
ds, data_size = create_yolo_dataset(image_dir=args.data_root,
anno_path=args.annFile,
is_training=True,
def train():
"""Train function."""
args = parse_args()
devid = int(os.getenv('DEVICE_ID', '0'))
context.set_context(mode=context.GRAPH_MODE,
enable_auto_mixed_precision=True,
device_target=args.device_target,
save_graphs=False,
device_id=devid)
loss_meter = AverageMeter('loss')
network = YOLOV4CspDarkNet53(is_training=True)
# default is kaiming-normal
default_recursive_init(network)
if args.pretrained_backbone:
pretrained_backbone_slice = args.pretrained_backbone.split('/')
backbone_ckpt_file = pretrained_backbone_slice[
len(pretrained_backbone_slice) - 1]
local_backbone_ckpt_path = '/cache/' + backbone_ckpt_file
# download backbone checkpoint
mox.file.copy_parallel(src_url=args.pretrained_backbone,
dst_url=local_backbone_ckpt_path)
args.pretrained_backbone = local_backbone_ckpt_path
load_yolov4_params(args, network)
network = YoloWithLossCell(network)
args.logger.info('finish get network')
config = ConfigYOLOV4CspDarkNet53()
config.label_smooth = args.label_smooth
config.label_smooth_factor = args.label_smooth_factor
if args.training_shape:
config.multi_scale = [convert_training_shape(args)]
if args.resize_rate:
config.resize_rate = args.resize_rate
# data download
local_data_path = '/cache/data'
local_ckpt_path = '/cache/ckpt_file'
print('Download data.')
mox.file.copy_parallel(src_url=args.data_url, dst_url=local_data_path)
ds, data_size = create_yolo_dataset(
image_dir=os.path.join(local_data_path, 'images'),
anno_path=os.path.join(local_data_path, 'annotation.json'),
is_training=True,
batch_size=args.per_batch_size,
max_epoch=args.max_epoch,
device_num=args.group_size,
rank=args.rank,
config=config)
args.logger.info('Finish loading dataset')
args.steps_per_epoch = int(data_size / args.per_batch_size /
args.group_size)
if not args.ckpt_interval:
args.ckpt_interval = args.steps_per_epoch * 10
lr = get_lr(args)
opt = Momentum(params=get_param_groups(network),
learning_rate=Tensor(lr),
momentum=args.momentum,
weight_decay=args.weight_decay,
loss_scale=args.loss_scale)
is_gpu = context.get_context("device_target") == "GPU"
if is_gpu:
loss_scale_value = 1.0
loss_scale = FixedLossScaleManager(loss_scale_value,
drop_overflow_update=False)
network = amp.build_train_network(network,
optimizer=opt,
loss_scale_manager=loss_scale,
level="O2",
keep_batchnorm_fp32=False)
keep_loss_fp32(network)
else:
network = TrainingWrapper(network, opt)
network.set_train()
# checkpoint save
ckpt_max_num = 10
ckpt_config = CheckpointConfig(save_checkpoint_steps=args.ckpt_interval,
keep_checkpoint_max=ckpt_max_num)
ckpt_cb = ModelCheckpoint(config=ckpt_config,
directory=local_ckpt_path,
prefix='yolov4')
cb_params = _InternalCallbackParam()
cb_params.train_network = network
cb_params.epoch_num = ckpt_max_num
cb_params.cur_epoch_num = 1
run_context = RunContext(cb_params)
ckpt_cb.begin(run_context)
old_progress = -1
t_end = time.time()
data_loader = ds.create_dict_iterator(output_numpy=True, num_epochs=1)
for i, data in enumerate(data_loader):
images = data["image"]
input_shape = images.shape[2:4]
images = Tensor.from_numpy(images)
batch_y_true_0 = Tensor.from_numpy(data['bbox1'])
batch_y_true_1 = Tensor.from_numpy(data['bbox2'])
batch_y_true_2 = Tensor.from_numpy(data['bbox3'])
batch_gt_box0 = Tensor.from_numpy(data['gt_box1'])
batch_gt_box1 = Tensor.from_numpy(data['gt_box2'])
batch_gt_box2 = Tensor.from_numpy(data['gt_box3'])
input_shape = Tensor(tuple(input_shape[::-1]), ms.float32)
loss = network(images, batch_y_true_0, batch_y_true_1, batch_y_true_2,
batch_gt_box0, batch_gt_box1, batch_gt_box2,
input_shape)
loss_meter.update(loss.asnumpy())
# ckpt progress
cb_params.cur_step_num = i + 1 # current step number
cb_params.batch_num = i + 2
ckpt_cb.step_end(run_context)
if i % args.log_interval == 0:
time_used = time.time() - t_end
epoch = int(i / args.steps_per_epoch)
fps = args.per_batch_size * (
i - old_progress) * args.group_size / time_used
if args.rank == 0:
args.logger.info(
'epoch[{}], iter[{}], {}, {:.2f} imgs/sec, lr:{}'.format(
epoch, i, loss_meter, fps, lr[i]))
t_end = time.time()
loss_meter.reset()
old_progress = i
if (i + 1) % args.steps_per_epoch == 0:
cb_params.cur_epoch_num += 1
args.logger.info('==========end training===============')
# upload checkpoint files
print('Upload checkpoint.')
mox.file.copy_parallel(src_url=local_ckpt_path, dst_url=args.train_url)