| 313 | |
| 314 | |
| 315 | def get_arg_parser(): |
| 316 | parser = argparse.ArgumentParser() |
| 317 | parser.add_argument('--job', |
| 318 | help='train or test model', |
| 319 | type=str, |
| 320 | choices=["train", "test"], |
| 321 | default='train') |
| 322 | parser.add_argument('--seed', help='set random seed', type=int, default=3) |
| 323 | parser.add_argument('--data_location', |
| 324 | help='Full path of train data', |
| 325 | required=False, |
| 326 | default='./data') |
| 327 | parser.add_argument('--output_dir', |
| 328 | help='Full path to logs & model output directory', |
| 329 | required=False, |
| 330 | default='./result') |
| 331 | parser.add_argument('--model_dir', |
| 332 | help='Full path to test model directory', |
| 333 | required=False) |
| 334 | parser.add_argument('--save_steps', |
| 335 | help='set the number of steps on saving checkpoints', |
| 336 | type=int, |
| 337 | default=0) |
| 338 | parser.add_argument('--test_steps', |
| 339 | help='set the number of steps on test training model', |
| 340 | type=int, |
| 341 | default=100) |
| 342 | parser.add_argument('--steps', |
| 343 | help='set the number of steps on train dataset', |
| 344 | type=int, |
| 345 | default=0) |
| 346 | parser.add_argument('--timeline', |
| 347 | help='number of steps on saving timeline. Default 0', |
| 348 | type=int, |
| 349 | default=0) |
| 350 | parser.add_argument('--bf16', |
| 351 | help='enable DeepRec BF16 in deep model. Default FP32', |
| 352 | action='store_true') |
| 353 | parser.add_argument('--no_eval', |
| 354 | help='not evaluate trained model by eval dataset.', |
| 355 | action='store_true') |
| 356 | return parser |
| 357 | |
| 358 | |
| 359 | if __name__ == '__main__': |