浏览代码

[skip ci] adding gpu config settings

/release_3_distributed
Anupam Bhatnagar 4 年前
当前提交
beb7cb74
共有 2 个文件被更改,包括 3 次插入1 次删除
  1. 2
      ml-agents/mlagents/tf_utils/tf.py
  2. 2
      ml-agents/mlagents/trainers/stats.py

2
ml-agents/mlagents/tf_utils/tf.py


# Everywhere else is caught by the banned-modules setting for flake8
import tensorflow as tf # noqa I201
from distutils.version import LooseVersion
import horovod.tensorflow as hvd
# LooseVersion handles things "1.2.3a" or "4.5.6-rc7" fairly sensibly.

"""
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(hvd.local_rank())
# For multi-GPU training, set allow_soft_placement to True to allow
# placing the operation into an alternative device automatically
# to prevent from exceptions if the device doesn't suppport the operation

2
ml-agents/mlagents/trainers/stats.py


stats_summary = values["Environment/Cumulative Reward"]
rank = hvd.rank()
logger.info(
"Horovod Rank: {}."
"Horovod Rank: {}. "
"{}: Step: {}. "
"Time Elapsed: {:0.3f} s "
"Mean "

正在加载...
取消
保存