def train(config): # Build Train Data data = TrainData(config.graph_work_path) train_iter = BatchGraphGenerator(graph_wrappers=[1], batch_size=config.batch_size, data=data, samples=config.samples, num_workers=config.sample_workers, feed_name_list=None, use_pyreader=False, phase="train", graph_data_path=config.graph_work_path, shuffle=True, neg_type=config.neg_type) train_ds = Dataset.from_generator_func(train_iter).repeat(config.epochs) dev_ds = Dataset.from_generator_func(train_iter) ernie_cfg_dict, ernie_param_path = PretrainedModelLoader.from_pretrained( config.ernie_name) if "warm_start_from" not in config: warm_start_from = ernie_param_path else: ernie_param_path = config.ernie_param_path if "ernie_config" not in config: config.ernie_config = ernie_cfg_dict ws = propeller.WarmStartSetting(predicate_fn=lambda v: os.path.exists( os.path.join(warm_start_from, v.name)), from_dir=warm_start_from) train_ds.name = "train" train_ds.data_shapes = [[-1] + list(shape[1:]) for shape in train_ds.data_shapes] dev_ds.name = "dev" dev_ds.data_shapes = [[-1] + list(shape[1:]) for shape in dev_ds.data_shapes] tokenizer = load_tokenizer(config.ernie_name) config.cls_id = tokenizer.cls_id propeller.train.train_and_eval( model_class_or_model_fn=ERNIESageLinkPredictModel, params=config, run_config=config, train_dataset=train_ds, eval_dataset={"eval": dev_ds}, warm_start_setting=ws, )
.padded_batch(hparams.batch_size) \ train_ds.data_shapes = shapes train_ds.data_types = types dev_ds.data_shapes = shapes dev_ds.data_types = types test_ds.data_shapes = shapes test_ds.data_types = types varname_to_warmstart = re.compile( r'^encoder.*[wb]_0$|^.*embedding$|^.*bias$|^.*scale$|^pooled_fc.[wb]_0$' ) ws = propeller.WarmStartSetting( predicate_fn=lambda v: varname_to_warmstart.match(v.name) and os. path.exists(os.path.join(param_path, v.name)), from_dir=param_path, ) best_exporter = propeller.train.exporter.BestExporter( os.path.join(run_config.model_dir, 'best'), cmp_fn=lambda old, new: new['dev']['acc'] > old['dev']['acc']) propeller.train.train_and_eval(model_class_or_model_fn=model_fn, params=hparams, run_config=run_config, train_dataset=train_ds, eval_dataset={ 'dev': dev_ds, 'test': test_ds }, warm_start_setting=ws,
def train(args, pretrained_model_config=None): log.info("loading data") raw_dataset = GraphPropPredDataset(name=args.dataset_name) args.num_class = raw_dataset.num_tasks args.eval_metric = raw_dataset.eval_metric args.task_type = raw_dataset.task_type train_ds = MolDataset(args, raw_dataset) args.eval_steps = math.ceil(len(train_ds) / args.batch_size) log.info("Total %s steps (eval_steps) every epoch." % (args.eval_steps)) fn = MgfCollateFn(args) train_loader = Dataloader(train_ds, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=args.shuffle, stream_shuffle_size=args.shuffle_size, collate_fn=fn) # for evaluating eval_train_loader = train_loader eval_train_loader = PDataset.from_generator_func(eval_train_loader) train_loader = multi_epoch_dataloader(train_loader, args.epochs) train_loader = PDataset.from_generator_func(train_loader) if args.warm_start_from is not None: # warm start setting def _fn(v): if not isinstance(v, F.framework.Parameter): return False if os.path.exists(os.path.join(args.warm_start_from, v.name)): return True else: return False ws = propeller.WarmStartSetting(predicate_fn=_fn, from_dir=args.warm_start_from) else: ws = None def cmp_fn(old, new): if old['eval'][args.metrics] - new['eval'][args.metrics] > 0: log.info("best %s eval result: %s" % (args.metrics, new['eval'])) return True else: return False if args.log_id is not None: save_best_model = int(args.log_id) == 5 else: save_best_model = True best_exporter = propeller.exporter.BestResultExporter( args.output_dir, (cmp_fn, save_best_model)) eval_datasets = {"eval": eval_train_loader} propeller.train.train_and_eval( model_class_or_model_fn=MgfModel, params=pretrained_model_config, run_config=args, train_dataset=train_loader, eval_dataset=eval_datasets, warm_start_setting=ws, exporters=[best_exporter], )