浏览代码

[skip ci] adding horovod rank

/distributed-training
Anupam Bhatnagar 4 年前
当前提交
56d82152
共有 1 个文件被更改,包括 4 次插入1 次删除
  1. 5
      ml-agents/mlagents/trainers/stats.py

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


from mlagents_envs.logging_util import get_logger
from mlagents_envs.timers import set_gauge
from mlagents.tf_utils import tf, generate_session_config
import horovod.tensorflow as hvd
logger = get_logger(__name__)

if "Environment/Cumulative Reward" in values:
stats_summary = values["Environment/Cumulative Reward"]
rank = hvd.rank()
"{}: Step: {}. "
"Horovod Rank: {}, {}: Step: {}. "
rank,
category,
step,
time.time() - self.training_start_time,

正在加载...
取消
保存