Python v1.RunOptions方法代码示例

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import RunOptions [as 别名]
def save_checkpoint_chrome_trace(dataset: str, model_name: str, log_base: PathLike, batch_size: int = 32):
    def trace_solver_solution(save_path: PathLike, train_ds, solver):
        import tensorflow.compat.v1 as tf1
        from tensorflow.python.client import timeline

        data_iter = train_ds.__iter__()
        data_list = [x.numpy() for x in data_iter.next()]
        with tf1.Session() as sess:
            sqrtn_fn, *_ = _build_model_via_solver(dataset, model_name, train_ds.element_spec, solver)
            out = sqrtn_fn(*[tf1.convert_to_tensor(x) for x in data_list])

            run_meta = tf1.RunMetadata()
            sess.run(tf1.global_variables_initializer())
            sess.run(out, options=tf1.RunOptions(trace_level=tf1.RunOptions.FULL_TRACE), run_metadata=run_meta)
            t1 = timeline.Timeline(run_meta.step_stats)
            lctf = t1.generate_chrome_trace_format()

        with Path(save_path).open("w") as f:
            f.write(lctf)

    log_base = Path(log_base)
    log_base.mkdir(parents=True, exist_ok=True)
    train_ds, test_ds = get_data(dataset, batch_size=batch_size)
    trace_solver_solution(log_base / "check_all.json", train_ds, solve_checkpoint_all)
    trace_solver_solution(log_base / "check_sqrtn_noap.json", train_ds, solve_chen_sqrtn_noap) 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import RunOptions [as 别名]
def fit(self, feed, session=None):
    """Training step for observed source language example."""
    if session is None:
      session = tf.get_default_session()
    run_options = tf.RunOptions(
        report_tensor_allocations_upon_oom=True)
    _, cost, summary, last_state = session.run(
        [self.training_update, self.unregularized_loss, self.training_summary,
         self.last_state],
        feed_dict=feed, options=run_options)
    return cost, summary, last_state 

# 需要导入模块: from tensorflow.compat import v1 [as 别名]
# 或者: from tensorflow.compat.v1 import RunOptions [as 别名]
def benchmark(self, image_arrays, trace_filename=None):
    """Benchmark inference latency/throughput.

    Args:
      image_arrays: a list of images in numpy array format.
      trace_filename: If None, specify the filename for saving trace.
    """
    if not self.sess:
      self.build()

    # init session
    self.sess.run(
        self.signitures['prediction'],
        feed_dict={self.signitures['image_arrays']: image_arrays})

    start = time.perf_counter()
    for _ in range(10):
      self.sess.run(
          self.signitures['prediction'],
          feed_dict={self.signitures['image_arrays']: image_arrays})
    end = time.perf_counter()
    inference_time = (end - start) / 10

    print('Per batch inference time: ', inference_time)
    print('FPS: ', self.batch_size / inference_time)

    if trace_filename:
      run_options = tf.RunOptions()
      run_options.trace_level = tf.RunOptions.FULL_TRACE
      run_metadata = tf.RunMetadata()
      self.sess.run(
          self.signitures['prediction'],
          feed_dict={self.signitures['image_arrays']: image_arrays},
          options=run_options,
          run_metadata=run_metadata)
      with tf.io.gfile.GFile(trace_filename, 'w') as trace_file:
        trace = timeline.Timeline(step_stats=run_metadata.step_stats)
        trace_file.write(trace.generate_chrome_trace_format(show_memory=True))