MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / main

Function main

modelzoo/dbmtl/train.py:599–715  ·  view source on GitHub ↗
(tf_config=None, server=None)

Source from the content-addressed store, hash-verified

597 print("ACC = {}\nAUC = {}".format(eval_acc, eval_auc))
598
599def main(tf_config=None, server=None):
600 # check dataset and count data set size
601 print("Checking dataset...")
602 train_file = args.data_location + '/taobao_train_data'
603 test_file = args.data_location + '/taobao_test_data'
604 if args.parquet_dataset and not args.tf:
605 train_file += '.parquet'
606 test_file += '.parquet'
607 if (not os.path.exists(train_file)) or (not os.path.exists(test_file)):
608 print("Dataset does not exist in the given data_location.")
609 sys.exit()
610 no_of_training_examples = 0
611 no_of_test_examples = 0
612 if args.parquet_dataset and not args.tf:
613 import pyarrow.parquet as pq
614 no_of_training_examples = pq.read_table(train_file).num_rows
615 no_of_test_examples = pq.read_table(test_file).num_rows
616 else:
617 no_of_training_examples = sum(1 for line in open(train_file))
618 no_of_test_examples = sum(1 for line in open(test_file))
619 print("Number of training dataset is {}".format(no_of_training_examples))
620 print("Number of test dataset is {}".format(no_of_test_examples))
621
622 # set batch size, epoch & steps
623 batch_size = math.ceil(
624 args.batch_size / args.micro_batch
625 ) if args.micro_batch and not args.tf else args.batch_size
626
627 if args.steps == 0:
628 no_of_epochs = 100
629 train_steps = math.ceil(
630 (float(no_of_epochs) * no_of_training_examples) / batch_size)
631 else:
632 no_of_epochs = math.ceil(
633 (float(batch_size) * args.steps) / no_of_training_examples)
634 train_steps = args.steps
635 test_steps = math.ceil(float(no_of_test_examples) / batch_size)
636 print("The training steps is {}".format(train_steps))
637 print("The testing steps is {}".format(test_steps))
638
639 # set fixed random seed
640 tf.set_random_seed(args.seed)
641
642 # ste directory path for checkpoint_dir
643 model_dir = os.path.join(args.output_dir,
644 'model_DBMTL_' + str(int(time.time())))
645 checkpoint_dir = args.checkpoint if args.checkpoint else model_dir
646 print("Saving model checkpoints to = " + checkpoint_dir)
647
648 # create data pipeline of train & test dataset
649 train_dataset = build_model_input(train_file, batch_size, no_of_epochs)
650 test_dataset = build_model_input(test_file, batch_size, 1)
651
652 dataset_output_types = tf.data.get_output_types(train_dataset)
653 dataset_output_shapes = tf.data.get_output_shapes(test_dataset)
654 iterator = tf.data.Iterator.from_structure(dataset_output_types,
655 dataset_output_shapes)
656 next_element = iterator.get_next()

Callers 1

train.pyFile · 0.70

Calls 14

sumFunction · 0.85
DBMTLClass · 0.85
exitMethod · 0.80
timeMethod · 0.80
from_structureMethod · 0.80
build_model_inputFunction · 0.70
build_feature_colsFunction · 0.70
trainFunction · 0.70
evalFunction · 0.70
formatMethod · 0.45
joinMethod · 0.45
get_nextMethod · 0.45

Tested by

no test coverage detected