def predict_wrapper(args,
exe,
bert_config,
test_prog=None,
data_loader=None,
fetch_list=None):
# Context to do validation.
data_path = args.test_set_dir if args.do_test else args.validation_set_dir
data_reader = DataReader(data_path,
vocab_path=args.vocab_path,
batch_size=args.batch_size,
in_tokens=args.in_tokens,
voc_size=bert_config['vocab_size'],
shuffle_files=False,
epoch=1,
max_seq_len=args.max_seq_len,
is_test=True)
data_loader.set_batch_generator(data_reader.data_generator())
if args.do_test:
assert args.init_checkpoint is not None, "[FATAL] Please use --init_checkpoint '/path/to/checkpoints' \
to specify you pretrained model checkpoints"
init_pretraining_params(exe, args.init_checkpoint, test_prog)
def predict(exe=exe, data_loader=data_loader):
data_loader.start()
cost = 0
lm_cost = 0
acc = 0
steps = 0
time_begin = time.time()
while True:
try:
each_next_acc, each_mask_lm_cost, each_total_cost = exe.run(
fetch_list=fetch_list, program=test_prog)
acc += each_next_acc
lm_cost += each_mask_lm_cost
cost += each_total_cost
steps += 1
if args.do_test and steps % args.skip_steps == 0:
print("[test_set] steps: %d" % steps)
except fluid.core.EOFException:
data_loader.reset()
break
used_time = time.time() - time_begin
return cost, lm_cost, acc, steps, (args.skip_steps / used_time)
return predict
def train(args):
print("pretraining start")
bert_config = BertConfig(args.bert_config_path)
bert_config.print_config()
train_program = fluid.Program()
startup_prog = fluid.Program()
with fluid.program_guard(train_program, startup_prog):
with fluid.unique_name.guard():
train_data_loader, next_sent_acc, mask_lm_loss, total_loss = create_model(
bert_config=bert_config)
scheduled_lr, loss_scaling = optimization(
loss=total_loss,
warmup_steps=args.warmup_steps,
num_train_steps=args.num_train_steps,
learning_rate=args.learning_rate,
train_program=train_program,
startup_prog=startup_prog,
weight_decay=args.weight_decay,
scheduler=args.lr_scheduler,
use_fp16=args.use_fp16,
use_dynamic_loss_scaling=args.use_dynamic_loss_scaling,
init_loss_scaling=args.init_loss_scaling,
incr_every_n_steps=args.incr_every_n_steps,
decr_every_n_nan_or_inf=args.decr_every_n_nan_or_inf,
incr_ratio=args.incr_ratio,
decr_ratio=args.decr_ratio)
test_prog = fluid.Program()
with fluid.program_guard(test_prog, startup_prog):
with fluid.unique_name.guard():
test_data_loader, next_sent_acc, mask_lm_loss, total_loss = create_model(
bert_config=bert_config)
test_prog = test_prog.clone(for_test=True)
if args.use_cuda:
place = fluid.CUDAPlace(0)
dev_count = fluid.core.get_cuda_device_count()
else:
place = fluid.CPUPlace()
dev_count = int(os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
print("Device count %d" % dev_count)
nccl2_num_trainers = 1
nccl2_trainer_id = 0
print("args.is_distributed:", args.is_distributed)
if args.is_distributed:
worker_endpoints_env = os.getenv("worker_endpoints")
worker_endpoints = worker_endpoints_env.split(",")
trainers_num = len(worker_endpoints)
current_endpoint = os.getenv("current_endpoint")
trainer_id = worker_endpoints.index(current_endpoint)
if trainer_id == 0:
print("train_id == 0, sleep 60s")
time.sleep(60)
print("worker_endpoints:{} trainers_num:{} current_endpoint:{} \
trainer_id:{}".format(worker_endpoints, trainers_num,
current_endpoint, trainer_id))
# prepare nccl2 env.
config = fluid.DistributeTranspilerConfig()
config.mode = "nccl2"
t = fluid.DistributeTranspiler(config=config)
t.transpile(trainer_id,
trainers=worker_endpoints_env,
current_endpoint=current_endpoint,
program=train_program,
startup_program=startup_prog)
nccl2_num_trainers = trainers_num
nccl2_trainer_id = trainer_id
exe = fluid.Executor(place)
exe.run(startup_prog)
if args.init_checkpoint and args.init_checkpoint != "":
init_checkpoint(exe, args.init_checkpoint, train_program,
args.use_fp16)
data_reader = DataReader(data_dir=args.data_dir,
batch_size=args.batch_size,
in_tokens=args.in_tokens,
vocab_path=args.vocab_path,
voc_size=bert_config['vocab_size'],
epoch=args.epoch,
max_seq_len=args.max_seq_len,
generate_neg_sample=args.generate_neg_sample)
exec_strategy = fluid.ExecutionStrategy()
exec_strategy.use_experimental_executor = args.use_fast_executor
exec_strategy.num_threads = dev_count
exec_strategy.num_iteration_per_drop_scope = args.num_iteration_per_drop_scope
build_strategy = fluid.BuildStrategy()
if not sys.platform == "win32":
build_strategy.num_trainers = nccl2_num_trainers
elif nccl2_num_trainers > 1:
raise ValueError(
"Windows platform doesn't support distributed training!")
build_strategy.trainer_id = nccl2_trainer_id
# use_ngraph is for CPU only, please refer to README_ngraph.md for details
use_ngraph = os.getenv('FLAGS_use_ngraph')
if not use_ngraph:
train_compiled_program = fluid.CompiledProgram(
train_program).with_data_parallel(loss_name=total_loss.name,
exec_strategy=exec_strategy,
build_strategy=build_strategy)
if args.validation_set_dir and args.validation_set_dir != "":
predict = predict_wrapper(args,
exe,
bert_config,
test_prog=test_prog,
data_loader=test_data_loader,
fetch_list=[
next_sent_acc.name, mask_lm_loss.name,
total_loss.name
])
train_data_loader.set_batch_generator(data_reader.data_generator())
train_data_loader.start()
steps = 0
cost = []
lm_cost = []
acc = []
time_begin = time.time()
while steps < args.num_train_steps:
try:
steps += 1
skip_steps = args.skip_steps * nccl2_num_trainers
if nccl2_trainer_id != 0:
if use_ngraph:
exe.run(fetch_list=[], program=train_program)
else:
exe.run(fetch_list=[], program=train_compiled_program)
continue
if steps % args.skip_steps != 0:
if use_ngraph:
exe.run(fetch_list=[], program=train_program)
else:
exe.run(fetch_list=[], program=train_compiled_program)
else:
fetch_list = [
next_sent_acc.name, mask_lm_loss.name, total_loss.name,
scheduled_lr.name
]
if args.use_fp16:
fetch_list.append(loss_scaling.name)
if use_ngraph:
outputs = exe.run(fetch_list=fetch_list,
program=train_program)
else:
outputs = exe.run(fetch_list=fetch_list,
program=train_compiled_program)
if args.use_fp16:
each_next_acc, each_mask_lm_cost, each_total_cost, np_lr, np_scaling = outputs
else:
each_next_acc, each_mask_lm_cost, each_total_cost, np_lr = outputs
acc.extend(each_next_acc)
lm_cost.extend(each_mask_lm_cost)
cost.extend(each_total_cost)
time_end = time.time()
used_time = time_end - time_begin
epoch, current_file_index, total_file, current_file = data_reader.get_progress(
)
if args.verbose:
verbose = "feed_queue size: %d, " % train_data_loader.queue.size(
)
verbose += "current learning_rate: %f, " % np_lr[0]
if args.use_fp16:
verbose += "loss scaling: %f" % np_scaling[0]
print(verbose)
print(
"epoch: %d, progress: %d/%d, step: %d, loss: %f, "
"ppl: %f, next_sent_acc: %f, speed: %f steps/s, file: %s" %
(epoch, current_file_index, total_file, steps,
np.mean(np.array(cost)), np.mean(np.exp(
np.array(lm_cost))), np.mean(np.array(acc)),
skip_steps / used_time, current_file))
cost = []
lm_cost = []
acc = []
time_begin = time.time()
if steps % args.save_steps == 0:
save_path = os.path.join(args.checkpoints,
"step_" + str(steps))
fluid.save(program=train_program, model_path=save_path)
if args.validation_set_dir and steps % args.validation_steps == 0:
vali_cost, vali_lm_cost, vali_acc, vali_steps, vali_speed = predict(
)
print("[validation_set] epoch: %d, step: %d, "
"loss: %f, global ppl: %f, batch-averged ppl: %f, "
"next_sent_acc: %f, speed: %f steps/s" %
(epoch, steps, np.mean(np.array(vali_cost) / vali_steps),
np.exp(np.mean(np.array(vali_lm_cost) / vali_steps)),
np.mean(np.exp(np.array(vali_lm_cost) / vali_steps)),
np.mean(np.array(vali_acc) / vali_steps), vali_speed))
except fluid.core.EOFException:
train_data_loader.reset()
break
def train(args):
print("pretraining start")
bert_config = BertConfig(args.bert_config_path)
bert_config.print_config()
train_program = fluid.Program()
startup_prog = fluid.Program()
with fluid.program_guard(train_program, startup_prog):
with fluid.unique_name.guard():
train_pyreader, next_sent_acc, mask_lm_loss, total_loss = create_model(
pyreader_name='train_reader', bert_config=bert_config)
scheduled_lr = optimization(loss=total_loss,
warmup_steps=args.warmup_steps,
num_train_steps=args.num_train_steps,
learning_rate=args.learning_rate,
train_program=train_program,
startup_prog=startup_prog,
weight_decay=args.weight_decay,
scheduler=args.lr_scheduler,
use_fp16=args.use_fp16,
loss_scaling=args.loss_scaling)
fluid.memory_optimize(input_program=train_program,
skip_opt_set=[
next_sent_acc.name, mask_lm_loss.name,
total_loss.name
])
test_prog = fluid.Program()
with fluid.program_guard(test_prog, startup_prog):
with fluid.unique_name.guard():
test_pyreader, next_sent_acc, mask_lm_loss, total_loss = create_model(
pyreader_name='test_reader', bert_config=bert_config)
test_prog = test_prog.clone(for_test=True)
if args.use_cuda:
place = fluid.CUDAPlace(0)
dev_count = fluid.core.get_cuda_device_count()
else:
place = fluid.CPUPlace()
dev_count = int(os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
print("Device count %d" % dev_count)
if args.verbose:
if args.in_tokens:
lower_mem, upper_mem, unit = fluid.contrib.memory_usage(
program=train_program,
batch_size=args.batch_size // args.max_seq_len)
else:
lower_mem, upper_mem, unit = fluid.contrib.memory_usage(
program=train_program, batch_size=args.batch_size)
print("Theoretical memory usage in training: %.3f - %.3f %s" %
(lower_mem, upper_mem, unit))
nccl2_num_trainers = 1
nccl2_trainer_id = 0
print("args.is_distributed:", args.is_distributed)
if args.is_distributed:
worker_endpoints_env = os.getenv("worker_endpoints")
worker_endpoints = worker_endpoints_env.split(",")
trainers_num = len(worker_endpoints)
current_endpoint = os.getenv("current_endpoint")
trainer_id = worker_endpoints.index(current_endpoint)
if trainer_id == 0:
print("train_id == 0, sleep 60s")
time.sleep(60)
print("worker_endpoints:{} trainers_num:{} current_endpoint:{} \
trainer_id:{}".format(worker_endpoints, trainers_num,
current_endpoint, trainer_id))
# prepare nccl2 env.
config = fluid.DistributeTranspilerConfig()
config.mode = "nccl2"
t = fluid.DistributeTranspiler(config=config)
t.transpile(trainer_id,
trainers=worker_endpoints_env,
current_endpoint=current_endpoint,
program=train_program,
startup_program=startup_prog)
nccl2_num_trainers = trainers_num
nccl2_trainer_id = trainer_id
exe = fluid.Executor(place)
exe.run(startup_prog)
if args.init_checkpoint and args.init_checkpoint != "":
init_checkpoint(exe, args.init_checkpoint, train_program,
args.use_fp16)
data_reader = DataReader(data_dir=args.data_dir,
batch_size=args.batch_size,
in_tokens=args.in_tokens,
vocab_path=args.vocab_path,
voc_size=bert_config['vocab_size'],
epoch=args.epoch,
max_seq_len=args.max_seq_len,
generate_neg_sample=args.generate_neg_sample)
exec_strategy = fluid.ExecutionStrategy()
if args.use_fast_executor:
exec_strategy.use_experimental_executor = True
exec_strategy.num_threads = dev_count
build_strategy = fluid.BuildStrategy()
build_strategy.remove_unnecessary_lock = False
train_exe = fluid.ParallelExecutor(use_cuda=args.use_cuda,
loss_name=total_loss.name,
build_strategy=build_strategy,
exec_strategy=exec_strategy,
main_program=train_program,
num_trainers=nccl2_num_trainers,
trainer_id=nccl2_trainer_id)
if args.validation_set_dir and args.validation_set_dir != "":
predict = predict_wrapper(args,
exe,
bert_config,
test_prog=test_prog,
pyreader=test_pyreader,
fetch_list=[
next_sent_acc.name, mask_lm_loss.name,
total_loss.name
])
train_pyreader.decorate_tensor_provider(data_reader.data_generator())
train_pyreader.start()
steps = 0
cost = []
lm_cost = []
acc = []
time_begin = time.time()
while steps < args.num_train_steps:
try:
steps += nccl2_num_trainers
skip_steps = args.skip_steps * nccl2_num_trainers
if nccl2_trainer_id != 0:
train_exe.run(fetch_list=[])
continue
if steps % skip_steps != 0:
train_exe.run(fetch_list=[])
else:
each_next_acc, each_mask_lm_cost, each_total_cost, np_lr = train_exe.run(
fetch_list=[
next_sent_acc.name, mask_lm_loss.name, total_loss.name,
scheduled_lr.name
])
acc.extend(each_next_acc)
lm_cost.extend(each_mask_lm_cost)
cost.extend(each_total_cost)
print("feed_queue size", train_pyreader.queue.size())
time_end = time.time()
used_time = time_end - time_begin
epoch, current_file_index, total_file, current_file = data_reader.get_progress(
)
print("current learning_rate:%f" % np_lr[0])
print(
"epoch: %d, progress: %d/%d, step: %d, loss: %f, "
"ppl: %f, next_sent_acc: %f, speed: %f steps/s, file: %s" %
(epoch, current_file_index, total_file, steps,
np.mean(np.array(cost)), np.mean(np.exp(
np.array(lm_cost))), np.mean(np.array(acc)),
skip_steps / used_time, current_file))
cost = []
lm_cost = []
acc = []
time_begin = time.time()
if steps % args.save_steps == 0:
save_path = os.path.join(args.checkpoints,
"step_" + str(steps))
fluid.io.save_persistables(exe, save_path, train_program)
if args.validation_set_dir and steps % args.validation_steps == 0:
vali_cost, vali_lm_cost, vali_acc, vali_steps, vali_speed = predict(
)
print("[validation_set] epoch: %d, step: %d, "
"loss: %f, global ppl: %f, batch-averged ppl: %f, "
"next_sent_acc: %f, speed: %f steps/s" %
(epoch, steps, np.mean(np.array(vali_cost) / vali_steps),
np.exp(np.mean(np.array(vali_lm_cost) / vali_steps)),
np.mean(np.exp(np.array(vali_lm_cost) / vali_steps)),
np.mean(np.array(vali_acc) / vali_steps), vali_speed))
except fluid.core.EOFException:
train_pyreader.reset()
break