(tf_config=None, server=None)
| 337 | |
| 338 | |
| 339 | def main(tf_config=None, server=None): |
| 340 | # check dataset |
| 341 | print('Checking dataset') |
| 342 | train_file = args.data_location + '/train.csv' |
| 343 | test_file = args.data_location + '/eval.csv' |
| 344 | |
| 345 | if (not os.path.exists(train_file)) or (not os.path.exists(test_file)): |
| 346 | print( |
| 347 | '------------------------------------------------------------------------------------------' |
| 348 | ) |
| 349 | print( |
| 350 | "train.csv or eval.csv does not exist in the given data_location. Please provide valid path" |
| 351 | ) |
| 352 | print( |
| 353 | '------------------------------------------------------------------------------------------' |
| 354 | ) |
| 355 | sys.exit() |
| 356 | no_of_training_examples = sum(1 for line in open(train_file)) |
| 357 | no_of_test_examples = sum(1 for line in open(test_file)) |
| 358 | print("Numbers of training dataset is {}".format(no_of_training_examples)) |
| 359 | print("Numbers of test dataset is {}".format(no_of_test_examples)) |
| 360 | |
| 361 | # set params |
| 362 | # set batch size & steps |
| 363 | batch_size = args.batch_size |
| 364 | if args.steps == 0: |
| 365 | no_of_epochs = 10 |
| 366 | train_steps = math.ceil( |
| 367 | (float(no_of_epochs) * no_of_training_examples) / batch_size) |
| 368 | else: |
| 369 | no_of_epochs = math.ceil( |
| 370 | (float(batch_size) * args.steps) / no_of_training_examples) |
| 371 | train_steps = args.steps |
| 372 | test_steps = math.ceil(float(no_of_test_examples) / batch_size) |
| 373 | |
| 374 | # set fixed random seed |
| 375 | tf.set_random_seed(2021) |
| 376 | |
| 377 | # set directory path |
| 378 | model_dir = os.path.join(args.output_dir, |
| 379 | 'model_DeepFM_' + str(int(time.time()))) |
| 380 | checkpoint_dir = args.checkpoint if args.checkpoint else model_dir |
| 381 | print("Saving model checkpoints to " + checkpoint_dir) |
| 382 | |
| 383 | # create data pipline |
| 384 | wide_column, fm_column, deep_column = build_feature_cols() |
| 385 | train_dataset = generate_input_data(train_file, batch_size, no_of_epochs) |
| 386 | test_dataset = generate_input_data(test_file, batch_size, 1) |
| 387 | |
| 388 | iterator = tf.data.Iterator.from_structure(train_dataset.output_types, |
| 389 | test_dataset.output_shapes) |
| 390 | next_element = iterator.get_next() |
| 391 | |
| 392 | train_init_op = iterator.make_initializer(train_dataset) |
| 393 | test_init_op = iterator.make_initializer(test_dataset) |
| 394 | |
| 395 | # create variable partitioner for distributed training |
| 396 | num_ps_replicas = len(tf_config['ps_hosts']) if tf_config else 0 |
no test coverage detected