Method _write_cache

tensorflow/python/tpu/tensor_tracer.py:1199–1236

Writes the given caches as tensor summary. Args: step: Step tensor with dimension [num_cores]. **kwargs: The dictionary of tensors that needs to be written as summaries. Key and value pairs within kwargs correspond to the tag name, and tensor content that w

(step, **kwargs)

                       '--trace_dir=/model/dir')

    def _write_cache(step, **kwargs):
      """Writes the given caches as tensor summary.

      Args:
        step: Step tensor with dimension [num_cores].
        **kwargs: The dictionary of tensors that needs to be written as
          summaries. Key and value pairs within kwargs correspond to the tag
          name, and tensor content that will be written using summary.write.
          The trace_modes that use this function are:
            - summary: In summary mode, kwargs includes a single (tag, content)
            pair which are, _TT_SUMMARY_TAG and a tf.float32 signature_cache
            variable. The dimension of the signature_cache is:
              num_cores x num_traced_tensors x num_signatures.
            - full_tensor_summary: kwargs will include all traced tensors. Tag
            and content correspond to the name of the tensor, and its actual
            content.
      Returns:
        A tf.Operation that needs to be executed for the host call dependencies.
      """

      # TODO(deveci): Parametrize max_queue, so that flushing op can be called
      # less frequently.
      # Setting max_queue to 100 appears to be safe even when the number of
      # iterations are much lower, as the destructor of the writer will flushes
      # it.
      summary_write_ops = []
      with summary.create_file_writer_v2(
          self._parameters.trace_dir,
          filename_suffix=_TT_EVENT_FILE_SUFFIX,
          max_queue=_TT_SUMMARY_MAX_QUEUE).as_default():
        summary_metadata = summary_pb2.SummaryMetadata(
            plugin_data=summary_pb2.SummaryMetadata.PluginData(
                plugin_name=_TT_TENSORBOARD_PLUGIN_NAME))
        for key, value in kwargs.items():
          summary_write_ops.append(summary.write(
              _TT_SUMMARY_TAG + '/' + key, value, metadata=summary_metadata,
              step=step[0]))
      return control_flow_ops.group(summary_write_ops)

    step = array_ops.reshape(training_util.get_or_create_global_step(), [1])
    self._host_call_fn = {}