34 lines
890 B
Python
34 lines
890 B
Python
import logging
|
|
import os
|
|
import sys
|
|
|
|
import dllogger
|
|
import horovod.mxnet as hvd
|
|
|
|
|
|
def format_step(step):
|
|
if isinstance(step, str):
|
|
return step
|
|
s = ""
|
|
if len(step) > 0:
|
|
s += "Epoch: {} ".format(step[0])
|
|
if len(step) > 1:
|
|
s += "Iteration: {} ".format(step[1])
|
|
if len(step) > 2:
|
|
s += "Validation Iteration: {} ".format(step[2])
|
|
if len(step) == 0:
|
|
s = "Summary:"
|
|
return s
|
|
|
|
|
|
def setup_logging(args):
|
|
logging.basicConfig(level=logging.DEBUG, format='{asctime}:{levelname}: {message}', style='{')
|
|
if hvd.rank() == 0:
|
|
dllogger.init(backends=[
|
|
dllogger.StdOutBackend(dllogger.Verbosity.DEFAULT, step_format=format_step),
|
|
dllogger.JSONStreamBackend(
|
|
dllogger.Verbosity.VERBOSE, os.path.join(args.workspace, args.dllogger_log)),
|
|
])
|
|
else:
|
|
dllogger.init([])
|