From 9e268c0e5e0393e4ce19a4161ef2a2f8eb56797d Mon Sep 17 00:00:00 2001 From: Zheyu Ye <37728728+ZheyuYe@users.noreply.github.com> Date: Sat, 8 Aug 2020 09:28:36 +0800 Subject: [PATCH] Fix electra (#1291) * update Dockerfile * fix num_out_files * fix run_electra * Revert "update Dockerfile" This reverts commit 80593a26c56b3e8aedb0d3c69d20224d8e6df743. --- scripts/pretraining/data_preprocessing.py | 4 +--- scripts/pretraining/run_electra.py | 8 +++++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/scripts/pretraining/data_preprocessing.py b/scripts/pretraining/data_preprocessing.py index 067fbf1634..1f75e2f782 100644 --- a/scripts/pretraining/data_preprocessing.py +++ b/scripts/pretraining/data_preprocessing.py @@ -53,9 +53,7 @@ def main(args): random.shuffle(fnames) num_files = len(fnames) num_out_files = min(args.num_out_files, num_files) - file_volume = math.ceil(num_files / num_out_files) - splited_files = np.array_split(fnames, file_volume) - num_out_files = len(splited_files) + splited_files = np.array_split(fnames, num_out_files) output_files = [os.path.join( args.output, "owt-pretrain-record-{}.npz".format(str(i).zfill(4))) for i in range(num_out_files)] print("All preprocessed features will be saved in {} npz files".format(num_out_files)) diff --git a/scripts/pretraining/run_electra.py b/scripts/pretraining/run_electra.py index 2a67017106..78653bc3ad 100644 --- a/scripts/pretraining/run_electra.py +++ b/scripts/pretraining/run_electra.py @@ -472,9 +472,11 @@ def train(args): train_end_time - train_start_time)) if writer is not None: writer.close() - model_name = args.model_name.replace('google', 'gluon') - save_dir = os.path.join(args.output_dir, model_name) - final_save(model, save_dir, tokenizer) + + if local_rank == 0: + model_name = args.model_name.replace('google', 'gluon') + save_dir = os.path.join(args.output_dir, model_name) + final_save(model, save_dir, tokenizer) # TODO(zheyuye), Directly implement a metric for weighted accuracy