Contents

module TensorStream
  # Parses pbtext files and loads it as a graph
  class Yaml < TensorStream::Serializer
    include TensorStream::StringHelper
    include TensorStream::OpHelper

    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
  end
end

Version data entries

9 entries across 9 versions & 1 rubygems

Version	Path
tensor_stream-1.0.9	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.8	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.7	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.6	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.5	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.4	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.3	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.2	lib/tensor_stream/graph_serializers/yaml.rb
tensor_stream-1.0.1	lib/tensor_stream/graph_serializers/yaml.rb