class TensorStream::Yaml

Parses pbtext files and loads it as a graph

Public Instance Methods

get_string(tensor_or_graph, session = nil, graph_keys = nil) { |graph, k| ... } click to toggle source
# File lib/tensor_stream/graph_serializers/yaml.rb, line 7
def get_string(tensor_or_graph, session = nil, graph_keys = nil)
  graph = tensor_or_graph.is_a?(Tensor) ? tensor_or_graph.graph : tensor_or_graph
  serialized_arr = []

  node_keys = graph_keys.nil? ? graph.node_keys : graph.node_keys.select { |k| graph_keys.include?(k) }

  node_keys.each do |k|
    node = if block_given?
      yield graph, k
    else
      graph.get_tensor_by_name(k)
    end
    next unless node.is_a?(Operation)

    serialized_arr << node.to_h
  end

  serialized_arr.to_yaml
end