diff --git a/easy_rec/python/main.py b/easy_rec/python/main.py index c2158183b..373191633 100644 --- a/easy_rec/python/main.py +++ b/easy_rec/python/main.py @@ -105,12 +105,15 @@ def _create_estimator(pipeline_config, distribution=None, params={}): train_config = pipeline_config.train_config gpu_options = GPUOptions(allow_growth=True) # False) + logging.info('train_config.train_distribute=%s[value=%d]' % + (DistributionStrategy.Name(pipeline_config.train_config.train_distribute), + pipeline_config.train_config.train_distribute)) + + # set gpu options only under hvd scenes if hvd is not None and pipeline_config.train_config.train_distribute in [ DistributionStrategy.EmbeddingParallelStrategy, DistributionStrategy.SokStrategy, DistributionStrategy.HorovodStrategy ]: - logging.info('train_config.train_distribute=%s' % - pipeline_config.train_config.train_distribute) local_rnk = hvd.local_rank() gpus = tf.config.experimental.list_physical_devices('GPU') logging.info('local_rnk=%d num_gpus=%d' % (local_rnk, len(gpus)))