| 991 | |
| 992 | # Get parse |
| 993 | def get_arg_parser(): |
| 994 | parser = argparse.ArgumentParser() |
| 995 | parser.add_argument('--data_location', |
| 996 | help='Full path of train data', |
| 997 | required=False, |
| 998 | default='./data') |
| 999 | parser.add_argument('--steps', |
| 1000 | help='set the number of steps on train dataset', |
| 1001 | type=int, |
| 1002 | default=0) |
| 1003 | parser.add_argument('--batch_size', |
| 1004 | help='Batch size to train. Default is 512', |
| 1005 | type=int, |
| 1006 | default=2048) |
| 1007 | parser.add_argument('--output_dir', |
| 1008 | help='Full path to model output directory. \ |
| 1009 | Default to ./result. Covered by --checkpoint. ', |
| 1010 | required=False, |
| 1011 | default='./result') |
| 1012 | parser.add_argument('--checkpoint', |
| 1013 | help='Full path to checkpoints input/output. \ |
| 1014 | Default to ./result/$MODEL_TIMESTAMP', |
| 1015 | required=False) |
| 1016 | parser.add_argument('--save_steps', |
| 1017 | help='set the number of steps on saving checkpoints', |
| 1018 | type=int, |
| 1019 | default=0) |
| 1020 | parser.add_argument('--seed', |
| 1021 | help='set the random seed for tensorflow', |
| 1022 | type=int, |
| 1023 | default=2021) |
| 1024 | parser.add_argument('--optimizer', |
| 1025 | type=str, \ |
| 1026 | choices=['adam', 'adamasync', 'adagraddecay'], |
| 1027 | default='adamasync') |
| 1028 | parser.add_argument('--learning_rate', |
| 1029 | help='Learning rate for model', |
| 1030 | type=float, |
| 1031 | default=0.001) |
| 1032 | parser.add_argument('--keep_checkpoint_max', |
| 1033 | help='Maximum number of recent checkpoint to keep', |
| 1034 | type=int, |
| 1035 | default=1) |
| 1036 | parser.add_argument('--timeline', |
| 1037 | help='number of steps on saving timeline. Default 0', |
| 1038 | type=int, |
| 1039 | default=0) |
| 1040 | parser.add_argument('--protocol', |
| 1041 | type=str, |
| 1042 | choices=['grpc', 'grpc++', 'star_server'], |
| 1043 | default='grpc') |
| 1044 | parser.add_argument('--inter', |
| 1045 | help='set inter op parallelism threads.', |
| 1046 | type=int, |
| 1047 | default=0) |
| 1048 | parser.add_argument('--intra', |
| 1049 | help='set inter op parallelism threads.', |
| 1050 | type=int, |