def test():
"""The function of eval."""
start_time = time.time()
args = parse_args()
devid = int(os.getenv('DEVICE_ID')) if os.getenv('DEVICE_ID') else 0
context.set_context(mode=context.GRAPH_MODE,
device_target=args.device_target,
save_graphs=False,
device_id=devid)
# logger
args.outputs_dir = os.path.join(
args.log_path,
datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
rank_id = int(
os.environ.get('RANK_ID')) if os.environ.get('RANK_ID') else 0
args.logger = get_logger(args.outputs_dir, rank_id)
context.reset_auto_parallel_context()
parallel_mode = ParallelMode.STAND_ALONE
context.set_auto_parallel_context(parallel_mode=parallel_mode,
gradients_mean=True,
device_num=1)
args.logger.info('Creating Network....')
network = YOLOV3DarkNet53(is_training=False)
args.logger.info(args.pretrained)
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('yolo_network.'):
param_dict_new[key[13:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
args.logger.info('load_model {} success'.format(args.pretrained))
else:
args.logger.info('{} not exists or not a pre-trained file'.format(
args.pretrained))
assert FileNotFoundError(
'{} not exists or not a pre-trained file'.format(args.pretrained))
exit(1)
data_root = args.data_root
ann_file = args.annFile
config = ConfigYOLOV3DarkNet53()
if args.testing_shape:
config.test_img_shape = conver_testing_shape(args)
ds, data_size = create_yolo_dataset(data_root,
ann_file,
is_training=False,
batch_size=args.per_batch_size,
max_epoch=1,
device_num=1,
rank=rank_id,
shuffle=False,
config=config)
args.logger.info('testing shape : {}'.format(config.test_img_shape))
args.logger.info('totol {} images to eval'.format(data_size))
network.set_train(False)
# init detection engine
detection = DetectionEngine(args)
args.logger.info('Start inference....')
for i, data in enumerate(ds.create_dict_iterator(num_epochs=1)):
image = data["image"]
image_shape = data["image_shape"]
image_id = data["img_id"]
prediction = network(image)
output_big, output_me, output_small = prediction
output_big = output_big.asnumpy()
output_me = output_me.asnumpy()
output_small = output_small.asnumpy()
image_id = image_id.asnumpy()
image_shape = image_shape.asnumpy()
detection.detect([output_small, output_me, output_big],
args.per_batch_size, image_shape, image_id)
if i % 1000 == 0:
args.logger.info('Processing... {:.2f}% '.format(
i * args.per_batch_size / data_size * 100))
args.logger.info('Calculating mAP...')
detection.do_nms_for_results()
result_file_path = detection.write_result()
args.logger.info('result file path: {}'.format(result_file_path))
eval_result = detection.get_eval_result()
cost_time = time.time() - start_time
args.logger.info('\n=============coco eval result=========\n' +
eval_result)
args.logger.info('testing cost time {:.2f}h'.format(cost_time / 3600.))
def main():
parser = argparse.ArgumentParser(description="SSD training")
parser.add_argument("--only_create_dataset", type=ast.literal_eval, default=False,
help="If set it true, only create Mindrecord, default is False.")
parser.add_argument("--distribute", type=ast.literal_eval, default=False,
help="Run distribute, default is False.")
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("--lr", type=float, default=0.05, help="Learning rate, default is 0.05.")
parser.add_argument("--mode", type=str, default="sink", help="Run sink mode or not, default is sink.")
parser.add_argument("--dataset", type=str, default="coco", help="Dataset, defalut is coco.")
parser.add_argument("--epoch_size", type=int, default=500, help="Epoch size, default is 500.")
parser.add_argument("--batch_size", type=int, default=32, help="Batch size, default is 32.")
parser.add_argument("--pre_trained", type=str, default=None, help="Pretrained Checkpoint file path.")
parser.add_argument("--pre_trained_epoch_size", type=int, default=0, help="Pretrained epoch size.")
parser.add_argument("--save_checkpoint_epochs", type=int, default=10, help="Save checkpoint epochs, default is 10.")
parser.add_argument("--loss_scale", type=int, default=1024, help="Loss scale, default is 1024.")
parser.add_argument("--filter_weight", type=ast.literal_eval, default=False,
help="Filter weight parameters, default is False.")
args_opt = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id)
if args_opt.distribute:
device_num = args_opt.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True,
device_num=device_num)
init()
rank = args_opt.device_id % device_num
else:
rank = 0
device_num = 1
print("Start create dataset!")
# It will generate mindrecord file in args_opt.mindrecord_dir,
# and the file name is ssd.mindrecord0, 1, ... file_num.
prefix = "ssd.mindrecord"
mindrecord_dir = config.mindrecord_dir
mindrecord_file = os.path.join(mindrecord_dir, prefix + "0")
if not os.path.exists(mindrecord_file):
if not os.path.isdir(mindrecord_dir):
os.makedirs(mindrecord_dir)
if args_opt.dataset == "coco":
if os.path.isdir(config.coco_root):
print("Create Mindrecord.")
data_to_mindrecord_byte_image("coco", True, prefix)
print("Create Mindrecord Done, at {}".format(mindrecord_dir))
else:
print("coco_root not exits.")
elif args_opt.dataset == "voc":
if os.path.isdir(config.voc_dir):
print("Create Mindrecord.")
voc_data_to_mindrecord(mindrecord_dir, True, prefix)
print("Create Mindrecord Done, at {}".format(mindrecord_dir))
else:
print("voc_dir not exits.")
else:
if os.path.isdir(config.image_dir) and os.path.exists(config.anno_path):
print("Create Mindrecord.")
data_to_mindrecord_byte_image("other", True, prefix)
print("Create Mindrecord Done, at {}".format(mindrecord_dir))
else:
print("image_dir or anno_path not exits.")
if not args_opt.only_create_dataset:
loss_scale = float(args_opt.loss_scale)
# When create MindDataset, using the fitst mindrecord file, such as ssd.mindrecord0.
dataset = create_ssd_dataset(mindrecord_file, repeat_num=1,
batch_size=args_opt.batch_size, device_num=device_num, rank=rank)
dataset_size = dataset.get_dataset_size()
print("Create dataset done!")
backbone = ssd_mobilenet_v2()
ssd = SSD300(backbone=backbone, config=config)
net = SSDWithLossCell(ssd, config)
init_net_param(net)
# checkpoint
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_size * args_opt.save_checkpoint_epochs)
ckpoint_cb = ModelCheckpoint(prefix="ssd", directory=None, config=ckpt_config)
if args_opt.pre_trained:
if args_opt.pre_trained_epoch_size <= 0:
raise KeyError("pre_trained_epoch_size must be greater than 0.")
param_dict = load_checkpoint(args_opt.pre_trained)
if args_opt.filter_weight:
filter_checkpoint_parameter(param_dict)
load_param_into_net(net, param_dict)
lr = Tensor(get_lr(global_step=config.global_step,
lr_init=config.lr_init, lr_end=config.lr_end_rate * args_opt.lr, lr_max=args_opt.lr,
warmup_epochs=config.warmup_epochs,
total_epochs=args_opt.epoch_size,
steps_per_epoch=dataset_size))
opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr,
config.momentum, config.weight_decay, loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
callback = [TimeMonitor(data_size=dataset_size), LossMonitor(), ckpoint_cb]
model = Model(net)
dataset_sink_mode = False
if args_opt.mode == "sink":
print("In sink mode, one epoch return a loss.")
dataset_sink_mode = True
print("Start train SSD, the first epoch will be slower because of the graph compilation.")
model.train(args_opt.epoch_size, dataset, callbacks=callback, dataset_sink_mode=dataset_sink_mode)
def run_pretrain():
"""pre-train bert_clue"""
parser = argparse.ArgumentParser(description='bert pre_training')
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="1",
help="Epoch size, default is 1.")
parser.add_argument("--device_id",
type=int,
default=4,
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("--enable_save_ckpt",
type=str,
default="true",
help="Enable save checkpoint, default is true.")
parser.add_argument("--enable_lossscale",
type=str,
default="false",
help="Use lossscale or not, default is not.")
parser.add_argument("--do_shuffle",
type=str,
default="false",
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="100",
help="Sink steps for each epoch, default is 1.")
parser.add_argument("--save_checkpoint_path",
type=str,
default="",
help="Save checkpoint path")
parser.add_argument("--load_checkpoint_path",
type=str,
default="",
help="Load checkpoint file path")
parser.add_argument("--save_checkpoint_steps",
type=int,
default=1000,
help="Save checkpoint steps, "
"default is 1000.")
parser.add_argument("--train_steps",
type=int,
default=-1,
help="Training Steps, default is -1, "
"meaning run all steps according to epoch number.")
parser.add_argument("--save_checkpoint_num",
type=int,
default=1,
help="Save checkpoint numbers, default is 1.")
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,
save_graphs=True)
context.set_context(reserve_class_name_in_scope=False)
context.set_context(variable_memory_max_size="30GB")
ckpt_save_dir = args_opt.save_checkpoint_path
if args_opt.distribute == "true":
if args_opt.device_target == 'Ascend':
D.init('hccl')
device_num = args_opt.device_num
rank = args_opt.device_id % device_num
else:
D.init('nccl')
device_num = D.get_group_size()
rank = D.get_rank()
ckpt_save_dir = args_opt.save_checkpoint_path + 'ckpt_' + str(
rank) + '/'
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
mirror_mean=True,
device_num=device_num)
from mindspore.parallel._auto_parallel_context import auto_parallel_context
if bert_net_cfg.num_hidden_layers == 12:
if bert_net_cfg.use_relative_positions:
auto_parallel_context().set_all_reduce_fusion_split_indices(
[29, 58, 87, 116, 145, 174, 203, 217],
"hccl_world_groupsum1")
auto_parallel_context().set_all_reduce_fusion_split_indices(
[29, 58, 87, 116, 145, 174, 203, 217],
"hccl_world_groupsum3")
else:
auto_parallel_context().set_all_reduce_fusion_split_indices(
[28, 55, 82, 109, 136, 163, 190, 205],
"hccl_world_groupsum1")
auto_parallel_context().set_all_reduce_fusion_split_indices(
[28, 55, 82, 109, 136, 163, 190, 205],
"hccl_world_groupsum3")
elif bert_net_cfg.num_hidden_layers == 24:
if bert_net_cfg.use_relative_positions:
auto_parallel_context().set_all_reduce_fusion_split_indices(
[30, 90, 150, 210, 270, 330, 390, 421],
"hccl_world_groupsum1")
auto_parallel_context().set_all_reduce_fusion_split_indices(
[30, 90, 150, 210, 270, 330, 390, 421],
"hccl_world_groupsum3")
else:
auto_parallel_context().set_all_reduce_fusion_split_indices(
[38, 93, 148, 203, 258, 313, 368, 397],
"hccl_world_groupsum1")
auto_parallel_context().set_all_reduce_fusion_split_indices(
[38, 93, 148, 203, 258, 313, 368, 397],
"hccl_world_groupsum3")
else:
rank = 0
device_num = 1
if args_opt.device_target == 'GPU' and bert_net_cfg.compute_type != mstype.float32:
logger.warning('Gpu only support fp32 temporarily, run with fp32.')
bert_net_cfg.compute_type = mstype.float32
ds = create_bert_dataset(device_num, rank, args_opt.do_shuffle,
args_opt.data_dir, args_opt.schema_dir)
net_with_loss = BertNetworkWithLoss(bert_net_cfg, True)
new_repeat_count = args_opt.epoch_size * ds.get_dataset_size(
) // args_opt.data_sink_steps
if args_opt.train_steps > 0:
new_repeat_count = min(
new_repeat_count, args_opt.train_steps // args_opt.data_sink_steps)
else:
args_opt.train_steps = args_opt.epoch_size * ds.get_dataset_size()
logger.info("train steps: {}".format(args_opt.train_steps))
if cfg.optimizer == 'Lamb':
lr_schedule = BertLearningRate(
learning_rate=cfg.Lamb.learning_rate,
end_learning_rate=cfg.Lamb.end_learning_rate,
warmup_steps=cfg.Lamb.warmup_steps,
decay_steps=args_opt.train_steps,
power=cfg.Lamb.power)
params = net_with_loss.trainable_params()
decay_params = list(filter(cfg.Lamb.decay_filter, params))
other_params = list(
filter(lambda x: not cfg.Lamb.decay_filter(x), params))
group_params = [{
'params': decay_params,
'weight_decay': cfg.Lamb.weight_decay
}, {
'params': other_params
}, {
'order_params': params
}]
optimizer = Lamb(group_params,
learning_rate=lr_schedule,
eps=cfg.Lamb.eps)
elif cfg.optimizer == 'Momentum':
optimizer = Momentum(net_with_loss.trainable_params(),
learning_rate=cfg.Momentum.learning_rate,
momentum=cfg.Momentum.momentum)
elif cfg.optimizer == 'AdamWeightDecay':
lr_schedule = BertLearningRate(
learning_rate=cfg.AdamWeightDecay.learning_rate,
end_learning_rate=cfg.AdamWeightDecay.end_learning_rate,
warmup_steps=cfg.AdamWeightDecay.warmup_steps,
decay_steps=args_opt.train_steps,
power=cfg.AdamWeightDecay.power)
params = net_with_loss.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
}, {
'order_params': params
}]
optimizer = AdamWeightDecay(group_params,
learning_rate=lr_schedule,
eps=cfg.AdamWeightDecay.eps)
elif cfg.optimizer == "Thor":
lr = get_bert_lr()
damping = get_bert_damping()
optimizer = THOR(
filter(lambda x: x.requires_grad,
net_with_loss.get_parameters()), lr, cfg.Thor.momentum,
filter(lambda x: 'matrix_A' in x.name,
net_with_loss.get_parameters()),
filter(lambda x: 'matrix_G' in x.name,
net_with_loss.get_parameters()),
filter(lambda x: 'A_inv_max' in x.name,
net_with_loss.get_parameters()),
filter(lambda x: 'G_inv_max' in x.name,
net_with_loss.get_parameters()), cfg.Thor.weight_decay,
cfg.Thor.loss_scale, bert_net_cfg.num_hidden_layers,
bert_net_cfg.batch_size, damping)
else:
raise ValueError(
"Don't support optimizer {}, only support [Lamb, Momentum, AdamWeightDecay]"
.format(cfg.optimizer))
callback = [TimeMonitor(args_opt.data_sink_steps), LossCallBack()]
if args_opt.enable_save_ckpt == "true":
config_ck = CheckpointConfig(
save_checkpoint_steps=args_opt.save_checkpoint_steps,
keep_checkpoint_max=args_opt.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(prefix='checkpoint_bert',
directory=ckpt_save_dir,
config=config_ck)
callback.append(ckpoint_cb)
if args_opt.load_checkpoint_path:
param_dict = load_checkpoint(args_opt.load_checkpoint_path)
load_param_into_net(net_with_loss, param_dict)
if args_opt.enable_lossscale == "true":
update_cell = DynamicLossScaleUpdateCell(
loss_scale_value=cfg.loss_scale_value,
scale_factor=cfg.scale_factor,
scale_window=cfg.scale_window)
net_with_grads = BertTrainOneStepWithLossScaleCell(
net_with_loss, optimizer=optimizer, scale_update_cell=update_cell)
else:
net_with_grads = BertTrainOneStepCell(net_with_loss,
optimizer=optimizer)
model = Model(net_with_grads, frequency=cfg.Thor.frequency)
model.train(new_repeat_count,
ds,
callbacks=callback,
dataset_sink_mode=(args_opt.enable_data_sink == "true"),
sink_size=args_opt.data_sink_steps)
def set_parallel_env():
context.reset_auto_parallel_context()
MultiDevice.init()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL,
device_num=MultiDevice.get_group_size(),
gradients_mean=True)
def _test_context():
context.set_context(enable_sparse=True)
yield
context.set_context(enable_sparse=False)
context.reset_auto_parallel_context()
def run_train_pipeline(args_opt):
device_id = int(os.getenv("DEVICE_ID"))
rank_id = int(os.getenv("RANK_ID"))
local_rank = rank_id
print('local_rank:{}, device id:{} start to run...'.format(
local_rank, device_id),
flush=True)
context.set_context(save_graphs=False,
mode=context.GRAPH_MODE,
device_target="Ascend",
device_id=device_id)
context.set_context(variable_memory_max_size="31GB")
strategy_ckpt_save_file = "/cache/" + "strategy" + str(
local_rank) + ".ckpt"
if args_opt.distribute == "true":
D.init()
device_num = D.get_group_size()
rank = D.get_rank()
print("device_id is {}, rank_id is {}, device_num is {}".format(
device_id, rank, device_num))
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL,
gradients_mean=False,
device_num=device_num,
full_batch=True,
loss_repeated_mean=True,
enable_parallel_optimizer=bool(args_opt.optimizer_shard),
pipeline_stages=args_opt.stage_num,
strategy_ckpt_save_file=strategy_ckpt_save_file)
set_algo_parameters(elementwise_op_strategy_follow=True)
_set_multi_subgraphs()
else:
rank = 0
device_num = 1
model_parallel_num = args_opt.tensor_model_parallel_num
stage_device_num = int(device_num / args_opt.stage_num)
data_parallel_num = int(stage_device_num / model_parallel_num)
per_batch_size = args_opt.per_batch_size
batch_size = per_batch_size * data_parallel_num * args_opt.micro_size
config = PANGUALPHAConfig(data_parallel_num=data_parallel_num,
model_parallel_num=model_parallel_num,
batch_size=batch_size,
seq_length=args_opt.seq_length,
vocab_size=args_opt.vocab_size,
embedding_size=args_opt.embedding_size,
num_layers=args_opt.num_layers,
num_heads=args_opt.num_heads,
expand_ratio=4,
post_layernorm_residual=False,
dropout_rate=0.1,
compute_dtype=mstype.float16,
use_past=False,
self_layernorm=True,
forward_reduce_scatter=True,
stage_num=args_opt.stage_num,
micro_size=args_opt.micro_size,
word_emb_dp=False)
print("===config is: ", config, flush=True)
pangu_alpha = PANGUALPHAPipeline(config)
loss = CrossEntropyLoss(config)
pangu_alpha_with_loss = PANGUALPHAWithLossPipeline(config, pangu_alpha,
loss)
pangu_alpha_with_loss = VirtualDatasetOneInputCell(pangu_alpha_with_loss)
print("=====args_opt is: ", args_opt, flush=True)
lr = LearningRate(learning_rate=args_opt.start_lr,
end_learning_rate=args_opt.end_lr,
warmup_steps=args_opt.warmup_step,
decay_steps=args_opt.decay_steps)
per_stage_layers = config.num_layers // config.stage_num
per_stage_devices = device_num // config.stage_num
self_stage = rank_id // per_stage_devices
range_min = self_stage * per_stage_layers
range_max = range_min + per_stage_layers
if self_stage == 0:
params = [pangu_alpha.embedding_table]
params.extend(pangu_alpha.backbone.pangu_alpha_embedding.
position_embedding.trainable_params())
elif self_stage == config.stage_num - 1:
params = [pangu_alpha.embedding_table]
params.extend(pangu_alpha.backbone.layernorm.trainable_params())
params.extend(
pangu_alpha.backbone.top_query_embedding.trainable_params())
else:
params = []
for i in range(range_min, range_max):
params.extend(pangu_alpha.backbone.blocks[i].trainable_params())
decay_filter = lambda x: 'layernorm' not in x.name.lower(
) and "bias" not in x.name.lower()
decay_params = list(filter(decay_filter, params))
other_params = list(filter(lambda x: not decay_filter(x), params))
group_params = [{
'params': decay_params,
'weight_decay': args_opt.weight_decay
}, {
'params': other_params,
'weight_decay': 0.0
}, {
'order_params': params
}]
if args_opt.optimizer == "lamb":
optimizer = nn.Lamb(group_params, learning_rate=lr)
else:
optimizer = nn.AdamWeightDecay(group_params,
learning_rate=lr,
beta1=0.9,
beta2=0.95,
eps=1e-8)
save_steps = args_opt.save_steps
ckpt_dir = os.path.join(args_opt.ckpt_save_sir, f"rank_{str(local_rank)}")
if not os.path.exists(ckpt_dir):
Path(ckpt_dir).mkdir(parents=True, exist_ok=True)
ds = create_dataset(config.batch_size,
data_path=args_opt.data_url,
data_start_index=0)
epoch_num = args_opt.epoch_size
step_per_epoch = ds.get_dataset_size()
callback_size = args_opt.sink_size
actual_epoch_num = int(epoch_num * step_per_epoch / callback_size)
callback = [
TimeMonitor(callback_size),
LossCallBack(callback_size, local_rank, config.stage_num)
]
config_ck = CheckpointConfig(save_checkpoint_steps=save_steps,
keep_checkpoint_max=1,
integrated_save=False,
filter_prefix="accu_grads")
ckpoint_cb = ModelCheckpoint(prefix="PanguAlpha",
directory=ckpt_dir,
config=config_ck)
callback.append(ckpoint_cb)
loss_scale_value = math.pow(2, 32)
update_cell = DynamicLossScaleUpdateCell(loss_scale_value=loss_scale_value,
scale_factor=2,
scale_window=1000)
pangu_alpha_with_grads = PANGUALPHATrainPipelineWithLossScaleCell(
pangu_alpha_with_loss,
optimizer=optimizer,
config=config,
scale_update_cell=update_cell)
model = Model(pangu_alpha_with_grads)
de.config.set_sending_batches(2 * args_opt.sink_size)
model.train(actual_epoch_num,
ds,
callbacks=callback,
sink_size=callback_size,
dataset_sink_mode=True)
def compile(net):
_executor.compile(net, _x, _b)
context.reset_auto_parallel_context()
def train_process_bert_thor(q, device_id, epoch_size, device_num):
os.system("mkdir " + str(device_id))
os.chdir(str(device_id))
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
device_id=device_id,
save_graphs=False)
context.set_context(reserve_class_name_in_scope=False)
context.set_context(max_call_depth=3000)
os.environ['MINDSPORE_HCCL_CONFIG_PATH'] = MINDSPORE_HCCL_CONFIG_PATH
os.environ['RANK_ID'] = str(device_id)
os.environ['RANK_SIZE'] = str(device_num)
D.init()
rank = device_id % device_num
context.reset_auto_parallel_context()
_set_bert_all_reduce_split()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
device_num=device_num)
bert_net_cfg.num_hidden_layers = 4
data_set = create_bert_dataset(device_num=device_num,
rank=rank,
do_shuffle=False,
data_dir=DATASET_PATH,
schema_dir=None)
net_with_loss = BertNetworkWithLoss(bert_net_cfg, True)
new_repeat_count = epoch_size * data_set.get_dataset_size(
) // data_sink_steps
new_repeat_count = min(new_repeat_count, train_steps // data_sink_steps)
lr = get_bert_lr()
damping = get_bert_damping()
optimizer = THOR(
filter(lambda x: x.requires_grad,
net_with_loss.get_parameters()), lr, cfg.Thor.momentum,
filter(lambda x: 'matrix_A' in x.name, net_with_loss.get_parameters()),
filter(lambda x: 'matrix_G' in x.name, net_with_loss.get_parameters()),
cfg.Thor.weight_decay, cfg.Thor.loss_scale,
bert_net_cfg.num_hidden_layers, bert_net_cfg.batch_size, damping)
time_monitor_callback = TimeMonitor(data_sink_steps)
loss_callback = LossCallback()
callback = [time_monitor_callback, loss_callback]
if load_checkpoint_path:
param_dict = load_checkpoint(load_checkpoint_path)
load_param_into_net(net_with_loss, param_dict)
net_with_grads = BertTrainOneStepCell(net_with_loss, optimizer=optimizer)
model = Model(net_with_grads, frequency=cfg.Thor.frequency)
model.train(new_repeat_count,
data_set,
callbacks=callback,
dataset_sink_mode=True,
sink_size=data_sink_steps)
loss_list = loss_callback.loss_list
per_step_mseconds = time_monitor_callback.per_step_mseconds_list
q.put({'loss': loss_list, 'cost': per_step_mseconds})
def test_yolov3_darknet53():
devid = int(os.getenv('DEVICE_ID')) if os.getenv('DEVICE_ID') else 0
context.set_context(mode=context.GRAPH_MODE, enable_auto_mixed_precision=True,
device_target="Ascend", device_id=devid)
rank = 0
device_num = 1
lr_init = 0.001
epoch_size = 3
batch_size = 32
loss_scale = 1024
mindrecord_dir = DATA_DIR
# It will generate mindrecord file in args_opt.mindrecord_dir,
# and the file name is yolo.mindrecord0, 1, ... file_num.
if not os.path.isdir(mindrecord_dir):
raise KeyError("mindrecord path is not exist.")
data_root = os.path.join(mindrecord_dir, 'train2014')
annFile = os.path.join(mindrecord_dir, 'annotations/instances_train2014.json')
# print("yolov3 mindrecord is ", mindrecord_file)
if not os.path.exists(annFile):
print("instances_train2014 file is not exist.")
assert False
loss_meter = AverageMeter('loss')
context.reset_auto_parallel_context()
parallel_mode = ParallelMode.STAND_ALONE
context.set_auto_parallel_context(parallel_mode=parallel_mode, gradients_mean=True, device_num=1)
network = YOLOV3DarkNet53(is_training=True)
# default is kaiming-normal
default_recurisive_init(network)
network = YoloWithLossCell(network)
print('finish get network')
config = ConfigYOLOV3DarkNet53()
label_smooth = 0
label_smooth_factor = 0.1
config.label_smooth = label_smooth
config.label_smooth_factor = label_smooth_factor
# When create MindDataset, using the fitst mindrecord file, such as yolo.mindrecord0.
print("Create dataset begin!")
training_shape = [int(416), int(416)]
config.multi_scale = [training_shape]
num_samples = 256
ds, data_size = create_yolo_dataset(image_dir=data_root, anno_path=annFile, is_training=True,
batch_size=batch_size, max_epoch=epoch_size,
device_num=device_num, rank=rank, config=config, num_samples=num_samples)
print("Create dataset done!")
per_batch_size = batch_size
group_size = 1
print("data_size:", data_size)
steps_per_epoch = int(data_size / per_batch_size / group_size)
print("steps_per_epoch:", steps_per_epoch)
warmup_epochs = 0.
max_epoch = epoch_size
T_max = 1
eta_min = 0
lr = warmup_cosine_annealing_lr(lr_init,
steps_per_epoch,
warmup_epochs,
max_epoch,
T_max,
eta_min)
opt = Momentum(params=get_param_groups(network),
learning_rate=Tensor(lr),
momentum=0.9,
weight_decay=0.0005,
loss_scale=loss_scale)
network = TrainingWrapper(network, opt)
network.set_train()
old_progress = -1
t_end = time.time()
data_loader = ds.create_dict_iterator(output_numpy=True)
train_starttime = time.time()
time_used_per_epoch = 0
print("time:", time.time())
for i, data in enumerate(data_loader):
images = data["image"]
input_shape = images.shape[2:4]
print('iter[{}], shape{}'.format(i, input_shape[0]))
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())
if (i + 1) % steps_per_epoch == 0:
time_used = time.time() - t_end
epoch = int(i / steps_per_epoch)
fps = per_batch_size * (i - old_progress) * group_size / time_used
if rank == 0:
print(
'epoch[{}], iter[{}], {}, {:.2f} imgs/sec, lr:{}, time_used:{}'.format(epoch,
i, loss_meter, fps, lr[i],
time_used))
t_end = time.time()
loss_meter.reset()
old_progress = i
time_used_per_epoch = time_used
train_endtime = time.time() - train_starttime
print('train_time_used:{}'.format(train_endtime))
expect_loss_value = 3210.0
loss_value = re.findall(r"\d+\.?\d*", str(loss_meter))
print('loss_value:{}'.format(loss_value[0]))
assert float(loss_value[0]) < expect_loss_value
export_time_used = 20.0
print('time_used_per_epoch:{}'.format(time_used_per_epoch))
assert time_used_per_epoch < export_time_used
print('==========test case passed===========')
def train_alexnet():
print(config)
print('device id:', get_device_id())
print('device num:', get_device_num())
print('rank id:', get_rank_id())
print('job id:', get_job_id())
device_target = config.device_target
context.set_context(mode=context.GRAPH_MODE,
device_target=config.device_target)
context.set_context(save_graphs=False)
device_num = get_device_num()
if config.dataset_name == "cifar10":
if device_num > 1:
config.learning_rate = config.learning_rate * device_num
config.epoch_size = config.epoch_size * 2
elif config.dataset_name == "imagenet":
pass
else:
raise ValueError("Unsupported dataset.")
if device_num > 1:
context.reset_auto_parallel_context()
context.set_auto_parallel_context(device_num=device_num, \
parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True)
if device_target == "Ascend":
context.set_context(device_id=get_device_id())
init()
elif device_target == "GPU":
init()
else:
context.set_context(device_id=get_device_id())
if config.dataset_name == "cifar10":
ds_train = create_dataset_cifar10(config.data_path,
config.batch_size,
target=config.device_target)
elif config.dataset_name == "imagenet":
ds_train = create_dataset_imagenet(config.data_path, config.batch_size)
else:
raise ValueError("Unsupported dataset.")
if ds_train.get_dataset_size() == 0:
raise ValueError(
"Please check dataset size > 0 and batch_size <= dataset size")
network = AlexNet(config.num_classes, phase='train')
loss_scale_manager = None
metrics = None
step_per_epoch = ds_train.get_dataset_size(
) if config.sink_size == -1 else config.sink_size
if config.dataset_name == 'cifar10':
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
lr = Tensor(
get_lr_cifar10(0, config.learning_rate, config.epoch_size,
step_per_epoch))
opt = nn.Momentum(network.trainable_params(), lr, config.momentum)
metrics = {"Accuracy": Accuracy()}
elif config.dataset_name == 'imagenet':
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
lr = Tensor(
get_lr_imagenet(config.learning_rate, config.epoch_size,
step_per_epoch))
opt = nn.Momentum(params=get_param_groups(network),
learning_rate=lr,
momentum=config.momentum,
weight_decay=config.weight_decay,
loss_scale=config.loss_scale)
from mindspore.train.loss_scale_manager import DynamicLossScaleManager, FixedLossScaleManager
if config.is_dynamic_loss_scale == 1:
loss_scale_manager = DynamicLossScaleManager(init_loss_scale=65536,
scale_factor=2,
scale_window=2000)
else:
loss_scale_manager = FixedLossScaleManager(
config.loss_scale, drop_overflow_update=False)
else:
raise ValueError("Unsupported dataset.")
if device_target == "Ascend":
model = Model(network,
loss_fn=loss,
optimizer=opt,
metrics=metrics,
amp_level="O2",
keep_batchnorm_fp32=False,
loss_scale_manager=loss_scale_manager)
elif device_target == "GPU":
model = Model(network,
loss_fn=loss,
optimizer=opt,
metrics=metrics,
loss_scale_manager=loss_scale_manager)
else:
raise ValueError("Unsupported platform.")
if device_num > 1:
ckpt_save_dir = os.path.join(config.checkpoint_path + "_" +
str(get_rank()))
else:
ckpt_save_dir = config.checkpoint_path
time_cb = TimeMonitor(data_size=step_per_epoch)
config_ck = CheckpointConfig(
save_checkpoint_steps=config.save_checkpoint_steps,
keep_checkpoint_max=config.keep_checkpoint_max)
ckpoint_cb = ModelCheckpoint(prefix="checkpoint_alexnet",
directory=ckpt_save_dir,
config=config_ck)
print("============== Starting Training ==============")
model.train(config.epoch_size,
ds_train,
callbacks=[time_cb, ckpoint_cb,
LossMonitor()],
dataset_sink_mode=config.dataset_sink_mode,
sink_size=config.sink_size)
def train():
"""Train function."""
args = parse_args()
args.logger.save_args(args)
if args.need_profiler:
from mindspore.profiler.profiling import Profiler
profiler = Profiler(output_path=args.outputs_dir,
is_detail=True,
is_show_op_path=True)
loss_meter = AverageMeter('loss')
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 = YOLOV3DarkNet53(is_training=True)
# default is kaiming-normal
default_recurisive_init(network)
load_yolov3_quant_params(args, network)
config = ConfigYOLOV3DarkNet53()
# convert fusion network to quantization aware network
if config.quantization_aware:
network = quant.convert_quant_network(network,
bn_fold=True,
per_channel=[True, False],
symmetric=[True, False])
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 = [conver_training_shape(args)]
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,
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
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)
network = TrainingWrapper(network, opt)
network.set_train()
if args.rank_save_ckpt_flag:
# checkpoint save
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)
save_ckpt_path = os.path.join(args.outputs_dir,
'ckpt_' + str(args.rank) + '/')
ckpt_cb = ModelCheckpoint(config=ckpt_config,
directory=save_ckpt_path,
prefix='{}'.format(args.rank))
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)
shape_record = ShapeRecord()
for i, data in enumerate(data_loader):
images = data["image"]
input_shape = images.shape[2:4]
args.logger.info('iter[{}], shape{}'.format(i, input_shape[0]))
shape_record.set(input_shape)
images = Tensor.from_numpy(images)
annos = data["annotation"]
if args.group_size == 1:
batch_y_true_0, batch_y_true_1, batch_y_true_2, batch_gt_box0, batch_gt_box1, batch_gt_box2 = \
batch_preprocess_true_box(annos, config, input_shape)
else:
batch_y_true_0, batch_y_true_1, batch_y_true_2, batch_gt_box0, batch_gt_box1, batch_gt_box2 = \
batch_preprocess_true_box_single(annos, config, input_shape)
batch_y_true_0 = Tensor.from_numpy(batch_y_true_0)
batch_y_true_1 = Tensor.from_numpy(batch_y_true_1)
batch_y_true_2 = Tensor.from_numpy(batch_y_true_2)
batch_gt_box0 = Tensor.from_numpy(batch_gt_box0)
batch_gt_box1 = Tensor.from_numpy(batch_gt_box1)
batch_gt_box2 = Tensor.from_numpy(batch_gt_box2)
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())
if args.rank_save_ckpt_flag:
# 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 and args.rank_save_ckpt_flag:
cb_params.cur_epoch_num += 1
if args.need_profiler:
if i == 10:
profiler.analyse()
break
args.logger.info('==========end training===============')
def run_pretrain():
"""pre-train bert_clue"""
parser = argparse.ArgumentParser(description='bert pre_training')
parser.add_argument("--distribute",
type=str,
default="false",
help="Run distribute, default is false.")
parser.add_argument("--epoch_size",
type=int,
default="1",
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("--enable_save_ckpt",
type=str,
default="true",
help="Enable save checkpoint, default is true.")
parser.add_argument("--enable_lossscale",
type=str,
default="true",
help="Use lossscale or not, default is not.")
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("--checkpoint_path",
type=str,
default="",
help="Checkpoint file path")
parser.add_argument("--save_checkpoint_steps",
type=int,
default=1000,
help="Save checkpoint steps, "
"default is 1000.")
parser.add_argument("--save_checkpoint_num",
type=int,
default=1,
help="Save checkpoint numbers, default is 1.")
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="Ascend",
device_id=args_opt.device_id)
context.set_context(reserve_class_name_in_scope=False)
if args_opt.distribute == "true":
device_num = args_opt.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
mirror_mean=True,
device_num=device_num)
D.init()
rank = args_opt.device_id % device_num
else:
rank = 0
device_num = 1
ds, new_repeat_count = create_bert_dataset(args_opt.epoch_size, device_num,
rank, args_opt.do_shuffle,
args_opt.enable_data_sink,
args_opt.data_sink_steps,
args_opt.data_dir,
args_opt.schema_dir)
netwithloss = BertNetworkWithLoss(bert_net_cfg, True)
if cfg.optimizer == 'Lamb':
optimizer = Lamb(netwithloss.trainable_params(),
decay_steps=ds.get_dataset_size() *
ds.get_repeat_count(),
start_learning_rate=cfg.Lamb.start_learning_rate,
end_learning_rate=cfg.Lamb.end_learning_rate,
power=cfg.Lamb.power,
warmup_steps=cfg.Lamb.warmup_steps,
weight_decay=cfg.Lamb.weight_decay,
eps=cfg.Lamb.eps)
elif cfg.optimizer == 'Momentum':
optimizer = Momentum(netwithloss.trainable_params(),
learning_rate=cfg.Momentum.learning_rate,
momentum=cfg.Momentum.momentum)
elif cfg.optimizer == 'AdamWeightDecayDynamicLR':
optimizer = AdamWeightDecayDynamicLR(
netwithloss.trainable_params(),
decay_steps=ds.get_dataset_size() * ds.get_repeat_count(),
learning_rate=cfg.AdamWeightDecayDynamicLR.learning_rate,
end_learning_rate=cfg.AdamWeightDecayDynamicLR.end_learning_rate,
power=cfg.AdamWeightDecayDynamicLR.power,
weight_decay=cfg.AdamWeightDecayDynamicLR.weight_decay,
eps=cfg.AdamWeightDecayDynamicLR.eps,
warmup_steps=cfg.AdamWeightDecayDynamicLR.warmup_steps)
else:
raise ValueError(
"Don't support optimizer {}, only support [Lamb, Momentum, AdamWeightDecayDynamicLR]"
.format(cfg.optimizer))
callback = [TimeMonitor(ds.get_dataset_size()), LossCallBack()]
if args_opt.enable_save_ckpt == "true":
config_ck = CheckpointConfig(
save_checkpoint_steps=args_opt.save_checkpoint_steps,
keep_checkpoint_max=args_opt.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(prefix='checkpoint_bert',
config=config_ck)
callback.append(ckpoint_cb)
if args_opt.checkpoint_path:
param_dict = load_checkpoint(args_opt.checkpoint_path)
load_param_into_net(netwithloss, param_dict)
if args_opt.enable_lossscale == "true":
update_cell = DynamicLossScaleUpdateCell(
loss_scale_value=cfg.loss_scale_value,
scale_factor=cfg.scale_factor,
scale_window=cfg.scale_window)
netwithgrads = BertTrainOneStepWithLossScaleCell(
netwithloss, optimizer=optimizer, scale_update_cell=update_cell)
else:
netwithgrads = BertTrainOneStepCell(netwithloss, optimizer=optimizer)
model = Model(netwithgrads)
model.train(new_repeat_count,
ds,
callbacks=callback,
dataset_sink_mode=(args_opt.enable_data_sink == "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=True, device_id=devid)
if args.need_profiler:
from mindspore.profiler.profiling import Profiler
profiler = Profiler(output_path=args.outputs_dir, is_detail=True, is_show_op_path=True)
loss_meter = AverageMeter('loss')
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 = YOLOV3DarkNet53(is_training=True)
# default is kaiming-normal
default_recurisive_init(network)
load_yolov3_params(args, network)
network = YoloWithLossCell(network)
args.logger.info('finish get network')
config = ConfigYOLOV3DarkNet53()
config.label_smooth = args.label_smooth
config.label_smooth_factor = args.label_smooth_factor
if args.training_shape:
config.multi_scale = [conver_training_shape(args)]
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,
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
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()
if args.rank_save_ckpt_flag:
# checkpoint save
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)
save_ckpt_path = os.path.join(args.outputs_dir, 'ckpt_' + str(args.rank) + '/')
ckpt_cb = ModelCheckpoint(config=ckpt_config,
directory=save_ckpt_path,
prefix='{}'.format(args.rank))
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]
args.logger.info('iter[{}], shape{}'.format(i, input_shape[0]))
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())
if args.rank_save_ckpt_flag:
# 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 and args.rank_save_ckpt_flag:
cb_params.cur_epoch_num += 1
if args.need_profiler:
if i == 10:
profiler.analyse()
break
args.logger.info('==========end training===============')
def train():
"""training CenterNet"""
context.set_context(mode=context.GRAPH_MODE,
device_target=args_opt.device_target)
context.set_context(reserve_class_name_in_scope=False)
context.set_context(save_graphs=False)
ckpt_save_dir = args_opt.save_checkpoint_path
rank = 0
device_num = 1
num_workers = 8
if args_opt.device_target == "Ascend":
context.set_context(enable_auto_mixed_precision=False)
context.set_context(device_id=args_opt.device_id)
if args_opt.distribute == "true":
D.init()
device_num = args_opt.device_num
rank = args_opt.device_id % device_num
ckpt_save_dir = args_opt.save_checkpoint_path + 'ckpt_' + str(
get_rank()) + '/'
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
device_num=device_num)
_set_parallel_all_reduce_split()
else:
args_opt.distribute = "false"
args_opt.need_profiler = "false"
args_opt.enable_data_sink = "false"
# Start create dataset!
# mindrecord files will be generated at args_opt.mindrecord_dir such as centernet.mindrecord0, 1, ... file_num.
logger.info("Begin creating dataset for CenterNet")
coco = COCOHP(dataset_config,
run_mode="train",
net_opt=net_config,
enable_visual_image=(args_opt.visual_image == "true"),
save_path=args_opt.save_result_dir)
dataset = coco.create_train_dataset(
args_opt.mindrecord_dir,
args_opt.mindrecord_prefix,
batch_size=train_config.batch_size,
device_num=device_num,
rank=rank,
num_parallel_workers=num_workers,
do_shuffle=args_opt.do_shuffle == 'true')
dataset_size = dataset.get_dataset_size()
logger.info("Create dataset done!")
net_with_loss = CenterNetMultiPoseLossCell(net_config)
new_repeat_count = args_opt.epoch_size * dataset_size // args_opt.data_sink_steps
if args_opt.train_steps > 0:
new_repeat_count = min(
new_repeat_count, args_opt.train_steps // args_opt.data_sink_steps)
else:
args_opt.train_steps = args_opt.epoch_size * dataset_size
logger.info("train steps: {}".format(args_opt.train_steps))
optimizer = _get_optimizer(net_with_loss, dataset_size)
enable_static_time = args_opt.device_target == "CPU"
callback = [
TimeMonitor(args_opt.data_sink_steps),
LossCallBack(dataset_size, enable_static_time)
]
if args_opt.enable_save_ckpt == "true" and args_opt.device_id % min(
8, device_num) == 0:
config_ck = CheckpointConfig(
save_checkpoint_steps=args_opt.save_checkpoint_steps,
keep_checkpoint_max=args_opt.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(
prefix='checkpoint_centernet',
directory=None if ckpt_save_dir == "" else ckpt_save_dir,
config=config_ck)
callback.append(ckpoint_cb)
if args_opt.load_checkpoint_path:
param_dict = load_checkpoint(args_opt.load_checkpoint_path)
load_param_into_net(net_with_loss, param_dict)
if args_opt.device_target == "Ascend":
net_with_grads = CenterNetWithLossScaleCell(
net_with_loss,
optimizer=optimizer,
sens=train_config.loss_scale_value)
else:
net_with_grads = CenterNetWithoutLossScaleCell(net_with_loss,
optimizer=optimizer)
model = Model(net_with_grads)
model.train(new_repeat_count,
dataset,
callbacks=callback,
dataset_sink_mode=(args_opt.enable_data_sink == "true"),
sink_size=args_opt.data_sink_steps)
def test_resnet_and_resnet_thor_imagenet_4p():
# reset context
context.set_context(save_graphs=False,
enable_graph_kernel=False,
enable_sparse=False)
context.reset_auto_parallel_context()
context.reset_ps_context()
q = Queue()
q2 = Queue()
# resnet50
device_num = 4
epoch_size = 2
epoch_size_2 = 1
enable_hccl = True
process = []
process2 = []
for i in range(device_num):
device_id = i
process.append(
Process(target=train_process,
args=(q, device_id, epoch_size, device_num, enable_hccl)))
process2.append(
Process(target=train_process_thor,
args=(q2, device_id + 4, epoch_size_2, device_num,
enable_hccl)))
for i in range(device_num):
process[i].start()
process2[i].start()
print("Waiting for all subprocesses done...")
for i in range(device_num):
process[i].join()
process2[i].join()
# resnet
acc = 0.0
cost = 0.0
for i in range(device_num):
output = q.get()
acc += output['acc']
cost += output['cost']
acc = acc / device_num
cost = cost / device_num
for i in range(device_num):
os.system("rm -rf " + str(i))
print("End training...")
assert acc > 0.15
assert cost < 26
# THOR
thor_acc = 0.0
thor_cost = 0.0
for i in range(device_num):
output = q2.get()
thor_acc += output['acc']
thor_cost += output['cost']
thor_acc = thor_acc / device_num
thor_cost = thor_cost / device_num
for i in range(4, device_num + 4):
os.system("rm -rf " + str(i))
print("End training...")
assert thor_acc > 0.25
assert thor_cost < 22
def compile_net(net):
net.set_auto_parallel()
net.set_train()
_executor.compile(net, _x, _b)
context.reset_auto_parallel_context()
def main():
args_opt = get_args()
rank = 0
device_num = 1
if args_opt.run_platform == "CPU":
context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
else:
context.set_context(mode=context.GRAPH_MODE,
device_target=args_opt.run_platform,
device_id=args_opt.device_id)
if args_opt.distribute:
device_num = args_opt.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
device_num=device_num)
init()
if config.model == "ssd_resnet50_fpn":
context.set_auto_parallel_context(
all_reduce_fusion_config=[90, 183, 279])
if config.model == "ssd_vgg16":
context.set_auto_parallel_context(
all_reduce_fusion_config=[20, 41, 62])
else:
context.set_auto_parallel_context(
all_reduce_fusion_config=[29, 58, 89])
rank = get_rank()
mindrecord_file = create_mindrecord(args_opt.dataset, "ssd.mindrecord",
True)
if args_opt.only_create_dataset:
return
loss_scale = float(args_opt.loss_scale)
if args_opt.run_platform == "CPU":
loss_scale = 1.0
# When create MindDataset, using the fitst mindrecord file, such as ssd.mindrecord0.
use_multiprocessing = (args_opt.run_platform != "CPU")
dataset = create_ssd_dataset(mindrecord_file,
repeat_num=1,
batch_size=args_opt.batch_size,
device_num=device_num,
rank=rank,
use_multiprocessing=use_multiprocessing)
dataset_size = dataset.get_dataset_size()
print(f"Create dataset done! dataset size is {dataset_size}")
ssd = ssd_model_build(args_opt)
if ("use_float16" in config
and config.use_float16) or args_opt.run_platform == "GPU":
ssd.to_float(dtype.float16)
net = SSDWithLossCell(ssd, config)
# checkpoint
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_size *
args_opt.save_checkpoint_epochs)
save_ckpt_path = './ckpt_' + str(rank) + '/'
ckpoint_cb = ModelCheckpoint(prefix="ssd",
directory=save_ckpt_path,
config=ckpt_config)
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
if args_opt.filter_weight:
filter_checkpoint_parameter_by_list(param_dict,
config.checkpoint_filter_list)
load_param_into_net(net, param_dict, True)
lr = Tensor(
get_lr(global_step=args_opt.pre_trained_epoch_size * dataset_size,
lr_init=config.lr_init,
lr_end=config.lr_end_rate * args_opt.lr,
lr_max=args_opt.lr,
warmup_epochs=config.warmup_epochs,
total_epochs=args_opt.epoch_size,
steps_per_epoch=dataset_size))
if "use_global_norm" in config and config.use_global_norm:
opt = nn.Momentum(
filter(lambda x: x.requires_grad, net.get_parameters()), lr,
config.momentum, config.weight_decay, 1.0)
net = TrainingWrapper(net, opt, loss_scale, True)
else:
opt = nn.Momentum(
filter(lambda x: x.requires_grad, net.get_parameters()), lr,
config.momentum, config.weight_decay, loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
callback = [TimeMonitor(data_size=dataset_size), LossMonitor(), ckpoint_cb]
if args_opt.run_eval:
eval_net = SsdInferWithDecoder(ssd, Tensor(default_boxes), config)
eval_net.set_train(False)
mindrecord_file = create_mindrecord(args_opt.dataset,
"ssd_eval.mindrecord", False)
eval_dataset = create_ssd_dataset(mindrecord_file,
batch_size=args_opt.batch_size,
repeat_num=1,
is_training=False,
use_multiprocessing=False)
if args_opt.dataset == "coco":
anno_json = os.path.join(
config.coco_root,
config.instances_set.format(config.val_data_type))
elif args_opt.dataset == "voc":
anno_json = os.path.join(config.voc_root, config.voc_json)
else:
raise ValueError(
'SSD eval only support dataset mode is coco and voc!')
eval_param_dict = {
"net": eval_net,
"dataset": eval_dataset,
"anno_json": anno_json
}
eval_cb = EvalCallBack(apply_eval,
eval_param_dict,
interval=args_opt.eval_interval,
eval_start_epoch=args_opt.eval_start_epoch,
save_best_ckpt=True,
ckpt_directory=save_ckpt_path,
besk_ckpt_name="best_map.ckpt",
metrics_name="mAP")
callback.append(eval_cb)
model = Model(net)
dataset_sink_mode = False
if args_opt.mode == "sink" and args_opt.run_platform != "CPU":
print("In sink mode, one epoch return a loss.")
dataset_sink_mode = True
print(
"Start train SSD, the first epoch will be slower because of the graph compilation."
)
model.train(args_opt.epoch_size,
dataset,
callbacks=callback,
dataset_sink_mode=dataset_sink_mode)
def compile_net(net):
optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
_executor.compile(train_net, _x, _b)
context.reset_auto_parallel_context()
def run_transformer_train():
"""
Transformer training.
"""
parser = argparse_init()
args, _ = parser.parse_known_args()
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
device_id=args.device_id)
context.set_context(reserve_class_name_in_scope=False,
enable_auto_mixed_precision=False)
if args.distribute == "true":
device_num = args.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
mirror_mean=True,
parameter_broadcast=True,
device_num=device_num)
D.init()
rank_id = args.device_id % device_num
else:
device_num = 1
rank_id = 0
dataset = create_transformer_dataset(
epoch_count=1,
rank_size=device_num,
rank_id=rank_id,
do_shuffle=args.do_shuffle,
enable_data_sink=args.enable_data_sink,
dataset_path=args.data_path)
netwithloss = TransformerNetworkWithLoss(transformer_net_cfg, True)
if args.checkpoint_path:
parameter_dict = load_checkpoint(args.checkpoint_path)
load_param_into_net(netwithloss, parameter_dict)
lr = Tensor(
create_dynamic_lr(
schedule="constant*rsqrt_hidden*linear_warmup*rsqrt_decay",
training_steps=dataset.get_dataset_size() * args.epoch_size,
learning_rate=cfg.lr_schedule.learning_rate,
warmup_steps=cfg.lr_schedule.warmup_steps,
hidden_size=transformer_net_cfg.hidden_size,
start_decay_step=cfg.lr_schedule.start_decay_step,
min_lr=cfg.lr_schedule.min_lr), mstype.float32)
optimizer = Adam(netwithloss.trainable_params(), lr)
callbacks = [TimeMonitor(dataset.get_dataset_size()), LossCallBack()]
if args.enable_save_ckpt == "true":
if device_num == 1 or (device_num > 1 and rank_id == 0):
ckpt_config = CheckpointConfig(
save_checkpoint_steps=args.save_checkpoint_steps,
keep_checkpoint_max=args.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(prefix='transformer',
directory=args.save_checkpoint_path,
config=ckpt_config)
callbacks.append(ckpoint_cb)
if args.enable_lossscale == "true":
scale_manager = DynamicLossScaleManager(
init_loss_scale=cfg.init_loss_scale_value,
scale_factor=cfg.scale_factor,
scale_window=cfg.scale_window)
update_cell = scale_manager.get_update_cell()
netwithgrads = TransformerTrainOneStepWithLossScaleCell(
netwithloss, optimizer=optimizer, scale_update_cell=update_cell)
else:
netwithgrads = TransformerTrainOneStepCell(netwithloss,
optimizer=optimizer)
netwithgrads.set_train(True)
model = Model(netwithgrads)
enable_sink = (args.enable_data_sink == "true")
if enable_sink:
sink_size = args.save_checkpoint_steps
model.train(args.epoch_size * dataset.get_dataset_size() // sink_size,
dataset,
callbacks=callbacks,
dataset_sink_mode=enable_sink,
sink_size=sink_size)
else:
model.train(args.epoch_size,
dataset,
callbacks=callbacks,
dataset_sink_mode=enable_sink)
def run_pretrain():
"""pre-train bert_clue"""
parser = argparse.ArgumentParser(description='bert pre_training')
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",
choices=["true", "false"],
help="Run distribute, default is false.")
parser.add_argument("--epoch_size",
type=int,
default="1",
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("--enable_save_ckpt",
type=str,
default="true",
choices=["true", "false"],
help="Enable save checkpoint, default is true.")
parser.add_argument("--enable_lossscale",
type=str,
default="true",
choices=["true", "false"],
help="Use lossscale or not, default is not.")
parser.add_argument("--do_shuffle",
type=str,
default="true",
choices=["true", "false"],
help="Enable shuffle for dataset, default is true.")
parser.add_argument("--enable_data_sink",
type=str,
default="true",
choices=["true", "false"],
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(
"--accumulation_steps",
type=int,
default="1",
help=
"Accumulating gradients N times before weight update, default is 1.")
parser.add_argument("--save_checkpoint_path",
type=str,
default="",
help="Save checkpoint path")
parser.add_argument("--load_checkpoint_path",
type=str,
default="",
help="Load checkpoint file path")
parser.add_argument("--save_checkpoint_steps",
type=int,
default=1000,
help="Save checkpoint steps, "
"default is 1000.")
parser.add_argument("--train_steps",
type=int,
default=-1,
help="Training Steps, default is -1, "
"meaning run all steps according to epoch number.")
parser.add_argument("--save_checkpoint_num",
type=int,
default=1,
help="Save checkpoint numbers, default is 1.")
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")
parser.add_argument("--enable_graph_kernel",
type=str,
default="auto",
choices=["auto", "true", "false"],
help="Accelerate by graph kernel, default is auto.")
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(reserve_class_name_in_scope=False)
ckpt_save_dir = args_opt.save_checkpoint_path
if args_opt.distribute == "true":
if args_opt.device_target == 'Ascend':
D.init()
device_num = args_opt.device_num
rank = args_opt.device_id % device_num
else:
D.init()
device_num = D.get_group_size()
rank = D.get_rank()
ckpt_save_dir = args_opt.save_checkpoint_path + 'ckpt_' + str(
get_rank()) + '/'
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
device_num=device_num)
if args_opt.device_target == 'Ascend':
_set_bert_all_reduce_split()
else:
rank = 0
device_num = 1
is_auto_enable_graph_kernel = _auto_enable_graph_kernel(
args_opt.device_target, args_opt.enable_graph_kernel)
if args_opt.enable_graph_kernel == "true" or is_auto_enable_graph_kernel:
context.set_context(enable_graph_kernel=True)
if args_opt.device_target == 'GPU' and bert_net_cfg.compute_type != mstype.float32 and \
not is_auto_enable_graph_kernel:
logger.warning('Gpu only support fp32 temporarily, run with fp32.')
bert_net_cfg.compute_type = mstype.float32
if args_opt.accumulation_steps > 1:
logger.info("accumulation steps: {}".format(
args_opt.accumulation_steps))
logger.info("global batch size: {}".format(
cfg.batch_size * args_opt.accumulation_steps))
if args_opt.enable_data_sink == "true":
args_opt.data_sink_steps *= args_opt.accumulation_steps
logger.info("data sink steps: {}".format(args_opt.data_sink_steps))
if args_opt.enable_save_ckpt == "true":
args_opt.save_checkpoint_steps *= args_opt.accumulation_steps
logger.info("save checkpoint steps: {}".format(
args_opt.save_checkpoint_steps))
ds = create_bert_dataset(device_num, rank, args_opt.do_shuffle,
args_opt.data_dir, args_opt.schema_dir)
net_with_loss = BertNetworkWithLoss(bert_net_cfg, True)
new_repeat_count = args_opt.epoch_size * ds.get_dataset_size(
) // args_opt.data_sink_steps
if args_opt.train_steps > 0:
train_steps = args_opt.train_steps * args_opt.accumulation_steps
new_repeat_count = min(new_repeat_count,
train_steps // args_opt.data_sink_steps)
else:
args_opt.train_steps = args_opt.epoch_size * ds.get_dataset_size(
) // args_opt.accumulation_steps
logger.info("train steps: {}".format(args_opt.train_steps))
optimizer = _get_optimizer(args_opt, net_with_loss)
callback = [
TimeMonitor(args_opt.data_sink_steps),
LossCallBack(ds.get_dataset_size())
]
if args_opt.enable_save_ckpt == "true" and args_opt.device_id % min(
8, device_num) == 0:
config_ck = CheckpointConfig(
save_checkpoint_steps=args_opt.save_checkpoint_steps,
keep_checkpoint_max=args_opt.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(
prefix='checkpoint_bert',
directory=None if ckpt_save_dir == "" else ckpt_save_dir,
config=config_ck)
callback.append(ckpoint_cb)
if args_opt.load_checkpoint_path:
param_dict = load_checkpoint(args_opt.load_checkpoint_path)
load_param_into_net(net_with_loss, param_dict)
if args_opt.enable_lossscale == "true":
update_cell = DynamicLossScaleUpdateCell(
loss_scale_value=cfg.loss_scale_value,
scale_factor=cfg.scale_factor,
scale_window=cfg.scale_window)
if args_opt.accumulation_steps <= 1:
net_with_grads = BertTrainOneStepWithLossScaleCell(
net_with_loss,
optimizer=optimizer,
scale_update_cell=update_cell)
else:
accumulation_steps = args_opt.accumulation_steps
net_with_grads = BertTrainAccumulateStepsWithLossScaleCell(
net_with_loss,
optimizer=optimizer,
scale_update_cell=update_cell,
accumulation_steps=accumulation_steps,
enable_global_norm=cfg.enable_global_norm)
else:
net_with_grads = BertTrainOneStepCell(net_with_loss,
optimizer=optimizer)
model = Model(net_with_grads)
model.train(new_repeat_count,
ds,
callbacks=callback,
dataset_sink_mode=(args_opt.enable_data_sink == "true"),
sink_size=args_opt.data_sink_steps)
def main():
args_opt = get_args()
rank = 0
device_num = 1
if args_opt.run_platform == "CPU":
context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
else:
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.run_platform, device_id=args_opt.device_id)
if args_opt.distribute:
device_num = args_opt.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True,
device_num=device_num)
init()
if config.model == "ssd_resnet50_fpn":
context.set_auto_parallel_context(all_reduce_fusion_config=[90, 183, 279])
else:
context.set_auto_parallel_context(all_reduce_fusion_config=[29, 58, 89])
rank = get_rank()
mindrecord_file = create_mindrecord(args_opt.dataset, "ssd.mindrecord", True)
if args_opt.only_create_dataset:
return
loss_scale = float(args_opt.loss_scale)
if args_opt.run_platform == "CPU":
loss_scale = 1.0
# When create MindDataset, using the fitst mindrecord file, such as ssd.mindrecord0.
use_multiprocessing = (args_opt.run_platform != "CPU")
dataset = create_ssd_dataset(mindrecord_file, repeat_num=1, batch_size=args_opt.batch_size,
device_num=device_num, rank=rank, use_multiprocessing=use_multiprocessing)
dataset_size = dataset.get_dataset_size()
print(f"Create dataset done! dataset size is {dataset_size}")
ssd = ssd_model_build(args_opt)
if ("use_float16" in config and config.use_float16) or args_opt.run_platform == "GPU":
ssd.to_float(dtype.float16)
net = SSDWithLossCell(ssd, config)
# checkpoint
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_size * args_opt.save_checkpoint_epochs)
save_ckpt_path = './ckpt_' + str(rank) + '/'
ckpoint_cb = ModelCheckpoint(prefix="ssd", directory=save_ckpt_path, config=ckpt_config)
if args_opt.pre_trained:
param_dict = load_checkpoint(args_opt.pre_trained)
if args_opt.filter_weight:
filter_checkpoint_parameter_by_list(param_dict, config.checkpoint_filter_list)
load_param_into_net(net, param_dict, True)
lr = Tensor(get_lr(global_step=args_opt.pre_trained_epoch_size * dataset_size,
lr_init=config.lr_init, lr_end=config.lr_end_rate * args_opt.lr, lr_max=args_opt.lr,
warmup_epochs=config.warmup_epochs,
total_epochs=args_opt.epoch_size,
steps_per_epoch=dataset_size))
if "use_global_norm" in config and config.use_global_norm:
opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr,
config.momentum, config.weight_decay, 1.0)
net = TrainingWrapper(net, opt, loss_scale, True)
else:
opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr,
config.momentum, config.weight_decay, loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
callback = [TimeMonitor(data_size=dataset_size), LossMonitor(), ckpoint_cb]
model = Model(net)
dataset_sink_mode = False
if args_opt.mode == "sink" and args_opt.run_platform != "CPU":
print("In sink mode, one epoch return a loss.")
dataset_sink_mode = True
print("Start train SSD, the first epoch will be slower because of the graph compilation.")
model.train(args_opt.epoch_size, dataset, callbacks=callback, dataset_sink_mode=dataset_sink_mode)
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", choices=["true", "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", choices=["true", "false"],
help="Enable shuffle for dataset, default is true.")
parser.add_argument("--enable_data_sink", type=str, default="true", choices=["true", "false"],
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")
parser.add_argument("--dataset_type", type=str, default="tfrecord",
help="dataset type tfrecord/mindrecord, default is tfrecord")
args_opt = parser.parse_args()
if args_opt.device_target == "Ascend":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
elif args_opt.device_target == "GPU":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target)
else:
raise Exception("Target error, GPU or Ascend is supported.")
context.set_context(reserve_class_name_in_scope=False)
save_ckpt_dir = os.path.join(args_opt.save_ckpt_path,
datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
if args_opt.distribute == "true":
if args_opt.device_target == 'Ascend':
D.init()
device_num = args_opt.device_num
rank = args_opt.device_id % device_num
else:
D.init()
device_num = D.get_group_size()
rank = D.get_rank()
save_ckpt_dir = save_ckpt_dir + '_ckpt_' + str(rank)
context.reset_auto_parallel_context()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True,
device_num=device_num)
else:
rank = 0
device_num = 1
if not os.path.exists(save_ckpt_dir):
os.makedirs(save_ckpt_dir)
enable_loss_scale = True
if args_opt.device_target == "GPU":
if bert_student_net_cfg.compute_type != mstype.float32:
logger.warning('Compute about the student only support float32 temporarily, run with float32.')
bert_student_net_cfg.compute_type = mstype.float32
# Backward of the network are calculated using fp32,
# and the loss scale is not necessary
enable_loss_scale = False
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)
if args_opt.dataset_type == "tfrecord":
dataset_type = DataType.TFRECORD
elif args_opt.dataset_type == "mindrecord":
dataset_type = DataType.MINDRECORD
else:
raise Exception("dataset format is not supported yet")
dataset = create_tinybert_dataset('gd', common_cfg.batch_size, device_num, rank,
args_opt.do_shuffle, args_opt.data_dir, args_opt.schema_dir,
data_type=dataset_type)
dataset_size = dataset.get_dataset_size()
print('dataset size: ', dataset_size)
print("dataset repeatcount: ", dataset.get_repeat_count())
if args_opt.enable_data_sink == "true":
repeat_count = args_opt.epoch_size * 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: not common_cfg.AdamWeightDecay.decay_filter(x), 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)]
if enable_loss_scale:
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)
else:
netwithgrads = BertTrainCell(netwithloss, optimizer=optimizer)
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)
def test_six_matmul_save():
class NetWithLoss(nn.Cell):
def __init__(self, network):
super(NetWithLoss, self).__init__()
self.loss = VirtualLoss()
self.network = network
def construct(self, x1, x6):
predict = self.network(x1, x6)
return self.loss(predict)
class GradWrap(nn.Cell):
def __init__(self, network):
super(GradWrap, self).__init__()
self.network = network
def construct(self, x1, x6):
return C.grad_all(self.network)(x1, x6)
class Net(nn.Cell):
def __init__(self, strategy1, strategy2, strategy3, strategy4,
strategy5, strategy6):
super().__init__()
self.matmul1 = P.MatMul().set_strategy(strategy1)
self.matmul2 = P.MatMul().set_strategy(strategy2)
self.matmul3 = P.MatMul().set_strategy(strategy3)
self.matmul4 = P.MatMul().set_strategy(strategy4)
self.matmul5 = P.MatMul().set_strategy(strategy5)
self.matmul6 = P.MatMul().set_strategy(strategy6)
self.weight1 = Parameter(Tensor(np.ones([32, 64]),
dtype=ms.float32),
name="weight1")
self.weight2 = Parameter(Tensor(np.ones([64, 64]),
dtype=ms.float32),
name="weight2")
self.weight3 = Parameter(Tensor(np.ones([64, 128]),
dtype=ms.float32),
name="weight3")
self.weight4 = Parameter(Tensor(np.ones([128, 64]),
dtype=ms.float32),
name="weight4")
self.weight5 = Parameter(Tensor(np.ones([64, 128]),
dtype=ms.float32),
name="weight5")
self.weight6 = Parameter(Tensor(np.ones([32, 128]),
dtype=ms.float32),
name="weight6")
def construct(self, x1, x6):
out = self.matmul1(x1, self.weight1)
out = self.matmul2(out, self.weight2)
out = self.matmul3(out, self.weight3)
out = self.matmul4(out, self.weight4)
out = self.matmul5(out, self.weight5)
out = out + self.weight6
out = self.matmul6(out, x6)
return out
reset_auto_parallel_context()
set_auto_parallel_context(device_num=8,
global_rank=0,
strategy_ckpt_save_file="./strategy_stage1.ckpt")
strategy1 = ((8, 1), (1, 1))
strategy2 = ((1, 8), (8, 1))
strategy3 = ((2, 2), (2, 2))
strategy4 = ((1, 1), (1, 8))
strategy5 = ((4, 2), (2, 1))
strategy6 = ((4, 1), (1, 2))
net = GradWrap(
NetWithLoss(
Net(strategy1, strategy2, strategy3, strategy4, strategy5,
strategy6)))
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
x1 = Tensor(np.ones([32, 32]), dtype=ms.float32)
x6 = Tensor(np.ones([128, 32]), dtype=ms.float32)
_executor.compile(net, x1, x6)
def main():
parser = argparse.ArgumentParser(description="YOLOv3 train")
parser.add_argument("--only_create_dataset",
type=bool,
default=False,
help="If set it true, only create "
"Mindrecord, default is false.")
parser.add_argument("--distribute",
type=bool,
default=False,
help="Run distribute, default is false.")
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("--lr",
type=float,
default=0.001,
help="Learning rate, default is 0.001.")
parser.add_argument("--mode",
type=str,
default="sink",
help="Run sink mode or not, default is sink")
parser.add_argument("--epoch_size",
type=int,
default=10,
help="Epoch size, default is 10")
parser.add_argument("--batch_size",
type=int,
default=32,
help="Batch size, default is 32.")
parser.add_argument("--pre_trained",
type=str,
default=None,
help="Pretrained checkpoint file path")
parser.add_argument("--pre_trained_epoch_size",
type=int,
default=0,
help="Pretrained epoch size")
parser.add_argument("--save_checkpoint_epochs",
type=int,
default=5,
help="Save checkpoint epochs, default is 5.")
parser.add_argument("--loss_scale",
type=int,
default=1024,
help="Loss scale, default is 1024.")
parser.add_argument(
"--mindrecord_dir",
type=str,
default="./Mindrecord_train",
help=
"Mindrecord directory. If the mindrecord_dir is empty, it wil generate mindrecord file by"
"image_dir and anno_path. Note if mindrecord_dir isn't empty, it will use mindrecord_dir "
"rather than image_dir and anno_path. Default is ./Mindrecord_train")
parser.add_argument("--image_dir",
type=str,
default="",
help="Dataset directory, "
"the absolute image path is joined by the image_dir "
"and the relative path in anno_path")
parser.add_argument("--anno_path",
type=str,
default="",
help="Annotation path.")
args_opt = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
device_id=args_opt.device_id)
if args_opt.distribute:
device_num = args_opt.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
mirror_mean=True,
device_num=device_num)
init()
rank = args_opt.device_id % device_num
else:
rank = 0
device_num = 1
print("Start create dataset!")
# It will generate mindrecord file in args_opt.mindrecord_dir,
# and the file name is yolo.mindrecord0, 1, ... file_num.
if not os.path.isdir(args_opt.mindrecord_dir):
os.makedirs(args_opt.mindrecord_dir)
prefix = "yolo.mindrecord"
mindrecord_file = os.path.join(args_opt.mindrecord_dir, prefix + "0")
if not os.path.exists(mindrecord_file):
if os.path.isdir(args_opt.image_dir) and os.path.exists(
args_opt.anno_path):
print("Create Mindrecord.")
data_to_mindrecord_byte_image(args_opt.image_dir,
args_opt.anno_path,
args_opt.mindrecord_dir,
prefix=prefix,
file_num=8)
print("Create Mindrecord Done, at {}".format(
args_opt.mindrecord_dir))
else:
print("image_dir or anno_path not exits.")
if not args_opt.only_create_dataset:
loss_scale = float(args_opt.loss_scale)
# When create MindDataset, using the fitst mindrecord file, such as yolo.mindrecord0.
dataset = create_yolo_dataset(mindrecord_file,
repeat_num=args_opt.epoch_size,
batch_size=args_opt.batch_size,
device_num=device_num,
rank=rank)
dataset_size = dataset.get_dataset_size()
print("Create dataset done!")
net = yolov3_resnet18(ConfigYOLOV3ResNet18())
net = YoloWithLossCell(net, ConfigYOLOV3ResNet18())
init_net_param(net, "XavierUniform")
# checkpoint
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_size *
args_opt.save_checkpoint_epochs)
ckpoint_cb = ModelCheckpoint(prefix="yolov3",
directory=None,
config=ckpt_config)
if args_opt.pre_trained:
if args_opt.pre_trained_epoch_size <= 0:
raise KeyError(
"pre_trained_epoch_size must be greater than 0.")
param_dict = load_checkpoint(args_opt.pre_trained)
load_param_into_net(net, param_dict)
total_epoch_size = 60
if args_opt.distribute:
total_epoch_size = 160
lr = Tensor(
get_lr(learning_rate=args_opt.lr,
start_step=args_opt.pre_trained_epoch_size * dataset_size,
global_step=total_epoch_size * dataset_size,
decay_step=1000,
decay_rate=0.95,
steps=True))
opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()),
lr,
loss_scale=loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
callback = [
TimeMonitor(data_size=dataset_size),
LossMonitor(), ckpoint_cb
]
model = Model(net)
dataset_sink_mode = False
if args_opt.mode == "sink":
print("In sink mode, one epoch return a loss.")
dataset_sink_mode = True
print(
"Start train YOLOv3, the first epoch will be slower because of the graph compilation."
)
model.train(args_opt.epoch_size,
dataset,
callbacks=callback,
dataset_sink_mode=dataset_sink_mode)
def run_transformer_train():
"""
Transformer training.
"""
parser = argparse_init()
args, _ = parser.parse_known_args()
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id)
context.set_context(save_graphs=True, reserve_class_name_in_scope=False, enable_auto_mixed_precision=False)
if args.distribute == "true":
device_num = args.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, mirror_mean=True,
parameter_broadcast=True, device_num=device_num)
D.init()
rank_id = args.device_id % device_num
else:
device_num = 1
rank_id = 0
dataset, repeat_count = create_transformer_dataset(epoch_count=args.epoch_size, rank_size=device_num,
rank_id=rank_id, do_shuffle=args.do_shuffle,
enable_data_sink=args.enable_data_sink,
dataset_path=args.data_path)
netwithloss = TransformerNetworkWithLoss(transformer_net_cfg, True)
if args.checkpoint_path:
parameter_dict = load_checkpoint(args.checkpoint_path)
else:
parameter_dict = {}
params = netwithloss.trainable_params()
for param in params:
name = param.name
value = param.default_input
if isinstance(value, Tensor):
if name.endswith(".gamma"):
parameter_dict[name] = Parameter(one_weight(value.asnumpy().shape), name=name)
elif name.endswith(".beta") or name.endswith(".bias"):
parameter_dict[name] = Parameter(zero_weight(value.asnumpy().shape), name=name)
elif "embedding" in name:
parameter_dict[name] = Parameter(normal_weight(value.asnumpy().shape,
transformer_net_cfg.hidden_size), name=name)
else:
parameter_dict[name] = Parameter(weight_variable(value.asnumpy().shape), name=name)
load_param_into_net(netwithloss, parameter_dict)
lr = Tensor(create_dynamic_lr(schedule="constant*rsqrt_hidden*linear_warmup*rsqrt_decay",
training_steps=dataset.get_dataset_size()*args.epoch_size,
learning_rate=cfg.lr_schedule.learning_rate,
warmup_steps=cfg.lr_schedule.warmup_steps,
hidden_size=transformer_net_cfg.hidden_size), mstype.float32)
optimizer = Adam(netwithloss.trainable_params(), lr)
callbacks = [TimeMonitor(dataset.get_dataset_size()), LossCallBack()]
if args.enable_save_ckpt == "true":
ckpt_config = CheckpointConfig(save_checkpoint_steps=args.save_checkpoint_steps,
keep_checkpoint_max=args.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(prefix='transformer', directory=args.save_checkpoint_path, config=ckpt_config)
callbacks.append(ckpoint_cb)
if args.enable_lossscale == "true":
scale_manager = DynamicLossScaleManager(init_loss_scale=cfg.init_loss_scale_value,
scale_factor=cfg.scale_factor,
scale_window=cfg.scale_window)
update_cell = scale_manager.get_update_cell()
netwithgrads = TransformerTrainOneStepWithLossScaleCell(netwithloss, optimizer=optimizer,
scale_update_cell=update_cell)
else:
netwithgrads = TransformerTrainOneStepCell(netwithloss, optimizer=optimizer)
netwithgrads.set_train(True)
model = Model(netwithgrads)
model.train(repeat_count, dataset, callbacks=callbacks, dataset_sink_mode=(args.enable_data_sink == "true"))
def teardown_module():
context.reset_auto_parallel_context()
_reset_op_id()
def get_args(phase):
"""Define the common options that are used in both training and test."""
parser = argparse.ArgumentParser(description='Cycle GAN.')
# basic parameters
parser.add_argument('--model', type=str, default="resnet", choices=("resnet", "unet"), \
help='generator model, should be in [resnet, unet].')
parser.add_argument('--platform', type=str, default="Ascend", choices=("Ascend", "GPU", "CPU"), \
help='run platform, only support GPU, CPU and Ascend')
parser.add_argument("--device_id",
type=int,
default=0,
help="device id, default is 0.")
parser.add_argument("--lr",
type=float,
default=0.0002,
help="learning rate, default is 0.0002.")
parser.add_argument('--pool_size', type=int, default=50, \
help='the size of image buffer that stores previously generated images, default is 50.')
parser.add_argument('--lr_policy', type=str, default='linear', choices=("linear", "constant"), \
help='learning rate policy, default is linear')
parser.add_argument("--image_size",
type=int,
default=256,
help="input image_size, default is 256.")
parser.add_argument('--batch_size',
type=int,
default=1,
help='batch_size, default is 1.')
parser.add_argument('--max_epoch',
type=int,
default=200,
help='epoch size for training, default is 200.')
parser.add_argument('--n_epochs', type=int, default=100, \
help='number of epochs with the initial learning rate, default is 100')
parser.add_argument("--beta1",
type=float,
default=0.5,
help="Adam beta1, default is 0.5.")
parser.add_argument('--init_type', type=str, default='normal', choices=("normal", "xavier"), \
help='network initialization, default is normal.')
parser.add_argument('--init_gain', type=float, default=0.02, \
help='scaling factor for normal, xavier and orthogonal, default is 0.02.')
# model parameters
parser.add_argument('--in_planes',
type=int,
default=3,
help='input channels, default is 3.')
parser.add_argument('--ngf',
type=int,
default=64,
help='generator model filter numbers, default is 64.')
parser.add_argument(
'--gl_num',
type=int,
default=9,
help='generator model residual block numbers, default is 9.')
parser.add_argument(
'--ndf',
type=int,
default=64,
help='discriminator model filter numbers, default is 64.')
parser.add_argument('--dl_num', type=int, default=3, \
help='discriminator model residual block numbers, default is 3.')
parser.add_argument('--slope',
type=float,
default=0.2,
help='leakyrelu slope, default is 0.2.')
parser.add_argument('--norm_mode', type=str, default="instance", choices=("batch", "instance"), \
help='norm mode, default is instance.')
parser.add_argument('--lambda_A', type=float, default=10.0, \
help='weight for cycle loss (A -> B -> A), default is 10.')
parser.add_argument('--lambda_B', type=float, default=10.0, \
help='weight for cycle loss (B -> A -> B), default is 10.')
parser.add_argument('--lambda_idt', type=float, default=0.5, \
help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the '
'weight of the identity mapping loss. For example, if the weight of the identity loss '
'should be 10 times smaller than the weight of the reconstruction loss,'
'please set lambda_identity = 0.1, default is 0.5.')
parser.add_argument('--gan_mode', type=str, default='lsgan', choices=("lsgan", "vanilla"), \
help='the type of GAN loss, default is lsgan.')
parser.add_argument('--pad_mode', type=str, default='REFLECT', choices=("CONSTANT", "REFLECT", "SYMMETRIC"), \
help='the type of Pad, default is REFLECT.')
parser.add_argument('--need_dropout', type=ast.literal_eval, default=True, \
help='whether need dropout, default is True.')
# distillation learning parameters
parser.add_argument('--kd', type=ast.literal_eval, default=False, \
help='knowledge distillation learning or not, default is False.')
parser.add_argument('--t_ngf', type=int, default=64, \
help='teacher network generator model filter numbers when `kd` is True, default is 64.')
parser.add_argument('--t_gl_num', type=int, default=9, \
help='teacher network generator model residual block numbers when `kd` is True, default is 9.')
parser.add_argument('--t_slope', type=float, default=0.2, \
help='teacher network leakyrelu slope when `kd` is True, default is 0.2.')
parser.add_argument('--t_norm_mode', type=str, default="instance", choices=("batch", "instance"), \
help='teacher network norm mode when `kd` is True, default is instance.')
parser.add_argument("--GT_A_ckpt", type=str, default=None, \
help="teacher network pretrained checkpoint file path of G_A when `kd` is True.")
parser.add_argument("--GT_B_ckpt", type=str, default=None, \
help="teacher network pretrained checkpoint file path of G_B when `kd` is True.")
# additional parameters
parser.add_argument('--device_num',
type=int,
default=1,
help='device num, default is 1.')
parser.add_argument("--G_A_ckpt",
type=str,
default=None,
help="pretrained checkpoint file path of G_A.")
parser.add_argument("--G_B_ckpt",
type=str,
default=None,
help="pretrained checkpoint file path of G_B.")
parser.add_argument("--D_A_ckpt",
type=str,
default=None,
help="pretrained checkpoint file path of D_A.")
parser.add_argument("--D_B_ckpt",
type=str,
default=None,
help="pretrained checkpoint file path of D_B.")
parser.add_argument("--save_checkpoint_epochs",
type=int,
default=1,
help="Save checkpoint epochs, default is 10.")
parser.add_argument("--print_iter",
type=int,
default=100,
help="log print iter, default is 100.")
parser.add_argument('--need_profiler', type=ast.literal_eval, default=False, \
help='whether need profiler, default is False.')
parser.add_argument('--save_graphs', type=ast.literal_eval, default=False, \
help='whether save graphs, default is False.')
parser.add_argument('--outputs_dir', type=str, default='./outputs', \
help='models are saved here, default is ./outputs.')
parser.add_argument('--dataroot', default=None, \
help='path of images (should have subfolders trainA, trainB, testA, testB, etc).')
parser.add_argument('--save_imgs', type=ast.literal_eval, default=True, \
help='whether save imgs when epoch end, if True result images will generate in '
'`outputs_dir/imgs`, default is True.')
if phase == "export":
parser.add_argument("--file_name",
type=str,
default="cyclegan",
help="output file name prefix.")
parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', \
help='file format')
args = parser.parse_args()
if args.device_num > 1 and args.platform != "CPU":
context.set_context(mode=context.GRAPH_MODE,
device_target=args.platform,
save_graphs=args.save_graphs)
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
device_num=args.device_num)
init()
args.rank = get_rank()
else:
context.set_context(mode=context.GRAPH_MODE,
device_target=args.platform,
save_graphs=args.save_graphs,
device_id=args.device_id)
args.rank = 0
args.device_num = 1
if args.platform != "GPU":
args.pad_mode = "CONSTANT"
if phase != "train" and (args.G_A_ckpt is None or args.G_B_ckpt is None):
raise ValueError('Must set G_A_ckpt and G_B_ckpt in predict phase!')
if args.kd:
if args.GT_A_ckpt is None or args.GT_B_ckpt is None:
raise ValueError(
'Must set GT_A_ckpt, GT_B_ckpt in knowledge distillation!')
if args.norm_mode == "instance" or (args.kd
and args.t_norm_mode == "instance"):
args.batch_size = 1
if args.dataroot is None and (phase in ["train", "predict"]):
raise ValueError('Must set dataroot!')
args.n_epochs_decay = args.max_epoch - args.n_epochs
args.phase = phase
return args
rank_size = int(os.environ.get("RANK_SIZE", 1))
set_seed(1)
if __name__ == '__main__':
data_config = DataConfig()
model_config = ModelConfig()
train_config = TrainConfig()
if rank_size > 1:
if args_opt.device_target == "Ascend":
device_id = int(os.getenv('DEVICE_ID'))
context.set_context(mode=context.GRAPH_MODE,
device_target=args_opt.device_target,
device_id=device_id)
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True)
init()
rank_id = int(os.environ.get('RANK_ID'))
elif args_opt.device_target == "GPU":
init()
context.set_context(mode=context.GRAPH_MODE,
device_target=args_opt.device_target)
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
device_num=get_group_size(),
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True)
rank_id = get_rank()
else:
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)