tensor_stream 0.6.1 → 0.7.0

data/lib/tensor_stream/exceptions.rb

@@ -0,0 +1,6 @@
+module TensorStream
+  class TensorStreamError < RuntimeError; end
+  class KeyError < TensorStreamError; end
+  class ValueError < TensorStreamError; end
+  class InvalidArgumentError < TensorStreamError; end
+end

data/lib/tensor_stream/graph.rb

@@ -1,13 +1,15 @@
 module TensorStream
   # A class that defines a TensorStream graph
   class Graph
-    attr_accessor :nodes, :collections, :eager_execution, :random_seed, :constants
+    attr_accessor :nodes, :node_keys, :collections, :eager_execution, :random_seed, :constants
 
     def initialize
       @eager_execution = false
       @nodes = {}
+      @node_keys = []
       @collections = {
-        :"#{GraphKeys::GLOBAL_VARIABLES}" => []
+        :"#{GraphKeys::GLOBAL_VARIABLES}" => [],
+        :"#{GraphKeys::TRAINABLE_VARIABLES}" => []
       }
       @constants = {}
     end
@@ -19,8 +21,10 @@ module TensorStream
       @op_counter = 0
       @random_seed = nil
       @nodes = {}
+      @node_keys = []
       @collections = {
-        :"#{GraphKeys::GLOBAL_VARIABLES}" => []
+        :"#{GraphKeys::GLOBAL_VARIABLES}" => [],
+        :"#{GraphKeys::TRAINABLE_VARIABLES}" => []
       }
       @constants = {}
     end
@@ -83,6 +87,7 @@ module TensorStream
                   end
 
       node.device = get_device_scope
+      @node_keys << node.name
       @nodes[node.name] = node
       @constants[node.name] = node if node.is_const
       node.send(:propagate_outputs)
@@ -98,6 +103,11 @@ module TensorStream
       @nodes[name]
     end
 
+    def get_tensor_by_name(name)
+      raise TensorStream::KeyError, "#{name} not found" unless @nodes.key?(name)
+      get_node(name)
+    end
+
     def add_node!(name, node)
       @nodes[name] = node
       node
@@ -120,12 +130,12 @@ module TensorStream
       add_node(node)
     end
 
-    def control_dependencies(control_inputs = [], &block)
+    def control_dependencies(control_inputs = [])
       Thread.current["ts_graph_#{object_id}"] ||= {}
       Thread.current["ts_graph_#{object_id}"][:control_dependencies] ||= []
       Thread.current["ts_graph_#{object_id}"][:control_dependencies] << Operation.new(:no_op, *control_inputs)
       begin
-        block.call
+        yield
       ensure
         Thread.current["ts_graph_#{object_id}"][:control_dependencies].pop
       end
@@ -201,6 +211,11 @@ module TensorStream
       TensorStream::Pbtext.new.get_string(self)
     end
 
+    def self.parse_from_string(buffer)
+      builder = TensorStream::GraphBuilder.new(Graph.new)
+      builder.build(buffer)
+    end
+
     def graph_def_versions
       "producer: 26"
     end
@@ -208,7 +223,7 @@ module TensorStream
     protected
 
     def _variable_scope
-      return OpenStruct.new(name: '', reuse: false, initializer: nil) if Thread.current[:tensor_stream_variable_scope].nil? || Thread.current[:tensor_stream_variable_scope].empty?
+      return VariableScope.new(name: '', reuse: false, initializer: nil) if Thread.current[:tensor_stream_variable_scope].nil? || Thread.current[:tensor_stream_variable_scope].empty?
       scope = Thread.current[:tensor_stream_variable_scope].last
       scope
     end

data/lib/tensor_stream/graph_builder.rb

@@ -0,0 +1,67 @@
+module TensorStream
+  class GraphBuilder
+    include TensorStream::OpHelper
+    include TensorStream::StringHelper
+
+    attr_accessor :graph
+
+    def initialize(graph)
+      @graph = graph
+    end
+
+    def build(buffer)
+      protobuf = TensorStream::Protobuf.new
+      parsed_tree = protobuf.load_from_string(buffer)
+      parsed_tree.each do |node|
+        next unless node['type'] == 'node'
+        # puts "build #{node['name']}"
+        options = protobuf.options_evaluator(node)
+        options[:name] = node['name']
+        options[:__graph] = @graph
+        value = options.delete('value')
+        options = symbolize_keys(options)
+        case node['op']
+        when 'Const'
+          dimension = shape_eval(value)
+          rank = dimension.size
+          options[:value] = value
+          options[:const] = true
+          TensorStream::Tensor.new(options[:dtype] || options[:T], rank, dimension, options)
+        when 'VariableV2'
+          # evaluate options
+          shape = options[:shape]
+          TensorStream::Variable.new(options[:dtype] || options[:T], nil, shape, nil, options)
+        when 'Placeholder'
+          shape = options[:shape]
+          TensorStream::Placeholder.new(options[:dtype] || options[:T], nil, shape, options)
+        else
+          op = underscore(node['op']).to_sym
+          unless TensorStream::Evaluator::RubyEvaluator.ops.keys.include?(op)
+            puts "warning unsupported op #{op}"
+          end
+          # map input tensor
+          inputs = node['input'].map do |input|
+            input[0] = '' if input.start_with?('^')
+
+            input_indexed, index = input.split(':')
+
+            tensor = if index && index.to_i > 0
+                       @graph.get_tensor_by_name(input_indexed)[index.to_i]
+                     else
+                       @graph.get_tensor_by_name(input)
+                     end
+
+            raise "tensor not found by name #{input}" if tensor.nil?
+
+            tensor
+          end
+
+          options[:data_type] = options.delete(:T)
+          TensorStream::Operation.new(op, *inputs, options)
+        end
+      end
+
+      @graph
+    end
+  end
+end

data/lib/tensor_stream/graph_deserializers/protobuf.rb

@@ -0,0 +1,271 @@
+require 'yaml'
+
+module TensorStream
+  # A .pb graph deserializer
+  class Protobuf
+    def initialize
+    end
+
+    def load_from_string(buffer)
+      evaluate_lines(buffer.split("\n").map(&:strip))
+    end
+
+    ##
+    # parsers a protobuf file and spits out
+    # a ruby hash
+    def load(pbfile)
+      f = File.new(pbfile, 'r')
+      lines = []
+      while !f.eof? && (str = f.readline.strip)
+        lines << str
+      end
+      evaluate_lines(lines)
+    end
+
+    def parse_value(value_node)
+      if value_node['tensor']
+        evaluate_tensor_node(value_node['tensor'])
+      end
+    end
+
+    def evaluate_tensor_node(node)
+      if !node['shape'].empty? && node['tensor_content']
+        content = node['tensor_content']
+        unpacked = eval(%Q{"#{content}"})
+
+        if node['dtype'] == 'DT_FLOAT'
+          TensorShape.reshape(unpacked.unpack('f*'), node['shape'])
+        elsif node['dtype'] == 'DT_INT32'
+          TensorShape.reshape(unpacked.unpack('l*'), node['shape'])
+        elsif node['dtype'] == 'DT_STRING'
+          node['string_val']
+        else
+          raise "unknown dtype #{node['dtype']}"
+        end
+      else
+
+        val = if node['dtype'] == 'DT_FLOAT'
+          node['float_val'] ? node['float_val'].to_f : []
+        elsif node['dtype'] == 'DT_INT32'
+          node['int_val'] ? node['int_val'].to_i : []
+        elsif node['dtype'] == 'DT_STRING'
+          node['string_val']
+        else
+          raise "unknown dtype #{node['dtype']}"
+        end
+
+        if node['shape'] == [1]
+          [val]
+        else
+          val
+        end
+      end
+    end
+
+    def map_type_to_ts(attr_value)
+      case(attr_value)
+      when 'DT_FLOAT'
+        :float32
+      when 'DT_INT32'
+        :int32
+      when 'DT_INT64'
+        :int64
+      when 'DT_STRING'
+        :string
+      when 'DT_BOOL'
+        :boolean
+      else
+        raise "unknown type #{attr_value}"
+      end
+    end
+
+    def options_evaluator(node)
+      return {} if node['attributes'].nil?
+
+      node['attributes'].map do |attribute|
+        attr_type, attr_value = attribute['value'].collect { |k, v| [k, v] }.flatten(1)
+
+        if attr_type == 'tensor'
+          attr_value = evaluate_tensor_node(attr_value)
+        elsif attr_type == 'type'
+          attr_value = map_type_to_ts(attr_value)
+        elsif attr_type == 'b'
+          attr_value = attr_value == 'true'
+        end
+
+        [attribute['key'], attr_value]
+      end.to_h
+    end
+
+    protected
+
+    def evaluate_lines(lines = [])
+      block = []
+      node = {}
+      node_attr = {}
+      dim = []
+      state = :top
+
+      lines.each do |str|
+        case(state)
+        when :top
+          node['type'] = parse_node_name(str)
+          state = :node_context
+          next
+        when :node_context
+          if str == 'attr {'
+            state = :attr_context
+            node_attr = {}
+            node['attributes'] ||= []
+            node['attributes'] << node_attr
+            next
+          elsif str == '}'
+            state = :top
+            block << node
+            node = {}
+            next
+          else
+            key, value = str.split(':', 2)
+            if key == 'input'
+              node['input'] ||= []
+              node['input'] << process_value(value.strip)
+            else
+              node[key] = process_value(value.strip)
+            end
+          end
+        when :attr_context
+          if str == 'value {'
+            state = :value_context
+            node_attr['value'] = {}
+            next
+          elsif str == '}'
+            state = :node_context
+            next
+          else
+            key, value = str.split(':', 2)
+            node_attr[key] = process_value(value.strip)
+          end
+        when :value_context
+          if str == 'list {'
+            state = :list_context
+            node_attr['value'] = []
+            next
+          elsif str == 'shape {'
+            state = :shape_context
+            node_attr['value']['shape'] = []
+            next
+          elsif str == 'tensor {'
+            state = :tensor_context
+            node_attr['value']['tensor'] = {}
+            next
+          elsif str == '}'
+            state = :attr_context
+            next
+          else
+            key, value = str.split(':', 2)
+            if key == 'dtype'
+              node_attr['value']['dtype'] = value.strip
+            elsif key == 'type'
+              node_attr['value']['type'] = value.strip
+            else
+              node_attr['value'][key] = process_value(value.strip)
+            end
+          end
+        when :list_context
+          if str == '}'
+            state = :value_context
+            next
+          else
+            key, value = str.split(':', 2)
+            node_attr['value'] << { key => value}
+          end
+        when :tensor_context
+          if str == 'tensor_shape {'
+            state = :tensor_shape_context
+            node_attr['value']['tensor']['shape'] = []
+            next
+          elsif str == '}'
+            state = :value_context
+            next
+          else
+            key, value = str.split(':', 2)
+            if node_attr['value']['tensor'][key] && !node_attr['value']['tensor'][key].is_a?(Array)
+              node_attr['value']['tensor'][key] = [node_attr['value']['tensor'][key]]
+              node_attr['value']['tensor'][key] << process_value(value.strip)
+            elsif node_attr['value']['tensor'][key]
+              node_attr['value']['tensor'][key] << process_value(value.strip)
+            else
+              node_attr['value']['tensor'][key] = process_value(value.strip)
+            end
+          end
+        when :tensor_shape_context
+          if str == 'dim {'
+            state = :tensor_shape_dim_context
+            next
+          elsif str == '}'
+            state = :tensor_context
+            next
+          end
+        when :shape_context
+          if str == '}'
+            state = :value_context
+            next
+          elsif str == 'dim {'
+            state = :shape_dim_context
+            next
+          end
+        when :shape_dim_context
+          if str == '}'
+            state = :shape_context
+            next
+          else
+            key, value = str.split(':', 2)
+            node_attr['value']['shape'] << value.strip.to_i
+          end
+        when :tensor_shape_dim_context
+          if str == '}'
+            state = :tensor_shape_context
+            next
+          else
+            key, value = str.split(':', 2)
+            node_attr['value']['tensor']['shape'] << value.strip.to_i
+          end
+        end
+      end
+
+      block
+    end
+
+    def parse_node_name(str)
+      name = str.split(' ')[0]
+    end
+
+    def process_value(value)
+      if value.start_with?('"')
+        unescape(value.gsub!(/\A"|"\Z/, ''))
+      else
+        unescape(value)
+      end
+    end
+
+    UNESCAPES = {
+      'a' => "\x07", 'b' => "\x08", 't' => "\x09",
+      'n' => "\x0a", 'v' => "\x0b", 'f' => "\x0c",
+      'r' => "\x0d", 'e' => "\x1b", "\\\\" => "\x5c",
+      "\"" => "\x22", "'" => "\x27"
+    }
+
+    def unescape(str)
+      # Escape all the things
+      str.gsub(/\\(?:([#{UNESCAPES.keys.join}])|u([\da-fA-F]{4}))|\\0?x([\da-fA-F]{2})/) {
+        if $1
+          $1 == '\\' ? '\\' : UNESCAPES[$1]
+        elsif $2 # escape \u0000 unicode
+          ["#$2".hex].pack('U*')
+        elsif $3 # escape \0xff or \xff
+          [$3].pack('H2')
+        end
+      }
+    end
+  end
+end

data/lib/tensor_stream/graph_keys.rb

@@ -2,5 +2,6 @@ module TensorStream
   class GraphKeys
     GLOBAL_VARIABLES = 'variables'.freeze
     TRAINABLE_VARIABLES = 'trainable_variables'.freeze
+    GLOBAL_STEP = 'global_step'.freeze
   end
 end

data/lib/tensor_stream/graph_serializers/pbtext.rb

@@ -6,7 +6,8 @@ module TensorStream
     def get_string(tensor_or_graph, session = nil)
       graph = tensor_or_graph.is_a?(Tensor) ? tensor_or_graph.graph : tensor_or_graph
       @lines = []
-      graph.nodes.each do |k, node|
+      graph.node_keys.each do |k|
+        node = graph.get_tensor_by_name(k)
         @lines << "node {"
         @lines << "  name: #{node.name.to_json}"
         if node.is_a?(TensorStream::Operation)
@@ -16,16 +17,16 @@ module TensorStream
             @lines << "  input: #{input.name.to_json}"
           end
           # type
-          pb_attr('T', "dtype: #{sym_to_protobuf_type(node.data_type)}")
+          pb_attr('T', "type: #{sym_to_protobuf_type(node.data_type)}")
           process_options(node)
         elsif node.is_a?(TensorStream::Tensor) && node.is_const
           @lines << "  op: \"Const\""
           # type
-          pb_attr('T', "dtype: #{sym_to_protobuf_type(node.data_type)}")
+          pb_attr('T', "type: #{sym_to_protobuf_type(node.data_type)}")
           pb_attr('value', tensor_value(node))
         elsif node.is_a?(TensorStream::Variable)
           @lines << "  op: \"VariableV2\""
-          pb_attr('T', "dtype: #{sym_to_protobuf_type(node.data_type)}")
+          pb_attr('T', "type: #{sym_to_protobuf_type(node.data_type)}")
           pb_attr('shape', shape_buf(node, 'shape'))
           process_options(node)
         end
@@ -42,16 +43,16 @@ module TensorStream
     def process_options(node)
       return if node.options.nil?
       node.options.each do |k, v|
-        next if %w[name].include?(k.to_s)
+        next if %w[name].include?(k.to_s) || k.to_s.start_with?('__')
         @lines << "  attr {"
         @lines << "    key: \"#{k}\""
         @lines << "    value {"
         if (v.is_a?(TrueClass) || v.is_a?(FalseClass))
-          @lines << "       b: #{v.to_s}"
+          @lines << "      b: #{v.to_s}"
         elsif (v.is_a?(Integer))
-          @lines << "       int_val: #{v}"
+          @lines << "      int_val: #{v}"
         elsif (v.is_a?(Float))
-          @lines << "       float_val: #{v}"
+          @lines << "      float_val: #{v}"
         end
         @lines << "    }"
         @lines << "  }"
@@ -65,7 +66,7 @@ module TensorStream
     def pack_arr_int(int_arr)
       int_arr.flatten.pack('l*').bytes.map { |b| b.chr =~ /[^[:print:]]/ ? "\\#{sprintf("%o", b).rjust(3, '0')}" : b.chr  }.join
     end
-  
+
     def shape_buf(tensor, shape_type = 'tensor_shape')
       arr = []
       arr << "  #{shape_type} {"
@@ -77,6 +78,7 @@ module TensorStream
       arr << "  }"
       arr
     end
+
     def tensor_value(tensor)
       arr = []
       arr << "tensor {"
@@ -146,5 +148,4 @@ module TensorStream
       @lines << "  }"
     end
   end
-
 end