choices=["AIR", "ONNX", "MINDIR"],
default="AIR",
help="file format")
parser.add_argument("--device_target",
type=str,
choices=["Ascend", "GPU", "CPU"],
default="Ascend",
help="device target")
args = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
if args.device_target == "Ascend":
context.set_context(device_id=args.device_id)
if __name__ == "__main__":
network = DenseNet121(config.num_classes)
param_dict = load_checkpoint(args.ckpt_file)
param_dict_new = {}
for key, value in param_dict.items():
if key.startswith("moments."):
continue
elif key.startswith("network."):
param_dict_new[key[8:]] = value
else:
param_dict_new[key] = value
load_param_into_net(network, param_dict_new)
network.add_flags_recursive(fp16=True)
def test(cloud_args=None):
"""
network eval function. Get top1 and top5 ACC from classification.
The result will be save at [./outputs] by default.
"""
args = parse_args(cloud_args)
context.set_context(mode=context.GRAPH_MODE,
device_target=args.device_target,
save_graphs=True)
if args.device_target == 'Ascend':
devid = int(os.getenv('DEVICE_ID'))
context.set_context(device_id=devid)
# init distributed
if args.is_distributed:
init()
args.rank = get_rank()
args.group_size = get_group_size()
args.outputs_dir = os.path.join(
args.log_path,
datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
args.logger = get_logger(args.outputs_dir, args.rank)
args.logger.save_args(args)
# network
args.logger.important_info('start create network')
if os.path.isdir(args.pretrained):
models = list(glob.glob(os.path.join(args.pretrained, '*.ckpt')))
f = lambda x: -1 * int(
os.path.splitext(os.path.split(x)[-1])[0].split('-')[-1].split('_')
[0])
args.models = sorted(models, key=f)
else:
args.models = [
args.pretrained,
]
for model in args.models:
de_dataset = classification_dataset(args.data_dir,
image_size=args.image_size,
per_batch_size=args.per_batch_size,
max_epoch=1,
rank=args.rank,
group_size=args.group_size,
mode='eval')
eval_dataloader = de_dataset.create_tuple_iterator()
network = DenseNet121(args.num_classes)
param_dict = load_checkpoint(model)
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(model))
if args.device_target == 'Ascend':
network.add_flags_recursive(fp16=True)
img_tot = 0
top1_correct = 0
top5_correct = 0
network.set_train(False)
for data, gt_classes in eval_dataloader:
output = network(Tensor(data, mstype.float32))
output = output.asnumpy()
gt_classes = gt_classes.asnumpy()
top1_output = np.argmax(output, (-1))
top5_output = np.argsort(output)[:, -5:]
t1_correct = np.equal(top1_output, gt_classes).sum()
top1_correct += t1_correct
top5_correct += get_top5_acc(top5_output, gt_classes)
img_tot += args.per_batch_size
results = [[top1_correct], [top5_correct], [img_tot]]
args.logger.info('before results={}'.format(results))
if args.is_distributed:
results = generate_results(model, args.rank, args.group_size,
top1_correct, top5_correct, img_tot)
results = np.array(results)
else:
results = np.array(results)
args.logger.info('after results={}'.format(results))
top1_correct = results[0, 0]
top5_correct = results[1, 0]
img_tot = results[2, 0]
acc1 = 100.0 * top1_correct / img_tot
acc5 = 100.0 * top5_correct / img_tot
args.logger.info(
'after allreduce eval: top1_correct={}, tot={}, acc={:.2f}%'.
format(top1_correct, img_tot, acc1))
args.logger.info(
'after allreduce eval: top5_correct={}, tot={}, acc={:.2f}%'.
format(top5_correct, img_tot, acc5))
if args.is_distributed:
release()
def test(cloud_args=None):
"""
network eval function. Get top1 and top5 ACC from classification.
The result will be save at [./outputs] by default.
"""
args = parse_args(cloud_args)
# init distributed
if args.is_distributed:
init()
args.rank = get_rank()
args.group_size = get_group_size()
args.outputs_dir = os.path.join(
args.log_path,
datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
args.logger = get_logger(args.outputs_dir, args.rank)
args.logger.save_args(args)
# network
args.logger.important_info('start create network')
if os.path.isdir(args.pretrained):
models = list(glob.glob(os.path.join(args.pretrained, '*.ckpt')))
f = lambda x: -1 * int(
os.path.splitext(os.path.split(x)[-1])[0].split('-')[-1].split('_')
[0])
args.models = sorted(models, key=f)
else:
args.models = [
args.pretrained,
]
for model in args.models:
de_dataset = classification_dataset(args.data_dir,
image_size=args.image_size,
per_batch_size=args.per_batch_size,
max_epoch=1,
rank=args.rank,
group_size=args.group_size,
mode='eval')
eval_dataloader = de_dataset.create_tuple_iterator()
network = DenseNet121(args.num_classes)
param_dict = load_checkpoint(model)
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(model))
network.add_flags_recursive(fp16=True)
img_tot = 0
top1_correct = 0
top5_correct = 0
network.set_train(False)
for data, gt_classes in eval_dataloader:
output = network(Tensor(data, mstype.float32))
output = output.asnumpy()
gt_classes = gt_classes.asnumpy()
top1_output = np.argmax(output, (-1))
top5_output = np.argsort(output)[:, -5:]
t1_correct = np.equal(top1_output, gt_classes).sum()
top1_correct += t1_correct
top5_correct += get_top5_acc(top5_output, gt_classes)
img_tot += args.per_batch_size
results = [[top1_correct], [top5_correct], [img_tot]]
args.logger.info('before results={}'.format(results))
if args.is_distributed:
model_md5 = model.replace('/', '')
tmp_dir = '../cache'
if not os.path.exists(tmp_dir):
os.mkdir(tmp_dir)
top1_correct_npy = '{}/top1_rank_{}_{}.npy'.format(
tmp_dir, args.rank, model_md5)
top5_correct_npy = '{}/top5_rank_{}_{}.npy'.format(
tmp_dir, args.rank, model_md5)
img_tot_npy = '{}/img_tot_rank_{}_{}.npy'.format(
tmp_dir, args.rank, model_md5)
np.save(top1_correct_npy, top1_correct)
np.save(top5_correct_npy, top5_correct)
np.save(img_tot_npy, img_tot)
while True:
rank_ok = True
for other_rank in range(args.group_size):
top1_correct_npy = '{}/top1_rank_{}_{}.npy'.format(
tmp_dir, other_rank, model_md5)
top5_correct_npy = '{}/top5_rank_{}_{}.npy'.format(
tmp_dir, other_rank, model_md5)
img_tot_npy = '{}/img_tot_rank_{}_{}.npy'.format(
tmp_dir, other_rank, model_md5)
if not os.path.exists(top1_correct_npy) or not os.path.exists(top5_correct_npy) \
or not os.path.exists(img_tot_npy):
rank_ok = False
if rank_ok:
break
top1_correct_all = 0
top5_correct_all = 0
img_tot_all = 0
for other_rank in range(args.group_size):
top1_correct_npy = '{}/top1_rank_{}_{}.npy'.format(
tmp_dir, other_rank, model_md5)
top5_correct_npy = '{}/top5_rank_{}_{}.npy'.format(
tmp_dir, other_rank, model_md5)
img_tot_npy = '{}/img_tot_rank_{}_{}.npy'.format(
tmp_dir, other_rank, model_md5)
top1_correct_all += np.load(top1_correct_npy)
top5_correct_all += np.load(top5_correct_npy)
img_tot_all += np.load(img_tot_npy)
results = [[top1_correct_all], [top5_correct_all], [img_tot_all]]
results = np.array(results)
else:
results = np.array(results)
args.logger.info('after results={}'.format(results))
top1_correct = results[0, 0]
top5_correct = results[1, 0]
img_tot = results[2, 0]
acc1 = 100.0 * top1_correct / img_tot
acc5 = 100.0 * top5_correct / img_tot
args.logger.info(
'after allreduce eval: top1_correct={}, tot={}, acc={:.2f}%'.
format(top1_correct, img_tot, acc1))
args.logger.info(
'after allreduce eval: top5_correct={}, tot={}, acc={:.2f}%'.
format(top5_correct, img_tot, acc5))
if args.is_distributed:
release()
def create_network(name, *args, **kwargs):
if name == 'densenet121':
return DenseNet121(*args, **kwargs)
raise NotImplementedError(f"{name} is not implemented in the repo")
def train(cloud_args=None):
"""training process"""
args = parse_args(cloud_args)
context.set_context(mode=context.GRAPH_MODE,
enable_auto_mixed_precision=True,
device_target=args.device_target,
save_graphs=False)
if args.device_target == 'Ascend':
devid = int(os.getenv('DEVICE_ID'))
context.set_context(device_id=devid)
# init distributed
if args.is_distributed:
init()
args.rank = get_rank()
args.group_size = get_group_size()
if args.is_dynamic_loss_scale == 1:
args.loss_scale = 1 # for dynamic loss scale can not set loss scale in momentum opt
# select for master rank save ckpt or all rank save, compatible for model parallel
args.rank_save_ckpt_flag = 0
if args.is_save_on_master:
if args.rank == 0:
args.rank_save_ckpt_flag = 1
else:
args.rank_save_ckpt_flag = 1
# logger
args.outputs_dir = os.path.join(
args.ckpt_path,
datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
args.logger = get_logger(args.outputs_dir, args.rank)
# dataloader
de_dataset = classification_dataset(args.data_dir, args.image_size,
args.per_batch_size, args.max_epoch,
args.rank, args.group_size)
de_dataset.map_model = 4
args.steps_per_epoch = de_dataset.get_dataset_size()
args.logger.save_args(args)
# network
args.logger.important_info('start create network')
# get network and init
network = DenseNet121(args.num_classes)
# loss
if not args.label_smooth:
args.label_smooth_factor = 0.0
criterion = CrossEntropy(smooth_factor=args.label_smooth_factor,
num_classes=args.num_classes)
# 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))
# lr scheduler
if args.lr_scheduler == 'exponential':
lr_scheduler = MultiStepLR(args.lr,
args.lr_epochs,
args.lr_gamma,
args.steps_per_epoch,
args.max_epoch,
warmup_epochs=args.warmup_epochs)
elif args.lr_scheduler == 'cosine_annealing':
lr_scheduler = CosineAnnealingLR(args.lr,
args.T_max,
args.steps_per_epoch,
args.max_epoch,
warmup_epochs=args.warmup_epochs,
eta_min=args.eta_min)
else:
raise NotImplementedError(args.lr_scheduler)
lr_schedule = lr_scheduler.get_lr()
# optimizer
opt = Momentum(params=get_param_groups(network),
learning_rate=Tensor(lr_schedule),
momentum=args.momentum,
weight_decay=args.weight_decay,
loss_scale=args.loss_scale)
# mixed precision training
criterion.add_flags_recursive(fp32=True)
# package training process, adjust lr + forward + backward + optimizer
train_net = BuildTrainNetwork(network, criterion)
if args.is_distributed:
parallel_mode = ParallelMode.DATA_PARALLEL
else:
parallel_mode = ParallelMode.STAND_ALONE
if args.is_dynamic_loss_scale == 1:
loss_scale_manager = DynamicLossScaleManager(init_loss_scale=65536,
scale_factor=2,
scale_window=2000)
else:
loss_scale_manager = FixedLossScaleManager(args.loss_scale,
drop_overflow_update=False)
context.set_auto_parallel_context(parallel_mode=parallel_mode,
device_num=args.group_size,
gradients_mean=True)
if args.device_target == 'Ascend':
model = Model(train_net,
optimizer=opt,
metrics=None,
loss_scale_manager=loss_scale_manager,
amp_level="O3")
elif args.device_target == 'GPU':
model = Model(train_net,
optimizer=opt,
metrics=None,
loss_scale_manager=loss_scale_manager,
amp_level="O0")
else:
raise ValueError("Unsupported device target.")
# checkpoint save
progress_cb = ProgressMonitor(args)
callbacks = [
progress_cb,
]
if args.rank_save_ckpt_flag:
ckpt_max_num = args.max_epoch * args.steps_per_epoch // args.ckpt_interval
ckpt_config = CheckpointConfig(
save_checkpoint_steps=args.ckpt_interval,
keep_checkpoint_max=ckpt_max_num)
ckpt_cb = ModelCheckpoint(config=ckpt_config,
directory=args.outputs_dir,
prefix='{}'.format(args.rank))
callbacks.append(ckpt_cb)
model.train(args.max_epoch, de_dataset, callbacks=callbacks)