Set some ENV for these DeepRec's features enabled by ENV. More Detail information is shown in https://deeprec.readthedocs.io/zh/latest/index.html. START_STATISTIC_STEP & STOP_STATISTIC_STEP: On CPU platform, DeepRec supports memory optimization in both stand-alone and distribute
()
| 1003 | # This func is used to set ENV and enable these features. |
| 1004 | # A triple quotes comment is used to introduce these features and play an emphasizing role. |
| 1005 | def set_env_for_DeepRec(): |
| 1006 | ''' |
| 1007 | Set some ENV for these DeepRec's features enabled by ENV. |
| 1008 | More Detail information is shown in https://deeprec.readthedocs.io/zh/latest/index.html. |
| 1009 | START_STATISTIC_STEP & STOP_STATISTIC_STEP: On CPU platform, DeepRec supports memory optimization |
| 1010 | in both stand-alone and distributed trainging. It's default to open, and the |
| 1011 | default start and stop steps of collection is 1000 and 1100. Reduce the initial |
| 1012 | cold start time by the following settings. |
| 1013 | MALLOC_CONF: On CPU platform, DeepRec can use memory optimization with the jemalloc library. |
| 1014 | Please preload libjemalloc.so by `LD_PRELOAD=./libjemalloc.so.2 python ...` |
| 1015 | ''' |
| 1016 | os.environ['START_STATISTIC_STEP'] = '100' |
| 1017 | os.environ['STOP_STATISTIC_STEP'] = '110' |
| 1018 | os.environ['MALLOC_CONF']= \ |
| 1019 | 'background_thread:true,metadata_thp:auto,dirty_decay_ms:20000,muzzy_decay_ms:20000' |
| 1020 | if args.group_embedding == "collective": |
| 1021 | tf.config.experimental.enable_distributed_strategy(strategy="collective") |
| 1022 | if args.smartstaged and not args.tf: |
| 1023 | os.environ["TF_GPU_THREAD_COUNT"] = "16" |
| 1024 | |
| 1025 | if __name__ == '__main__': |
| 1026 | parser = get_arg_parser() |