tensor_stream 0.7.0 → 0.8.0

data/lib/tensor_stream/tensor_shape.rb

@@ -30,6 +30,10 @@ module TensorStream
       true
     end
 
+    def is_fully_defined?
+      known?
+    end
+
     def self.infer_shape(shape_a, shape_b)
       return shape_a if shape_b.nil?
       return shape_b if shape_a.nil?
@@ -69,6 +73,7 @@ module TensorStream
 
     def self.fix_inferred_elements(shape, total_size)
       return shape if shape.empty?
+      return nil if shape[0].is_a?(Tensor)
 
       current_size = shape.inject(1) { |product, n| n > 0 ? product * n : product }
       inferred_size = total_size.nil? ? nil : total_size / current_size

data/lib/tensor_stream/train/gradient_descent_optimizer.rb

@@ -1,7 +1,7 @@
 module TensorStream
   module Train
     # High Level implementation of the gradient descent algorithm
-    class GradientDescentOptimizer
+    class GradientDescentOptimizer < Optimizer
       include TensorStream::OpHelper
 
       attr_accessor :learning_rate
@@ -10,43 +10,10 @@ module TensorStream
         @learning_rate = learning_rate
       end
 
-      def minimize(loss, var_list: nil, grad_loss: nil, global_step: nil)
-        grads_and_vars = compute_gradients(loss, var_list: var_list, grad_loss: grad_loss)
-        apply_gradients(grads_and_vars, global_step: global_step)
-      end
-
-      ##
-      # Apply gradients to variables.
-      # This is the second part of minimize(). It returns an Operation that applies gradients.
-      def apply_gradients(grads_and_vars, global_step: nil)
-        apply_ops = grads_and_vars.map do |grad, var|
-          i_op(:apply_gradient_descent, var, TensorStream.cast(@learning_rate, grad.data_type), grad)
-        end
-
-        if global_step.nil?
-          apply_ops
-        else
-          apply_ops + [global_step.assign_add(1)]
-        end
-      end
-
-      ##
-      # Compute gradients of loss for the variables in var_list.
-      #
-      # This is the first part of minimize(). It returns a list of (gradient, variable) pairs where "gradient" is the gradient for "variable".
-      def compute_gradients(loss, var_list: nil, grad_loss: nil)
-        trainable_vars = if var_list
-                           raise "var_list must be an array" unless var_list.is_a?(Array)
-                           var_list.each_with_index { |var, index| raise "var #{index} not a Variable" unless var.is_a?(Variable) }
+      protected
 
-                           var_list
-                         else
-                           loss.graph.get_collection(TensorStream::GraphKeys::TRAINABLE_VARIABLES)
-                         end
-        all_grads = grad_loss || TensorStream.gradients(loss, trainable_vars)
-        trainable_vars.each_with_index.collect do |var, index|
-          [all_grads[index], var]
-        end
+      def apply_dense(grad, var)
+        i_op(:apply_gradient_descent, var, TensorStream.cast(@learning_rate, grad.data_type), grad)
       end
     end
   end

data/lib/tensor_stream/train/momentum_optimizer.rb

@@ -0,0 +1,48 @@
+module TensorStream
+  module Train
+    # Optimizer that implements the Momentum algorithm. loosely based on the tensorflow implementation.
+    class MomentumOptimizer < Optimizer
+      include OpHelper
+
+      ##
+      # Construct a new Momentum optimizer.
+      #
+      # Args:
+      #   learning_rate: A Tensor or a floating point value that indicates the learning rate
+      #   momentum: A Tensor or a floating point value for the momentum
+      #   name: Optional name prefix
+      #   use_nesterov: boolean - Flag that indicates if nesterov momentum is to be used. http://jmlr.org/proceedings/papers/v28/sutskever13.pdf
+      #   use_locking: boolean - filler argument for compatibility, not used at the moment
+      def initialize(learning_rate, momentum, name: 'momentum', use_nesterov: false, use_locking: false)
+        @learning_rate = learning_rate
+        @momentum = momentum
+        @use_nesterov = use_nesterov
+        super(name: name, use_locking: use_locking)
+      end
+
+      protected
+
+      def prepare
+        @learning_rate_tensor = TensorStream.convert_to_tensor(@learning_rate, name: "learning_rate")
+        @momentum_tensor = TensorStream.convert_to_tensor(@momentum, name: "momentum")
+      end
+
+      def create_slots(var_list)
+        var_list.each do |v|
+          zeros_slot(v, "momentum", @name)
+        end
+      end
+
+      def apply_dense(grad, var)
+        mom = get_slot(var, "momentum")
+
+        _op(:apply_momentum, var, mom,
+            TensorStream.cast(@learning_rate_tensor, var.data_type),
+            grad,
+            TensorStream.cast(@momentum_tensor, var.data_type),
+            use_locking: @use_locking,
+            use_nesterov: @use_nesterov)
+      end
+    end
+  end
+end

data/lib/tensor_stream/train/optimizer.rb

@@ -0,0 +1,129 @@
+module TensorStream
+  module Train
+    # Base class for an optimizer
+    # This is a straight up port from the python version
+    class Optimizer
+      include SlotCreator
+
+      attr_reader :name
+
+      def initialize(name:, use_locking:)
+        @name = name
+        @use_locking = use_locking
+        raise TensorStream::ValueError, "Must specify the optimizer name" unless @name
+        @slots = {}
+      end
+
+      def minimize(loss, var_list: nil, grad_loss: nil, global_step: nil, name: nil)
+        grads_and_vars = compute_gradients(loss, var_list: var_list, grad_loss: grad_loss)
+        apply_gradients(grads_and_vars, global_step: global_step, name: name)
+      end
+
+      ##
+      # Apply gradients to variables.
+      # This is the second part of minimize(). It returns an Operation that applies gradients.
+      def apply_gradients(grads_and_vars, global_step: nil, name: nil)
+        varlist = grads_and_vars.map { |_grad, var| var }
+        create_slots(varlist)
+        TensorStream.name_scope(name, default: @name) do
+          prepare
+          apply_ops = grads_and_vars.map do |grad, var|
+            TensorStream.name_scope("update_" + var.op.name) do
+              apply_dense(grad, var)
+            end
+          end
+
+          if global_step.nil?
+            finish(apply_ops, name)
+          else
+            TensorStream.control_dependencies([finish(apply_ops, "update")]) do
+              global_step.assign_add(1)
+            end
+          end
+        end
+      end
+
+      ##
+      # Compute gradients of loss for the variables in var_list.
+      #
+      # This is the first part of minimize(). It returns a list of (gradient, variable) pairs where "gradient" is the gradient for "variable".
+      def compute_gradients(loss, var_list: nil, grad_loss: nil)
+        trainable_vars = if var_list
+                           raise "var_list must be an array" unless var_list.is_a?(Array)
+                           var_list.each_with_index { |var, index| raise "var #{index} not a Variable" unless var.is_a?(Variable) }
+
+                           var_list
+                         else
+                           loss.graph.get_collection(TensorStream::GraphKeys::TRAINABLE_VARIABLES)
+                         end
+        all_grads = grad_loss || TensorStream.gradients(loss, trainable_vars)
+        trainable_vars.each_with_index.collect do |var, index|
+          [all_grads[index], var]
+        end
+      end
+
+      def get_slot(var, name)
+        named_slots = @slots.fetch(name, nil)
+        return nil if named_slots.nil?
+        named_slots.fetch(var_key(var), nil)
+      end
+
+      def get_slot_names
+        @slots.keys.sort
+      end
+
+      protected
+
+      def finish(update_ops, name_scope)
+        TensorStream.group(update_ops, name: name_scope)
+      end
+
+      def create_slots(var_list)
+        # no implementation
+      end
+
+      def prepare
+        # no implementation
+      end
+
+      def apply_dense(grad, var)
+        raise TensorStream::NotImplementedError, "not implemented"
+      end
+
+      ##
+      # Find or create a slot initialized with 0.0.
+      #
+      # Args:
+      #   var: Variable - A Variable object
+      #   slot_name: string - Name fot the slot
+      #   op_name: string - Name to use when scoping the Variable that needs to be created
+      def zeros_slot(var, slot_name, op_name)
+        named_slots = slot_dict(slot_name)
+        if !named_slots.key?(var_key(var))
+          named_slots[var_key(var)] = create_zeros_slot(var, op_name)
+        end
+        named_slots[var_key(var)]
+      end
+
+      ##
+      # Returns a dict for caching slots created under the given name.
+      #
+      # Args:
+      # slot_name string Name for the slot
+      #
+      # Returns: A dict that maps primary 'Variable' objects to the slot created
+      def slot_dict(slot_name)
+        named_slots = @slots.fetch(slot_name, nil)
+        if named_slots.nil?
+          named_slots = {}
+          @slots[slot_name] = named_slots
+        end
+        named_slots
+      end
+
+      def var_key(var)
+        [var.op.graph, var.op.name]
+      end
+    end
+  end
+end

data/lib/tensor_stream/train/saver.rb

@@ -55,7 +55,6 @@ module TensorStream
         build_save: true,
         build_restore: true)
         saveables = _validate_and_slice_inputs(names_to_saveables)
-
       end
 
       def _validate_and_slice_inputs(names_to_saveables)

data/lib/tensor_stream/train/slot_creator.rb

@@ -0,0 +1,62 @@
+module TensorStream
+  module Train
+    module SlotCreator
+      include TensorStream::Utils
+
+      ##
+      # Helper function for creating a slot variable.
+      def create_slot_var(primary, val, scope)
+        slot = get_variable(scope, initializer: val, trainable: false,
+                                                  validate_shape: val.shape.is_fully_defined?)
+        slot
+      end
+
+      ##
+      # Create a slot initialized to the given value
+      #
+      # Args:
+      #   primary: Variable - The primary 'Variable' or 'Tensor'
+      #   val: Tensor - A `Tensor` specifying the initial value of the slot
+      #   name: String - Name to use for the slot variable
+      #   colocate_with_primary: Boolean - If true the slot is located
+      #                                    on the same device as `primary`
+      #
+      # Returns: A `Variable` object
+      def create_slot(primary, val, name, colocate_with_primary: true)
+        TensorStream.variable_scope(primary.op.name + "/" + name) do
+          if colocate_with_primary
+            TensorStream.colocate_with(primary) do
+              return create_slot_var(primary, val, "")
+            end
+          else
+            return create_slot_var(primary, val, "")
+          end
+        end
+      end
+
+      ##
+      # Create a slot initialized to 0 with same shape as the primary object.
+      #
+      # Args:
+      #   primary: The pirmary variable or Tensor
+      #   name: String - Name to use for the slot variable
+      #   dtype: Symbol - Type of the slot variable
+      #   colocate_with_primary: boolean - If true the slot is located on the same device as primary
+      #
+      # Returns:
+      #   A `Variable` object
+      def create_zeros_slot(primary, name, dtype: nil, colocate_with_primary: true)
+        dtype = primary.data_type if dtype.nil?
+        slot_shape = primary.shape
+        slot_shape = if slot_shape.is_fully_defined?
+                       slot_shape.shape
+                     else
+                       TensorStream.shape(primary.initialized_value)
+                     end
+        val = TensorStream.zeros(slot_shape, dtype: dtype)
+        create_slot(primary, val, name,
+                           colocate_with_primary: colocate_with_primary)
+      end
+    end
+  end
+end

data/lib/tensor_stream/train/utils.rb

@@ -6,24 +6,23 @@ module TensorStream
         target_graph = graph || TensorStream.get_default_graph
         raise TensorStream::ValueError, '"global_step" already exists.' unless get_global_step(target_graph).nil?
 
-        TensorStream.variable_scope.get_variable(
-          TensorStream::GraphKeys::GLOBAL_STEP, shape: [],
-                                                dtype: :int64,
-                                                initializer: TensorStream.zeros_initializer,
-                                                trainable: false,
-                                                collections: [TensorStream::GraphKeys::GLOBAL_VARIABLES,
-                                                              TensorStream::GraphKeys::GLOBAL_STEP])
+        TensorStream.variable_scope.get_variable(TensorStream::GraphKeys::GLOBAL_STEP, shape: [],
+                                                                                       dtype: :int64,
+                                                                                       initializer: TensorStream.zeros_initializer,
+                                                                                       trainable: false,
+                                                                                       collections: [TensorStream::GraphKeys::GLOBAL_VARIABLES,
+                                                                                                     TensorStream::GraphKeys::GLOBAL_STEP])
       end
 
       def get_global_step(graph = nil)
         target_graph = graph || TensorStream.get_default_graph
         global_step_tensors = target_graph.get_collection(TensorStream::GraphKeys::GLOBAL_STEP)
         global_step_tensor = if global_step_tensors.nil? || global_step_tensors.empty?
-                                begin
-                                  target_graph.get_tensor_by_name('global_step:0')
-                                rescue TensorStream::KeyError
-                                  nil
-                                end
+                               begin
+                                 target_graph.get_tensor_by_name('global_step:0')
+                               rescue TensorStream::KeyError
+                                 nil
+                               end
                              elsif global_step_tensors.size == 1
                                global_step_tensors[0]
                              else

data/lib/tensor_stream/trainer.rb

@@ -1,4 +1,7 @@
+require 'tensor_stream/train/slot_creator'
+require 'tensor_stream/train/optimizer'
 require 'tensor_stream/train/gradient_descent_optimizer'
+require 'tensor_stream/train/momentum_optimizer'
 require 'tensor_stream/train/saver'
 
 module TensorStream

data/lib/tensor_stream/utils.rb

@@ -54,14 +54,14 @@ module TensorStream
         trainable: trainable
       }
       tensor = if value.is_a?(String)
-        TensorStream::Variable.new(dtype || :string, 0, [], get_variable_scope, common_options)
-      elsif value.is_a?(Integer)
-        TensorStream::Variable.new(dtype || :int32, 0, [], get_variable_scope, common_options)
-      elsif value.is_a?(Float)
-        TensorStream::Variable.new(dtype || :float32, 0, [], get_variable_scope, common_options)
-      else
-        TensorStream::Variable.new(dtype || :float32, 0, nil, get_variable_scope, common_options)
-      end
+                 TensorStream::Variable.new(dtype || :string, 0, [], get_variable_scope, common_options)
+               elsif value.is_a?(Integer)
+                 TensorStream::Variable.new(dtype || :int32, 0, [], get_variable_scope, common_options)
+               elsif value.is_a?(Float)
+                 TensorStream::Variable.new(dtype || :float32, 0, [], get_variable_scope, common_options)
+               else
+                 TensorStream::Variable.new(dtype || :float32, 0, nil, get_variable_scope, common_options)
+               end
       op.inputs[0] = tensor
       tensor
     end
@@ -115,8 +115,13 @@ module TensorStream
       session
     end
 
-    def program(&block)
-      block.call(self)
+    def colocate_with(op, ignore_existing: false)
+      # noop for now
+      yield
+    end
+
+    def program
+      yield self
     end
 
     def layers
@@ -153,7 +158,7 @@ module TensorStream
       TensorStream::DynamicStitch.new(:dynamic_stitch, [indices, data], name: name)
     end
 
-    def get_variable(name, dtype: nil, shape: nil, initializer: nil, trainable: true, collections: nil)
+    def get_variable(name, dtype: nil, shape: nil, initializer: nil, trainable: true, collections: nil, validate_shape: false)
       get_variable_scope.get_variable(name, dtype: dtype, shape: shape, initializer: initializer, trainable: trainable, collections: collections)
     end
 
@@ -184,6 +189,8 @@ module TensorStream
       TensorStream.get_default_graph.get_collection(TensorStream::GraphKeys::TRAINABLE_VARIABLES)
     end
 
+    ##
+    # Sets random seed to use for the default graph
     def set_random_seed(seed)
       TensorStream.get_default_graph.random_seed = seed
     end

data/lib/tensor_stream/variable.rb

@@ -14,10 +14,9 @@ module TensorStream
       scope_name = variable_scope ? variable_scope.name : nil
       variable_scope_initializer = variable_scope ? variable_scope.initializer : nil
       @name = [scope_name, options[:name] || build_name].compact.reject(&:empty?).join('/')
-      @initalizer_tensor = options[:initializer] ? options[:initializer] : variable_scope_initializer || TensorStream.glorot_uniform_initializer
-      if shape.nil? && @initalizer_tensor && @initalizer_tensor.shape
-        shape = @initalizer_tensor.shape.shape
-      end
+      @initalizer_tensor = options[:initializer] || variable_scope_initializer || TensorStream.glorot_uniform_initializer
+      shape = @initalizer_tensor.shape.shape if shape.nil? && @initalizer_tensor && @initalizer_tensor.shape
+
       @shape = TensorShape.new(shape, rank)
       @trainable = options.fetch(:trainable, true)
       @graph.add_variable(self, options)
@@ -34,22 +33,30 @@ module TensorStream
       assign(init_op)
     end
 
+    def initialized_value
+      init_op = @initalizer_tensor.op
+      init_op.shape = @shape || init_op.shape
+      init_op.data_type = @data_type || init_op.data_type
+      init_op
+    end
+
     def assign(value, name: nil)
       _a, value = TensorStream.check_data_types(self, value)
-      Operation.new(:assign, self, value, name: name)
+      _op(:assign, self, value, name: name)
     end
 
     def read_value
-      if buffer
-        @value = buffer.to_ruby
-      end
-
+      @value = buffer.to_ruby if buffer
       @value
     end
 
-    def assign_add(value)
+    def assign_add(value, name: nil)
       _a, value = TensorStream.check_data_types(self, value)
-      Operation.new(:assign_add, self, value, data_type: data_type)
+      _op(:assign_add, self, value, data_type: data_type, name: name)
+    end
+
+    def op
+      @op ||= _op(:variable, self, data_type: data_type)
     end
 
     def to_math(_tensor, _name_only = false, _max_depth = 99, _unused = 0)
@@ -58,7 +65,7 @@ module TensorStream
 
     def assign_sub(value)
       _a, value = TensorStream.check_data_types(self, value)
-      Operation.new(:assign_sub, self, value)
+      _op(:assign_sub, self, value)
     end
 
     def self.variables_initializer(collection)