(strategy)
| 563 | # save_steps = args.save_steps if args.save_steps or args.no_eval else steps |
| 564 | |
| 565 | def model_main(strategy): |
| 566 | # create model |
| 567 | model = DeepFM(wide_column=wide_column, |
| 568 | fm_column=fm_column, |
| 569 | deep_column=deep_column, |
| 570 | optimizer_type=args.optimizer, |
| 571 | learning_rate=args.learning_rate, |
| 572 | bf16=args.bf16, |
| 573 | stock_tf=args.tf, |
| 574 | adaptive_emb=args.adaptive_emb, |
| 575 | input_layer_partitioner=input_layer_partitioner, |
| 576 | dense_layer_partitioner=dense_layer_partitioner, |
| 577 | strategy=strategy) |
| 578 | |
| 579 | |
| 580 | run_config = tf.estimator.RunConfig(session_config=sess_config) |
| 581 | estimator = strategy.estimator( |
| 582 | model_fn=model._call, model_dir=checkpoint_dir, config=run_config) |
| 583 | |
| 584 | if args.mode == "evaluate": |
| 585 | estimator.evaluate(input_fn=lambda: build_model_input(test_file, batch_size, 1), |
| 586 | steps=test_steps, |
| 587 | hooks=hooks) |
| 588 | |
| 589 | elif args.mode == "predict": |
| 590 | pred_result = estimator.predict(input_fn=lambda: build_model_input(test_file, batch_size, 1), |
| 591 | predict_keys=['score'], |
| 592 | hooks=hooks, |
| 593 | yield_single_examples=False) |
| 594 | print(next(pred_result)) |
| 595 | |
| 596 | elif args.mode == "train": |
| 597 | stop_hook = tf.train.StopAtStepHook(last_step=train_steps) |
| 598 | hooks.append(stop_hook) |
| 599 | if args.timeline > 0: |
| 600 | hooks.append( |
| 601 | tf.train.ProfilerHook(save_steps=args.timeline, |
| 602 | output_dir=checkpoint_dir)) |
| 603 | estimator.train_and_evaluate( |
| 604 | tf.estimator.TrainSpec( |
| 605 | input_fn=lambda: build_model_input( |
| 606 | train_file, batch_size, no_of_epochs), |
| 607 | max_steps=train_steps, |
| 608 | hooks=hooks), |
| 609 | tf.estimator.EvalSpec( |
| 610 | input_fn=lambda: build_model_input( |
| 611 | test_file, batch_size, 1), |
| 612 | hooks=hooks)) |
| 613 | else: |
| 614 | estimator.export_saved_model( |
| 615 | checkpoint_dir, |
| 616 | build_receive_fn) |
| 617 | |
| 618 | if group_embedding_type in ["sok", "hb"]: |
| 619 | from tensorflow.python.distribute.group_embedding_collective_strategy import CollectiveStrategy |
no test coverage detected