kfac_cifar.py 文件源码

def sessrun(*args, **kwargs):
  sess = u.get_default_session()
  if not GLOBAL_PROFILE:
    return sess.run(*args, **kwargs)

  run_metadata = tf.RunMetadata()

  kwargs['options'] = full_trace_options
  kwargs['run_metadata'] = run_metadata
  result = sess.run(*args, **kwargs)
  first_entry = args[0]
  if isinstance(first_entry, list):
    if len(first_entry) == 0 and len(args) == 1:
      return None
    first_entry = first_entry[0]
  name = first_entry.name
  name = name.replace('/', '-')

  tl = timeline.Timeline(run_metadata.step_stats)
  ctf = tl.generate_chrome_trace_format()
  with open('timelines/%s.json'%(name,), 'w') as f:
    f.write(ctf)
  with open('timelines/%s.pbtxt'%(name,), 'w') as f:
    f.write(str(run_metadata))
  return result