(tf_config=None, server=None)
| 849 | |
| 850 | |
| 851 | def main(tf_config=None, server=None): |
| 852 | # check dataset and count data set size |
| 853 | print("Checking dataset...") |
| 854 | train_file = args.data_location + '/local_train_splitByUser' |
| 855 | test_file = args.data_location + '/local_test_splitByUser' |
| 856 | train_neg_file = train_file + '_neg' |
| 857 | test_neg_file = test_file + '_neg' |
| 858 | if args.parquet_dataset and not args.tf: |
| 859 | train_file += '.parquet' |
| 860 | train_neg_file += '.parquet' |
| 861 | test_file += '.parquet' |
| 862 | test_neg_file += '.parquet' |
| 863 | if (not os.path.exists(train_file)) or (not os.path.exists(test_file)) or ( |
| 864 | not os.path.exists(train_neg_file)) or ( |
| 865 | not os.path.exists(test_neg_file)): |
| 866 | print("Dataset does not exist in the given data_location.") |
| 867 | sys.exit() |
| 868 | no_of_training_examples = 0 |
| 869 | no_of_test_examples = 0 |
| 870 | if args.parquet_dataset and not args.tf: |
| 871 | import pyarrow.parquet as pq |
| 872 | no_of_training_examples = pq.read_table(train_file).num_rows |
| 873 | no_of_test_examples = pq.read_table(test_file).num_rows |
| 874 | else: |
| 875 | no_of_training_examples = sum(1 for line in open(train_file)) |
| 876 | no_of_test_examples = sum(1 for line in open(test_file)) |
| 877 | print("Numbers of training dataset is {}".format(no_of_training_examples)) |
| 878 | print("Numbers of test dataset is {}".format(no_of_test_examples)) |
| 879 | |
| 880 | # set batch size, eporch & steps |
| 881 | batch_size = math.ceil( |
| 882 | args.batch_size / args.micro_batch |
| 883 | ) if args.micro_batch and not args.tf else args.batch_size |
| 884 | |
| 885 | if args.steps == 0: |
| 886 | no_of_epochs = 1 |
| 887 | train_steps = math.ceil( |
| 888 | (float(no_of_epochs) * no_of_training_examples) / batch_size) |
| 889 | else: |
| 890 | no_of_epochs = math.ceil( |
| 891 | (float(batch_size) * args.steps) / no_of_training_examples) |
| 892 | train_steps = args.steps |
| 893 | test_steps = math.ceil(float(no_of_test_examples) / batch_size) |
| 894 | print("The training steps is {}".format(train_steps)) |
| 895 | print("The testing steps is {}".format(test_steps)) |
| 896 | |
| 897 | # set fixed random seed |
| 898 | tf.set_random_seed(args.seed) |
| 899 | |
| 900 | # set directory path for checkpoint_dir |
| 901 | model_dir = os.path.join(args.output_dir, |
| 902 | 'model_DIEN_' + str(int(time.time()))) |
| 903 | checkpoint_dir = args.checkpoint if args.checkpoint else model_dir |
| 904 | print("Saving model checkpoints to " + checkpoint_dir) |
| 905 | |
| 906 | # create data pipline of train & test dataset |
| 907 | train_dataset = build_model_input(train_file, train_neg_file, |
| 908 | batch_size, no_of_epochs) |
no test coverage detected