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,
mirror_mean=True,
parameter_broadcast=True)
auto_parallel_context().set_all_reduce_fusion_split_indices([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 = nn.SoftmaxCrossEntropyWithLogits(
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=epoch_size,
batch_size=config.batch_size)
step_size = dataset.get_dataset_size()
eval_interval = config.eval_interval
dataset.__loop_size__ = step_size * eval_interval
# evalutation dataset
eval_dataset = create_dataset(dataset_path=eval_path,
do_train=False,
repeat_num=epoch_size,
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 = list(
filter(
lambda x: 'beta' not in x.name and 'gamma' not in x.name and 'bias'
not in x.name, net.trainable_params()))
no_decayed_params = [
param for param in net.trainable_params()
if param not in decayed_params
]
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)
# model init
print("init_start", device_id)
model.init(dataset, eval_dataset)
print("init_stop", device_id)
# 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})
value=args.batch_size * device_num)
mllogger.event(key="opt_name", value="lars")
mllogger.event(key="lars_opt_base_learning_rate", value=args.lr_max)
mllogger.event(key="lars_opt_end_learning_rate", value=args.lr_min)
mllogger.event(key="lars_opt_learning_rate_decay_poly_power",
value=args.poly_power)
mllogger.event(key="lars_opt_learning_rate_decay_steps",
value=step_size * (epoch_size - args.warmup_epochs))
mllogger.event(key="lars_epsilon", value=args.lars_epsilon)
mllogger.event(key="lars_opt_learning_rate_warmup_epochs",
value=args.warmup_epochs)
mllogger.event(key="lars_opt_momentum", value=args.momentum)
mllogger.event(key="lars_opt_weight_decay", value=args.weight_decay)
mllogger.start(key=mllog.constants.INIT_START)
model.init(dataset, eval_dataset)
sync_all_rank(device_num=device_num)
mllogger.end(key=mllog.constants.INIT_STOP)
# callbacks
state_cb = StateMonitor(data_size=step_size * eval_interval,
mllogger=mllogger,
tot_batch_size=args.batch_size * device_num,
lrs=lr_array,
device_id=device_id,
model=model,
eval_dataset=eval_dataset,
eval_interval=eval_interval)
cb = [
state_cb,
]
