def do_train(dataset=None,
network=None,
load_checkpoint_path="",
save_checkpoint_path="",
epoch_num=1):
""" do train """
if load_checkpoint_path == "":
raise ValueError(
"Pretrain model missed, finetune task must load pretrain model!")
steps_per_epoch = dataset.get_dataset_size()
# optimizer
if optimizer_cfg.optimizer == 'AdamWeightDecay':
lr_schedule = BertLearningRate(
learning_rate=optimizer_cfg.AdamWeightDecay.learning_rate,
end_learning_rate=optimizer_cfg.AdamWeightDecay.end_learning_rate,
warmup_steps=int(steps_per_epoch * epoch_num * 0.1),
decay_steps=steps_per_epoch * epoch_num,
power=optimizer_cfg.AdamWeightDecay.power)
params = net_with_loss.trainable_params()
decay_params = list(
filter(optimizer_cfg.AdamWeightDecay.decay_filter, params))
other_params = list(filter(lambda x: x not in decay_params, params))
group_params = [{
'params':
decay_params,
'weight_decay':
optimizer_cfg.AdamWeightDecay.weight_decay
}, {
'params': other_params,
'weight_decay': 0.0
}]
optimizer = AdamWeightDecay(group_params,
lr_schedule,
eps=optimizer_cfg.AdamWeightDecay.eps)
elif optimizer_cfg.optimizer == 'Lamb':
lr_schedule = BertLearningRate(
learning_rate=optimizer_cfg.Lamb.learning_rate,
end_learning_rate=optimizer_cfg.Lamb.end_learning_rate,
warmup_steps=int(steps_per_epoch * epoch_num * 0.1),
decay_steps=steps_per_epoch * epoch_num,
power=optimizer_cfg.Lamb.power)
optimizer = Lamb(network.trainable_params(), learning_rate=lr_schedule)
elif optimizer_cfg.optimizer == 'Momentum':
optimizer = Momentum(
network.trainable_params(),
learning_rate=optimizer_cfg.Momentum.learning_rate,
momentum=optimizer_cfg.Momentum.momentum)
else:
raise Exception(
"Optimizer not supported. support: [AdamWeightDecay, Lamb, Momentum]"
)
# load checkpoint into network
ckpt_config = CheckpointConfig(save_checkpoint_steps=steps_per_epoch,
keep_checkpoint_max=1)
ckpoint_cb = ModelCheckpoint(prefix="classifier",
directory=save_checkpoint_path,
config=ckpt_config)
param_dict = load_checkpoint(load_checkpoint_path)
load_param_into_net(network, param_dict)
update_cell = DynamicLossScaleUpdateCell(loss_scale_value=2**32,
scale_factor=2,
scale_window=1000)
netwithgrads = BertFinetuneCell(network,
optimizer=optimizer,
scale_update_cell=update_cell)
model = Model(netwithgrads)
callbacks = [
TimeMonitor(dataset.get_dataset_size()),
LossCallBack(), ckpoint_cb
]
model.train(epoch_num, dataset, callbacks=callbacks)
else:
loss = SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction='mean')
print("train args: ", args_opt, "\ncfg: ", config,
"\nparallel args: rank_id {}, device_id {}, rank_size {}".format(rank_id, device_id, rank_size))
dataset = create_dataset(dataset_path=args_opt.dataset_path, do_train=True,
repeat_num=epoch_size, batch_size=config.batch_size)
step_size = dataset.get_dataset_size()
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
load_param_into_net(net, param_dict)
loss_scale = FixedLossScaleManager(config.loss_scale, drop_overflow_update=False)
lr = Tensor(get_lr(global_step=0, lr_init=0, lr_end=0, lr_max=config.lr,
warmup_epochs=config.warmup_epochs, total_epochs=epoch_size, steps_per_epoch=step_size))
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr, config.momentum,
config.weight_decay, config.loss_scale)
model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale)
cb = None
if rank_id == 0:
cb = [Monitor(lr_init=lr.asnumpy())]
if config.save_checkpoint:
config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_epochs * step_size,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="mobilenet", directory=config.save_checkpoint_path, config=config_ck)
cb += [ckpt_cb]
model.train(epoch_size, dataset, callbacks=cb)
def train_process(q, device_id, epoch_size, device_num, enable_hccl):
os.system("mkdir " + str(device_id))
os.chdir(str(device_id))
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
save_graphs=False)
context.set_context(device_id=device_id)
os.environ['MINDSPORE_HCCL_CONFIG_PATH'] = MINDSPORE_HCCL_CONFIG_PATH
os.environ['RANK_ID'] = str(device_id)
os.environ['RANK_SIZE'] = str(device_num)
if enable_hccl:
context.set_auto_parallel_context(
device_num=device_num,
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
all_reduce_fusion_config=[107, 160])
init()
# network
net = resnet50(class_num=config.class_num)
# evaluation network
dist_eval_network = ClassifyCorrectCell(net)
if not config.use_label_smooth:
config.label_smooth_factor = 0.0
# loss
loss = CrossEntropySmooth(sparse=True,
reduction="mean",
smooth_factor=config.label_smooth_factor,
num_classes=config.class_num)
# train dataset
dataset = create_dataset(dataset_path=dataset_path,
do_train=True,
repeat_num=1,
batch_size=config.batch_size)
step_size = dataset.get_dataset_size()
eval_interval = config.eval_interval
dataset.__loop_size__ = step_size * eval_interval
# evaluation dataset
eval_dataset = create_dataset(dataset_path=eval_path,
do_train=False,
repeat_num=1,
batch_size=config.eval_batch_size)
# loss scale
loss_scale = FixedLossScaleManager(config.loss_scale,
drop_overflow_update=False)
# learning rate
lr = Tensor(
get_learning_rate(lr_init=config.lr_init,
lr_end=0.0,
lr_max=config.lr_max,
warmup_epochs=config.warmup_epochs,
total_epochs=config.epoch_size,
steps_per_epoch=step_size,
lr_decay_mode=config.lr_decay_mode))
# optimizer
decayed_params = []
no_decayed_params = []
for param in net.trainable_params():
if 'beta' not in param.name and 'gamma' not in param.name and 'bias' not in param.name:
decayed_params.append(param)
else:
no_decayed_params.append(param)
group_params = [{
'params': decayed_params,
'weight_decay': config.weight_decay
}, {
'params': no_decayed_params,
'weight_decay': 0.0
}, {
'order_params': net.trainable_params()
}]
if config.use_lars:
momentum = nn.Momentum(group_params,
lr,
config.momentum,
loss_scale=config.loss_scale,
use_nesterov=config.use_nesterov)
opt = nn.LARS(momentum,
epsilon=config.lars_epsilon,
coefficient=config.lars_coefficient,
lars_filter=lambda x: 'beta' not in x.name and 'gamma'
not in x.name and 'bias' not in x.name)
else:
opt = nn.Momentum(group_params,
lr,
config.momentum,
loss_scale=config.loss_scale,
use_nesterov=config.use_nesterov)
# model
model = Model(net,
loss_fn=loss,
optimizer=opt,
loss_scale_manager=loss_scale,
amp_level="O2",
keep_batchnorm_fp32=False,
metrics={
'acc':
DistAccuracy(batch_size=config.eval_batch_size,
device_num=device_num)
},
eval_network=dist_eval_network)
# callbacks
loss_cb = LossGet(1, step_size)
# train and eval
print("run_start", device_id)
acc = 0.0
time_cost = 0.0
for epoch_idx in range(0, int(epoch_size / eval_interval)):
model.train(1, dataset, callbacks=loss_cb)
eval_start = time.time()
output = model.eval(eval_dataset)
eval_cost = (time.time() - eval_start) * 1000
acc = float(output["acc"])
time_cost = loss_cb.get_per_step_time()
loss = loss_cb.get_loss()
print(
"the {} epoch's resnet result:\n "
"device{}, training loss {}, acc {}, "
"training per step cost {:.2f} ms, eval cost {:.2f} ms, total_cost {:.2f} ms"
.format(epoch_idx, device_id, loss, acc, time_cost, eval_cost,
time_cost * step_size + eval_cost))
q.put({'acc': acc, 'cost': time_cost})
def test_bert_performance():
"""test bert performance"""
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
reserve_class_name_in_scope=False)
ds, new_repeat_count = me_de_train_dataset(sink_mode=True)
version = os.getenv('VERSION', 'large')
batch_size = 16
config = get_config(version=version, batch_size=batch_size)
netwithloss = BertNetworkWithLoss(config, True)
optimizer = Lamb(netwithloss.trainable_params(),
decay_steps=ds.get_dataset_size() * new_repeat_count,
start_learning_rate=5e-5,
end_learning_rate=1e-9,
power=10.0,
warmup_steps=0,
weight_decay=0.01)
scale_window = 3
scale_manager = DynamicLossScaleManager(2**16, 2, scale_window)
netwithgrads = BertTrainOneStepWithLossScaleCell(
netwithloss,
optimizer=optimizer,
scale_update_cell=scale_manager.get_update_cell())
netwithgrads.set_train(True)
model = Model(netwithgrads)
callback = ModelCallback()
params = netwithloss.trainable_params()
for param in params:
param.init_data()
value = param.default_input
name = param.name
if isinstance(value, Tensor):
if name.split('.')[-1] in ['weight']:
if name.split('.')[-3] in ['cls2']:
logger.info(
"***************** BERT param name is 1 {}".format(
name))
param.default_input = weight_variable(
value.asnumpy().shape)
else:
logger.info(
"***************** BERT param name is 2 {}".format(
name))
tempshape = value.asnumpy().shape
shape = (tempshape[1], tempshape[0])
weight_value = weight_variable(shape).asnumpy()
param.default_input = Tensor(
np.transpose(weight_value, [1, 0]))
else:
logger.info(
"***************** BERT param name is 3 {}".format(name))
param.default_input = weight_variable(value.asnumpy().shape)
time_monitor_callback = TimeMonitor(ds.get_dataset_size())
model.train(new_repeat_count,
ds,
callbacks=[time_monitor_callback, callback],
dataset_sink_mode=True)
# assertion occurs while the loss value, overflow state or loss_scale value is wrong
loss_value = np.array(callback.loss_list)
expect_loss_value = [10.237753, 10.213153, 10.212972]
print("loss value: {}".format(loss_value))
assert np.allclose(loss_value, expect_loss_value, 0, 0.0005)
overflow = np.array(callback.overflow_list)
expect_overflow = [False, False, False]
print("overflow: {}".format(overflow))
assert (overflow == expect_overflow).all()
loss_scale = np.array(callback.lossscale_list)
expect_loss_scale = [16384.0, 16384.0, 16384.0]
print("loss scale: {}".format(loss_scale))
assert np.allclose(loss_scale, expect_loss_scale, 0, 0)
epoch_mseconds = np.array(time_monitor_callback.epoch_mseconds_list)[2]
expect_epoch_mseconds = 1726
print("epoch mseconds: {}".format(epoch_mseconds))
assert epoch_mseconds <= expect_epoch_mseconds + 5
per_step_mseconds = np.array(
time_monitor_callback.per_step_mseconds_list)[2]
expect_per_step_mseconds = 17
print("per step mseconds: {}".format(per_step_mseconds))
assert per_step_mseconds <= expect_per_step_mseconds + 1
def train_on_ascend():
config = config_ascend_quant
print("training args: {}".format(args_opt))
print("training configure: {}".format(config))
print("parallel args: rank_id {}, device_id {}, rank_size {}".format(
rank_id, device_id, rank_size))
epoch_size = config.epoch_size
# distribute init
if run_distribute:
context.set_auto_parallel_context(
device_num=rank_size,
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True)
init()
# define network
network = mobilenetV2(num_classes=config.num_classes)
# define loss
if config.label_smooth > 0:
loss = CrossEntropyWithLabelSmooth(smooth_factor=config.label_smooth,
num_classes=config.num_classes)
else:
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
# define dataset
dataset = create_dataset(dataset_path=args_opt.dataset_path,
do_train=True,
config=config,
device_target=args_opt.device_target,
repeat_num=1,
batch_size=config.batch_size)
step_size = dataset.get_dataset_size()
# load pre trained ckpt
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
load_nonquant_param_into_quant_net(network, param_dict)
# convert fusion network to quantization aware network
quantizer = QuantizationAwareTraining(bn_fold=True,
per_channel=[True, False],
symmetric=[True, False],
one_conv_fold=True)
network = quantizer.quantize(network)
# get learning rate
lr = Tensor(
get_lr(global_step=config.start_epoch * step_size,
lr_init=0,
lr_end=0,
lr_max=config.lr,
warmup_epochs=config.warmup_epochs,
total_epochs=epoch_size + config.start_epoch,
steps_per_epoch=step_size))
# define optimization
opt = nn.Momentum(
filter(lambda x: x.requires_grad, network.get_parameters()), lr,
config.momentum, config.weight_decay)
# define model
model = Model(network, loss_fn=loss, optimizer=opt)
print("============== Starting Training ==============")
callback = None
if rank_id == 0:
callback = [Monitor(lr_init=lr.asnumpy())]
if config.save_checkpoint:
config_ck = CheckpointConfig(
save_checkpoint_steps=config.save_checkpoint_epochs *
step_size,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="mobilenetV2",
directory=config.save_checkpoint_path,
config=config_ck)
callback += [ckpt_cb]
model.train(epoch_size, dataset, callbacks=callback)
print("============== End Training ==============")
def test_bert_performance():
"""test bert performance"""
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
reserve_class_name_in_scope=False)
ds, new_repeat_count, sink_size = me_de_train_dataset(sink_mode=True)
version = os.getenv('VERSION', 'large')
batch_size = 16
config = get_config(version=version, batch_size=batch_size)
netwithloss = BertNetworkWithLoss(config, True)
lr = BertLearningRate(decay_steps=sink_size * new_repeat_count,
learning_rate=5e-5,
end_learning_rate=1e-9,
power=10.0,
warmup_steps=0)
decay_filter = lambda x: 'layernorm' not in x.name.lower(
) and 'bias' not in x.name.lower()
no_decay_filter = lambda x: 'layernorm' in x.name.lower(
) or 'bias' in x.name.lower()
decay_params = list(filter(decay_filter, netwithloss.trainable_params()))
other_params = list(filter(no_decay_filter,
netwithloss.trainable_params()))
group_params = [{
'params': decay_params,
'weight_decay': 0.01
}, {
'params': other_params
}, {
'order_params': netwithloss.trainable_params()
}]
optimizer = Lamb(group_params, lr)
scale_window = 3
scale_manager = DynamicLossScaleManager(2**16, 2, scale_window)
netwithgrads = BertTrainOneStepWithLossScaleCell(
netwithloss,
optimizer=optimizer,
scale_update_cell=scale_manager.get_update_cell())
netwithgrads.set_train(True)
model = Model(netwithgrads)
callback = ModelCallback()
params = netwithloss.trainable_params()
for param in params:
value = param.default_input
name = param.name
if isinstance(value, Tensor):
if name.split('.')[-1] in ['weight']:
if name.split('.')[-3] in ['cls2']:
logger.info(
"***************** BERT param name is 1 {}".format(
name))
param.default_input = weight_variable(
value.asnumpy().shape)
else:
logger.info(
"***************** BERT param name is 2 {}".format(
name))
tempshape = value.asnumpy().shape
shape = (tempshape[1], tempshape[0])
weight_value = weight_variable(shape).asnumpy()
param.default_input = Tensor(
np.transpose(weight_value, [1, 0]))
else:
logger.info(
"***************** BERT param name is 3 {}".format(name))
param.default_input = weight_variable(value.asnumpy().shape)
time_monitor_callback = TimeMonitor(sink_size)
model.train(new_repeat_count,
ds,
callbacks=[time_monitor_callback, callback],
dataset_sink_mode=True,
sink_size=sink_size)
# assertion occurs while the loss value, overflow state or loss_scale value is wrong
loss_value = np.array(callback.loss_list)
expect_loss_value = [10.235566, 10.207392, 10.206976]
print("loss value: {}".format(loss_value))
assert np.allclose(loss_value, expect_loss_value, 0, 0.0005)
overflow = np.array(callback.overflow_list)
expect_overflow = [True, True, True]
print("overflow: {}".format(overflow))
assert (overflow == expect_overflow).all()
loss_scale = np.array(callback.lossscale_list)
expect_loss_scale = [262144.0, 262144.0, 262144.0]
print("loss scale: {}".format(loss_scale))
assert np.allclose(loss_scale, expect_loss_scale, 0, 0)
epoch_mseconds = np.array(time_monitor_callback.epoch_mseconds_list)[2]
expect_epoch_mseconds = 1400
print("epoch mseconds: {}".format(epoch_mseconds))
assert epoch_mseconds <= expect_epoch_mseconds + 5
per_step_mseconds = np.array(
time_monitor_callback.per_step_mseconds_list)[2]
expect_per_step_mseconds = 14
print("per step mseconds: {}".format(per_step_mseconds))
assert per_step_mseconds <= expect_per_step_mseconds + 1
opt = Momentum(params=get_param_groups(net),
learning_rate=Tensor(lr),
momentum=cfg.momentum,
weight_decay=cfg.weight_decay,
loss_scale=cfg.loss_scale)
if not cfg.use_label_smooth:
cfg.label_smooth_factor = 0.0
loss = CrossEntropySmooth(sparse=True, reduction="mean",
smooth_factor=cfg.label_smooth_factor, num_classes=cfg.num_classes)
if cfg.is_dynamic_loss_scale == 1:
loss_scale_manager = DynamicLossScaleManager(init_loss_scale=65536, scale_factor=2, scale_window=2000)
else:
loss_scale_manager = FixedLossScaleManager(cfg.loss_scale, drop_overflow_update=False)
if device_target == "Ascend":
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'},
amp_level="O2", keep_batchnorm_fp32=False, loss_scale_manager=loss_scale_manager)
else: # GPU
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'},
amp_level="O2", keep_batchnorm_fp32=True, loss_scale_manager=loss_scale_manager)
config_ck = CheckpointConfig(save_checkpoint_steps=batch_num * 5, keep_checkpoint_max=cfg.keep_checkpoint_max)
time_cb = TimeMonitor(data_size=batch_num)
ckpt_save_dir = "./ckpt_" + str(rank) + "/"
ckpoint_cb = ModelCheckpoint(prefix="train_googlenet_" + args_opt.dataset_name, directory=ckpt_save_dir,
config=config_ck)
loss_cb = LossMonitor()
model.train(cfg.epoch_size, dataset, callbacks=[time_cb, ckpoint_cb, loss_cb])
print("train success")
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.platform, save_graphs=False)
if os.getenv('DEVICE_ID', "not_set").isdigit():
context.set_context(device_id=int(os.getenv('DEVICE_ID')))
# init distributed
if args.is_distributed:
init()
args.rank = get_rank()
args.group_size = get_group_size()
parallel_mode = ParallelMode.DATA_PARALLEL
context.set_auto_parallel_context(parallel_mode=parallel_mode, device_num=args.group_size,
parameter_broadcast=True, mirror_mean=True)
else:
args.rank = 0
args.group_size = 1
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, compatiable 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, 1,
args.rank, args.group_size, num_parallel_workers=8)
de_dataset.map_model = 4 # !!!important
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 = get_network(args.backbone, args.num_classes, platform=args.platform)
if network is None:
raise NotImplementedError('not implement {}'.format(args.backbone))
# 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 = warmup_step_lr(args.lr,
args.lr_epochs,
args.steps_per_epoch,
args.warmup_epochs,
args.max_epoch,
gamma=args.lr_gamma,
)
elif args.lr_scheduler == 'cosine_annealing':
lr = warmup_cosine_annealing_lr(args.lr,
args.steps_per_epoch,
args.warmup_epochs,
args.max_epoch,
args.T_max,
args.eta_min)
else:
raise NotImplementedError(args.lr_scheduler)
# optimizer
opt = Momentum(params=get_param_groups(network),
learning_rate=Tensor(lr),
momentum=args.momentum,
weight_decay=args.weight_decay,
loss_scale=args.loss_scale)
# loss
if not args.label_smooth:
args.label_smooth_factor = 0.0
loss = CrossEntropy(smooth_factor=args.label_smooth_factor, num_classes=args.num_classes)
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)
if args.platform == "Ascend":
model = Model(network, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale_manager,
metrics={'acc'}, amp_level="O3")
else:
model = Model(network, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale_manager,
metrics={'acc'}, amp_level="O2")
# checkpoint save
progress_cb = ProgressMonitor(args)
callbacks = [progress_cb,]
if args.rank_save_ckpt_flag:
ckpt_config = CheckpointConfig(save_checkpoint_steps=args.ckpt_interval * args.steps_per_epoch,
keep_checkpoint_max=args.ckpt_save_max)
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, dataset_sink_mode=True)
def do_train(dataset=None,
network=None,
load_checkpoint_path="",
save_checkpoint_path=""):
""" do train """
if load_checkpoint_path == "":
raise ValueError(
"Pretrain model missed, finetune task must load pretrain model!")
steps_per_epoch = dataset.get_dataset_size()
epoch_num = dataset.get_repeat_count()
# optimizer
if optimizer_cfg.optimizer == 'AdamWeightDecayDynamicLR':
optimizer = AdamWeightDecayDynamicLR(
network.trainable_params(),
decay_steps=steps_per_epoch * epoch_num,
learning_rate=optimizer_cfg.AdamWeightDecayDynamicLR.learning_rate,
end_learning_rate=optimizer_cfg.AdamWeightDecayDynamicLR.
end_learning_rate,
power=optimizer_cfg.AdamWeightDecayDynamicLR.power,
warmup_steps=int(steps_per_epoch * epoch_num * 0.1),
weight_decay=optimizer_cfg.AdamWeightDecayDynamicLR.weight_decay,
eps=optimizer_cfg.AdamWeightDecayDynamicLR.eps)
elif optimizer_cfg.optimizer == 'Lamb':
optimizer = Lamb(
network.trainable_params(),
decay_steps=steps_per_epoch * epoch_num,
start_learning_rate=optimizer_cfg.Lamb.start_learning_rate,
end_learning_rate=optimizer_cfg.Lamb.end_learning_rate,
power=optimizer_cfg.Lamb.power,
weight_decay=optimizer_cfg.Lamb.weight_decay,
warmup_steps=int(steps_per_epoch * epoch_num * 0.1),
decay_filter=optimizer_cfg.Lamb.decay_filter)
elif optimizer_cfg.optimizer == 'Momentum':
optimizer = Momentum(
network.trainable_params(),
learning_rate=optimizer_cfg.Momentum.learning_rate,
momentum=optimizer_cfg.Momentum.momentum)
else:
raise Exception(
"Optimizer not supported. support: [AdamWeightDecayDynamicLR, Lamb, Momentum]"
)
# load checkpoint into network
ckpt_config = CheckpointConfig(save_checkpoint_steps=steps_per_epoch,
keep_checkpoint_max=1)
ckpoint_cb = ModelCheckpoint(prefix="squad",
directory=save_checkpoint_path,
config=ckpt_config)
param_dict = load_checkpoint(load_checkpoint_path)
load_param_into_net(network, param_dict)
update_cell = DynamicLossScaleUpdateCell(loss_scale_value=2**32,
scale_factor=2,
scale_window=1000)
netwithgrads = BertSquadCell(network,
optimizer=optimizer,
scale_update_cell=update_cell)
model = Model(netwithgrads)
callbacks = [
TimeMonitor(dataset.get_dataset_size()),
LossCallBack(), ckpoint_cb
]
model.train(epoch_num, dataset, callbacks=callbacks)
def do_train(dataset=None,
network=None,
load_checkpoint_path="",
save_checkpoint_path="",
epoch_num=1):
"""
Do train
Args:
dataset: the train dataset.
network: the network with loss
load_checkpoint_path: the file path which saved pretrain model checkpoint.
save_checkpoint_path: the file path which will save finetune model checkpoint.
epoch_num: the number of epoch
"""
if load_checkpoint_path == "":
raise ValueError(
"Pretrain model missed, finetune task must load pretrain model!")
steps_per_epoch = dataset.get_dataset_size()
# optimizer
if cfg.optimizer == 'AdamWeightDecay':
lr_schedule = GPT2LearningRate(
learning_rate=cfg.AdamWeightDecay.learning_rate,
end_learning_rate=cfg.AdamWeightDecay.end_learning_rate,
warmup_steps=int(steps_per_epoch * epoch_num * 0.1),
decay_steps=steps_per_epoch * epoch_num,
power=cfg.AdamWeightDecay.power)
params = network.trainable_params()
decay_params = list(filter(cfg.AdamWeightDecay.decay_filter, params))
other_params = list(
filter(lambda x: not cfg.AdamWeightDecay.decay_filter(x), params))
group_params = [{
'params': decay_params,
'weight_decay': cfg.AdamWeightDecay.weight_decay
}, {
'params': other_params,
'weight_decay': 0.0
}]
optimizer = AdamWeightDecay(group_params,
lr_schedule,
eps=cfg.AdamWeightDecay.eps)
elif cfg.optimizer == 'Lamb':
lr_schedule = GPT2LearningRate(
learning_rate=cfg.Lamb.learning_rate,
end_learning_rate=cfg.Lamb.end_learning_rate,
warmup_steps=int(steps_per_epoch * epoch_num * 0.1),
decay_steps=steps_per_epoch * epoch_num,
power=cfg.Lamb.power)
optimizer = Lamb(network.trainable_params(), lr_schedule)
elif cfg.optimizer == 'Momentum':
optimizer = Momentum(network.trainable_params(),
cfg.Momentum.learning_rate, cfg.Momentum.momentum)
else:
raise Exception(
"Optimizer not supported. support: [AdamWeightDecay, Lamb, Momentum]"
)
# load checkpoint into network
ckpt_config = CheckpointConfig(save_checkpoint_steps=steps_per_epoch,
keep_checkpoint_max=1)
prefix_name = "gpt2_" + "lambada_" + str(cfg.gpt2_network) + "_" + str(cfg.optimizer) + "_" \
+ str(epoch_num) + "_bs" + str(gpt2_net_cfg.batch_size)
ckpoint_cb = ModelCheckpoint(
prefix=prefix_name,
directory=None if save_checkpoint_path == "" else save_checkpoint_path,
config=ckpt_config)
param_dict = load_checkpoint(load_checkpoint_path)
final_param_dict = {}
for name, _ in param_dict.items():
final_param_dict['gpt2.gpt2.' + name] = param_dict[name]
final_param_dict['gpt2.dense1.weight'] = param_dict[
'gpt2_embedding_lookup.embedding_table']
load_param_into_net(network, final_param_dict)
print("Load pretrained parameter successfully!\n")
update_cell = DynamicLossScaleUpdateCell(loss_scale_value=2**32,
scale_factor=2,
scale_window=1000)
netwithgrads = GPT2FinetuneCell(network,
optimizer=optimizer,
scale_update_cell=update_cell)
netwithgrads.set_train(True)
loss_cb = LossMonitor(per_print_times=1)
model = Model(netwithgrads)
callbacks = [TimeMonitor(dataset.get_dataset_size()), loss_cb, ckpoint_cb]
print("==================== Starting Finetuning ====================")
model.train(epoch_num,
dataset,
callbacks=callbacks,
dataset_sink_mode=False)
print("==================== Finetuning Success ====================")
def train():
args = parse_args()
if args.device_target == "CPU":
context.set_context(mode=context.GRAPH_MODE, save_graphs=False, device_target="CPU")
else:
context.set_context(mode=context.GRAPH_MODE, enable_auto_mixed_precision=True, save_graphs=False,
device_target="Ascend", device_id=int(os.getenv('DEVICE_ID')))
# init multicards training
if args.is_distributed:
init()
args.rank = get_rank()
args.group_size = get_group_size()
parallel_mode = ParallelMode.DATA_PARALLEL
context.set_auto_parallel_context(parallel_mode=parallel_mode, gradients_mean=True, device_num=args.group_size)
# dataset
dataset = data_generator.SegDataset(image_mean=args.image_mean,
image_std=args.image_std,
data_file=args.data_file,
batch_size=args.batch_size,
crop_size=args.crop_size,
max_scale=args.max_scale,
min_scale=args.min_scale,
ignore_label=args.ignore_label,
num_classes=args.num_classes,
num_readers=2,
num_parallel_calls=4,
shard_id=args.rank,
shard_num=args.group_size)
dataset = dataset.get_dataset(repeat=1)
# network
if args.model == 'deeplab_v3_s16':
network = net_factory.nets_map[args.model]('train', args.num_classes, 16, args.freeze_bn)
elif args.model == 'deeplab_v3_s8':
network = net_factory.nets_map[args.model]('train', args.num_classes, 8, args.freeze_bn)
else:
raise NotImplementedError('model [{:s}] not recognized'.format(args.model))
# loss
loss_ = loss.SoftmaxCrossEntropyLoss(args.num_classes, args.ignore_label)
loss_.add_flags_recursive(fp32=True)
train_net = BuildTrainNetwork(network, loss_)
# load pretrained model
if args.ckpt_pre_trained:
param_dict = load_checkpoint(args.ckpt_pre_trained)
load_param_into_net(train_net, param_dict)
# optimizer
iters_per_epoch = dataset.get_dataset_size()
total_train_steps = iters_per_epoch * args.train_epochs
if args.lr_type == 'cos':
lr_iter = learning_rates.cosine_lr(args.base_lr, total_train_steps, total_train_steps)
elif args.lr_type == 'poly':
lr_iter = learning_rates.poly_lr(args.base_lr, total_train_steps, total_train_steps, end_lr=0.0, power=0.9)
elif args.lr_type == 'exp':
lr_iter = learning_rates.exponential_lr(args.base_lr, args.lr_decay_step, args.lr_decay_rate,
total_train_steps, staircase=True)
else:
raise ValueError('unknown learning rate type')
opt = nn.Momentum(params=train_net.trainable_params(), learning_rate=lr_iter, momentum=0.9, weight_decay=0.0001,
loss_scale=args.loss_scale)
# loss scale
manager_loss_scale = FixedLossScaleManager(args.loss_scale, drop_overflow_update=False)
amp_level = "O0" if args.device_target == "CPU" else "O3"
model = Model(train_net, optimizer=opt, amp_level=amp_level, loss_scale_manager=manager_loss_scale)
# callback for saving ckpts
time_cb = TimeMonitor(data_size=iters_per_epoch)
loss_cb = LossMonitor()
cbs = [time_cb, loss_cb]
if args.rank == 0:
config_ck = CheckpointConfig(save_checkpoint_steps=args.save_steps,
keep_checkpoint_max=args.keep_checkpoint_max)
ckpoint_cb = ModelCheckpoint(prefix=args.model, directory=args.train_dir, config=config_ck)
cbs.append(ckpoint_cb)
model.train(args.train_epochs, dataset, callbacks=cbs, dataset_sink_mode=(args.device_target != "CPU"))
def run_general_distill():
"""
run general distill
"""
parser = argparse.ArgumentParser(description='tinybert general distill')
parser.add_argument(
'--device_target',
type=str,
default='Ascend',
choices=['Ascend', 'GPU'],
help='device where the code will be implemented. (Default: Ascend)')
parser.add_argument("--distribute",
type=str,
default="false",
help="Run distribute, default is false.")
parser.add_argument("--epoch_size",
type=int,
default="3",
help="Epoch size, default is 1.")
parser.add_argument("--device_id",
type=int,
default=0,
help="Device id, default is 0.")
parser.add_argument("--device_num",
type=int,
default=1,
help="Use device nums, default is 1.")
parser.add_argument("--save_ckpt_step",
type=int,
default=100,
help="Enable data sink, default is true.")
parser.add_argument("--max_ckpt_num",
type=int,
default=1,
help="Enable data sink, default is true.")
parser.add_argument("--do_shuffle",
type=str,
default="true",
help="Enable shuffle for dataset, default is true.")
parser.add_argument("--enable_data_sink",
type=str,
default="true",
help="Enable data sink, default is true.")
parser.add_argument("--data_sink_steps",
type=int,
default=1,
help="Sink steps for each epoch, default is 1.")
parser.add_argument("--save_ckpt_path",
type=str,
default="",
help="Save checkpoint path")
parser.add_argument("--load_teacher_ckpt_path",
type=str,
default="",
help="Load checkpoint file path")
parser.add_argument("--data_dir",
type=str,
default="",
help="Data path, it is better to use absolute path")
parser.add_argument("--schema_dir",
type=str,
default="",
help="Schema path, it is better to use absolute path")
args_opt = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target=args_opt.device_target,
device_id=args_opt.device_id)
context.set_context(mode=context.GRAPH_MODE,
device_target=args_opt.device_target,
device_id=args_opt.device_id)
context.set_context(reserve_class_name_in_scope=False)
context.set_context(variable_memory_max_size="30GB")
save_ckpt_dir = os.path.join(
args_opt.save_ckpt_path,
datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
if not os.path.exists(save_ckpt_dir):
os.makedirs(save_ckpt_dir)
if args_opt.distribute == "true":
D.init('hccl')
device_num = args_opt.device_num
rank = args_opt.device_id % device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
mirror_mean=True,
device_num=device_num)
else:
rank = 0
device_num = 1
netwithloss = BertNetworkWithLoss_gd(
teacher_config=bert_teacher_net_cfg,
teacher_ckpt=args_opt.load_teacher_ckpt_path,
student_config=bert_student_net_cfg,
is_training=True,
use_one_hot_embeddings=False)
dataset = create_tinybert_dataset('gd', bert_teacher_net_cfg.batch_size,
device_num, rank, args_opt.do_shuffle,
args_opt.data_dir, args_opt.schema_dir)
dataset_size = dataset.get_dataset_size()
print('dataset size: ', dataset_size)
if args_opt.enable_data_sink == "true":
repeat_count = args_opt.epoch_size * dataset.get_dataset_size(
) // args_opt.data_sink_steps
time_monitor_steps = args_opt.data_sink_steps
else:
repeat_count = args_opt.epoch_size
time_monitor_steps = dataset_size
lr_schedule = BertLearningRate(
learning_rate=common_cfg.AdamWeightDecay.learning_rate,
end_learning_rate=common_cfg.AdamWeightDecay.end_learning_rate,
warmup_steps=int(dataset_size * args_opt.epoch_size / 10),
decay_steps=int(dataset_size * args_opt.epoch_size),
power=common_cfg.AdamWeightDecay.power)
params = netwithloss.trainable_params()
decay_params = list(filter(common_cfg.AdamWeightDecay.decay_filter,
params))
other_params = list(filter(lambda x: x not in decay_params, params))
group_params = [{
'params': decay_params,
'weight_decay': common_cfg.AdamWeightDecay.weight_decay
}, {
'params': other_params,
'weight_decay': 0.0
}, {
'order_params': params
}]
optimizer = AdamWeightDecay(group_params,
learning_rate=lr_schedule,
eps=common_cfg.AdamWeightDecay.eps)
callback = [
TimeMonitor(time_monitor_steps),
LossCallBack(),
ModelSaveCkpt(netwithloss.bert, args_opt.save_ckpt_step,
args_opt.max_ckpt_num, save_ckpt_dir)
]
update_cell = DynamicLossScaleUpdateCell(
loss_scale_value=common_cfg.loss_scale_value,
scale_factor=common_cfg.scale_factor,
scale_window=common_cfg.scale_window)
netwithgrads = BertTrainWithLossScaleCell(netwithloss,
optimizer=optimizer,
scale_update_cell=update_cell)
model = Model(netwithgrads)
model.train(repeat_count,
dataset,
callbacks=callback,
dataset_sink_mode=(args_opt.enable_data_sink == "true"),
sink_size=args_opt.data_sink_steps)
loss_scale_manager=loss_scale,
metrics={'acc'},
amp_level="O2",
keep_batchnorm_fp32=False)
else:
# GPU target
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()),
lr,
config.momentum,
config.weight_decay,
use_nesterov=True)
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
# define callbacks
time_cb = TimeMonitor(data_size=step_size)
loss_cb = LossMonitor()
cb = [time_cb, loss_cb]
if config.save_checkpoint:
config_ck = CheckpointConfig(
save_checkpoint_steps=config.save_checkpoint_epochs * step_size,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix=args_opt.net + '_' + args_opt.dataset,
directory=ckpt_save_dir,
config=config_ck)
cb += [ckpt_cb]
# train model
model.train(config.epoch_size - config.pretrain_epoch_size,
dataset,
callbacks=cb)
# init weight
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
load_param_into_net(net, param_dict)
lr = config.lr
lr = Tensor(lr, ms.float32)
# define opt
opt = Adam(params=net.trainable_params(), learning_rate=lr, eps=1e-07)
# define loss, model
loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction="sum")
model = Model(net, loss_fn=loss, optimizer=opt, metrics={"Accuracy": Accuracy()})
# define callbacks
time_cb = TimeMonitor(data_size=step_size)
loss_cb = LossMonitor()
cb = [time_cb, loss_cb]
if config.save_checkpoint:
if args_opt.rank_save_ckpt_flag == 1:
config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_steps,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="cnn_direction_model", directory=ckpt_save_dir, config=config_ck)
cb += [ckpt_cb]
# train model
model.train(config.epoch_size, dataset, callbacks=cb, dataset_sink_mode=False)
# define loss, model
loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr, momentum)
eval_net = nn.WithEvalCell(net, loss, AMP_LEVEL in ["O2", "O3"])
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'},amp_level=AMP_LEVEL, eval_network=eval_net,
eval_indexes=[0, 1, 2], keep_batchnorm_fp32=False)
# define callbacks
time_cb = TimeMonitor(data_size=step_size)
loss_cb = LossMonitor()
cb = [time_cb, loss_cb]
save_checkpoint = 5
if save_checkpoint:
save_checkpoint_epochs = 5
keep_checkpoint_max = 10
config_ck = CheckpointConfig(save_checkpoint_steps=save_checkpoint_epochs * step_size,
keep_checkpoint_max=keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="resnet", directory=ckpt_save_dir, config=config_ck)
cb += [ckpt_cb]
# train model
model.train(epoch_size, dataset, callbacks=cb, dataset_sink_mode=True)
# Eval model
eval_dataset_path = "./datasets/cifar10/test"
eval_data = create_dataset(eval_dataset_path,do_train=False)
acc = model.eval(eval_data,dataset_sink_mode=True)
print("Accuracy:",acc)
loss_scale = FixedLossScaleManager(config.loss_scale,
drop_overflow_update=False)
# Mixed precision
model = Model(net,
loss_fn=loss,
optimizer=opt,
loss_scale_manager=loss_scale,
metrics={'acc'},
amp_level="O2",
keep_batchnorm_fp32=False)
# define callbacks
time_cb = TimeMonitor(data_size=step_size)
loss_cb = LossMonitor()
cb = [time_cb, loss_cb]
if config.save_checkpoint:
config_ck = CheckpointConfig(
save_checkpoint_steps=config.save_checkpoint_epochs * step_size,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="mobilenetv1",
directory=ckpt_save_dir,
config=config_ck)
cb += [ckpt_cb]
# train model
model.train(config.epoch_size - config.pretrain_epoch_size,
dataset,
callbacks=cb,
sink_size=dataset.get_dataset_size(),
dataset_sink_mode=(not args_opt.parameter_server))
def train(cloud_args=None):
"""training process"""
args = parse_args(cloud_args)
# 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, compatiable 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)
model = Model(train_net,
optimizer=opt,
metrics=None,
loss_scale_manager=loss_scale_manager,
amp_level="O3")
# 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)
epoch_size = args_opt.epoch_size
net = resnet50(args_opt.batch_size, args_opt.num_classes)
loss = CrossEntropyLoss()
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()),
0.01, 0.9)
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
if args_opt.do_train:
dataset = create_dataset(epoch_size)
batch_num = dataset.get_dataset_size()
config_ck = CheckpointConfig(save_checkpoint_steps=batch_num * 5,
keep_checkpoint_max=10)
ckpoint_cb = ModelCheckpoint(prefix="train_resnet_cifar10",
directory="./",
config=config_ck)
time_cb = TimeMonitor(data_size=batch_num)
loss_cb = LossMonitor()
model.train(epoch_size,
dataset,
callbacks=[ckpoint_cb, loss_cb, time_cb])
if args_opt.do_eval:
if args_opt.checkpoint_path:
param_dict = load_checkpoint(args_opt.checkpoint_path)
load_param_into_net(net, param_dict)
net.set_train(False)
eval_dataset = create_dataset(1, training=False)
res = model.eval(eval_dataset)
print("result: ", res)
def train_on_gpu():
config = config_gpu_quant
print("training args: {}".format(args_opt))
print("training configure: {}".format(config))
# define network
network = mobilenetV2(num_classes=config.num_classes)
# define loss
if config.label_smooth > 0:
loss = CrossEntropyWithLabelSmooth(smooth_factor=config.label_smooth,
num_classes=config.num_classes)
else:
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
# define dataset
epoch_size = config.epoch_size
dataset = create_dataset(dataset_path=args_opt.dataset_path,
do_train=True,
config=config,
device_target=args_opt.device_target,
repeat_num=1,
batch_size=config.batch_size)
step_size = dataset.get_dataset_size()
# resume
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
load_nonquant_param_into_quant_net(network, param_dict)
# convert fusion network to quantization aware network
quantizer = QuantizationAwareTraining(bn_fold=True,
per_channel=[True, False],
symmetric=[False, False],
freeze_bn=1000000,
quant_delay=step_size * 2)
network = quantizer.quantize(network)
# get learning rate
loss_scale = FixedLossScaleManager(config.loss_scale,
drop_overflow_update=False)
lr = Tensor(
get_lr(global_step=config.start_epoch * step_size,
lr_init=0,
lr_end=0,
lr_max=config.lr,
warmup_epochs=config.warmup_epochs,
total_epochs=epoch_size + config.start_epoch,
steps_per_epoch=step_size))
# define optimization
opt = nn.Momentum(
filter(lambda x: x.requires_grad, network.get_parameters()), lr,
config.momentum, config.weight_decay, config.loss_scale)
# define model
model = Model(network,
loss_fn=loss,
optimizer=opt,
loss_scale_manager=loss_scale)
print("============== Starting Training ==============")
callback = [Monitor(lr_init=lr.asnumpy())]
ckpt_save_dir = config.save_checkpoint_path + "ckpt_" + str(
get_rank()) + "/"
if config.save_checkpoint:
config_ck = CheckpointConfig(
save_checkpoint_steps=config.save_checkpoint_epochs * step_size,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpt_cb = ModelCheckpoint(prefix="mobilenetV2",
directory=ckpt_save_dir,
config=config_ck)
callback += [ckpt_cb]
model.train(epoch_size, dataset, callbacks=callback)
print("============== End Training ==============")
