import os
import time
import multiprocessing
import numpy as np
import paddle.fluid as fluid
from reader.pretraining import ErnieDataReader
from model.ernie import ErnieModel, ErnieConfig
from optimization import optimization
from utils.args import print_arguments
from utils.init import init_checkpoint, init_pretraining_params
from pretrain_args import parser
args = parser.parse_args()
# yapf: enable.
def create_model(pyreader_name, ernie_config):
pyreader = fluid.layers.py_reader(
capacity=70,
shapes=[[-1, args.max_seq_len, 1], [-1, args.max_seq_len, 1],
[-1, args.max_seq_len, 1],
[-1, args.max_seq_len, args.max_seq_len], [-1, 1], [-1, 1],
[-1, 1]],
dtypes=[
'int64', 'int64', 'int64', 'float32', 'int64', 'int64', 'int64'
],
lod_levels=[0, 0, 0, 0, 0, 0, 0],
proc.wait()
if len(log_fns) > 0:
log_fns[i].close()
if proc.returncode != 0:
raise subprocess.CalledProcessError(
returncode=procs[i].returncode, cmd=cmds[i])
else:
log.info("proc %d finsh" % i)
except KeyboardInterrupt as e:
for p in procs:
log.info('killing %s' % p)
p.terminate()
def main(args):
if args.print_config:
print_arguments(args)
start_procs(args)
if __name__ == "__main__":
prepare_logger(log)
lanch_args = parser.parse_args()
pretraining_args = worker_parser.parse_args(
lanch_args.training_script_args)
init_path = pretraining_args.init_checkpoint
if init_path and not pretraining_args.use_fp16:
os.system('rename .master "" ' + init_path + '/*.master')
main(lanch_args)
