tensorflow-ruby 0.2.0

data/lib/tensorflow/data/dataset.rb

@@ -0,0 +1,106 @@
+module Tensorflow
+  module Data
+    class Dataset
+      # Copied from Python code
+      DEFAULT_READER_BUFFER_SIZE_BYTES = 256 * 1024  # 256 KB
+
+      include Enumerable
+
+      # TODO remove
+      attr_reader :output_types, :output_shapes, :variant_tensor
+
+      def self.to_tensor_array(values)
+        case values
+          when Numo::NArray
+            [Tensor.new(values)]
+          when Tensor
+            [values]
+          when Array
+            values.to_a.map do |v|
+              if v.is_a?(Tensor)
+                v
+              else
+                Tensor.new(v)
+              end
+            end
+          when Graph::Operation
+            [values]
+          else
+            raise(Error::UnimplementedError, "Unsupported dataset element: #{values}")
+        end
+      end
+
+      def self.from_tensors(tensors)
+        TensorDataset.new(tensors)
+      end
+
+      def self.from_tensor_slices(tensors)
+        TensorSliceDataset.new(tensors)
+      end
+
+      def initialize(variant_tensor)
+        @variant_tensor = variant_tensor
+      end
+
+      def to_ptr
+        @variant_tensor.to_ptr
+      end
+
+      def with_options(options)
+
+      end
+
+      def batch(batch_size, drop_remainder: false)
+        BatchDataset.new(self, batch_size, drop_remainder)
+      end
+
+      def shuffle(buffer_size)
+        ShuffleDataset.new(self, buffer_size)
+      end
+
+      def make_one_shot_iterator
+        OneShotIterator.new(self)
+      end
+
+      def make_initializable_iterator(shared_name: '')
+        InitializableIterator.new(self, shared_name: shared_name)
+      end
+
+      def each
+        iterator, deleter = RawOps.anonymous_iterator_v2(output_types: @output_types, output_shapes: @output_shapes)
+        RawOps.make_iterator(@variant_tensor, iterator)
+        begin
+          loop do
+            values = RawOps.iterator_get_next_sync(iterator, output_types: @output_types, output_shapes: @output_shapes)
+            yield values
+          end
+        rescue Error::OutOfRangeError
+        end
+      ensure
+        RawOps.delete_iterator(iterator, deleter) if iterator
+      end
+
+      # !!! DEBUG method. You don't want to use this method it because it iterates over
+      # the entire dataset and reads it into a ruby array in memory
+      def data
+        self.map do |slice|
+          if slice.is_a?(Array)
+            slice.map do |tensor|
+              tensor.value
+            end
+          else
+            slice.value
+          end
+        end
+      end
+
+      def map_func(func)
+        MapDataset.new(self, func)
+      end
+
+      def repeat(count)
+        RepeatDataset.new(self, 3)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/fixed_length_record_dataset.rb

@@ -0,0 +1,27 @@
+module Tensorflow
+  module Data
+    class FixedLengthRecordDataset < Dataset
+      def initialize(filenames, record_bytes, header_bytes: 0, footer_bytes: 0,
+                     buffer_size: DEFAULT_READER_BUFFER_SIZE_BYTES, compression_type: '', num_parallel_reads: 0)
+        @output_types = [:string]
+        @output_shapes = [[]]
+
+        record_bytes_tensor = Tensor.new(record_bytes, dtype: :int64)
+        header_bytes_tensor = Tensor.new(header_bytes, dtype: :int64)
+        footer_bytes_tensor = Tensor.new(footer_bytes, dtype: :int64)
+        buffer_size_tensor = Tensor.new(buffer_size, dtype: :int64)
+
+        variant_tensor = RawOps.fixed_length_record_dataset_v2(filenames,
+                                                               header_bytes_tensor,
+                                                               record_bytes_tensor,
+                                                               footer_bytes_tensor,
+                                                               buffer_size_tensor,
+                                                               compression_type)
+
+        super(variant_tensor)
+      end
+
+
+    end
+  end
+end

data/lib/tensorflow/data/iterator.rb

@@ -0,0 +1,76 @@
+module Tensorflow
+  module Data
+    class Iterator
+      attr_reader :output_types, :output_shapes
+
+      def self.from_structure(output_types, output_shapes=[], shared_name: '')
+        ReinitializableIterator.new(output_types, output_shapes, shared_name: shared_name)
+      end
+
+      def initialize(output_types, output_shapes=[])
+        @output_types = output_types
+        @output_shapes = output_shapes
+      end
+
+      def get_next
+        RawOps.iterator_get_next(@iterator, output_types: self.output_types, output_shapes: self.output_shapes)
+      end
+    end
+
+    class OneShotIterator < Iterator
+      def initialize(dataset)
+        super(dataset.output_types, dataset.output_shapes)
+        create_one_shot_iterator(dataset)
+      end
+
+      private
+
+      def create_one_shot_iterator(dataset)
+        function = make_dataset_function(dataset)
+        ExecutionContext.current.add_function(function)
+        @iterator = RawOps.one_shot_iterator(dataset_factory: function, output_types: self.output_types, output_shapes: self.output_shapes)
+      end
+
+      def make_dataset_function(dataset)
+        function = Graph::Graph.new.as_default do |func_graph|
+          optimize = RawOps.optimize_dataset(dataset.variant_tensor, ['noop_elimination'],
+                                             output_types: self.output_types, output_shapes: self.output_shapes)
+          func_graph.to_function('MakeDataset', nil, nil, [optimize])
+        end
+      end
+    end
+
+    class InitializableIterator < Iterator
+      attr_reader :initializer
+
+      def initialize(dataset, shared_name: '')
+        super(dataset.output_types, dataset.output_shapes)
+        create_initializable_iterator(dataset, shared_name)
+      end
+
+      private
+
+      def create_initializable_iterator(dataset, shared_name)
+        @iterator = RawOps.iterator_v2(shared_name: shared_name, output_types: self.output_types, output_shapes: self.output_shapes)
+        @initializer = RawOps.make_iterator(dataset.variant_tensor, @iterator)
+      end
+    end
+
+    class ReinitializableIterator < Iterator
+      def initialize(output_types, output_shapes, shared_name: '')
+        super(output_types, output_shapes)
+        create_iterator_from_structure(shared_name)
+      end
+
+      def make_initializer(dataset)
+        RawOps.make_iterator(dataset.variant_tensor, @iterator)
+      end
+
+      private
+
+      def create_iterator_from_structure(shared_name)
+        @iterator = RawOps.iterator_v2(shared_name: shared_name, output_types: self.output_types, output_shapes: self.output_shapes)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/map_dataset.rb

@@ -0,0 +1,17 @@
+module Tensorflow
+  module Data
+    class MapDataset < Dataset
+      def initialize(input_dataset, function, other_arguments: [])
+        @output_types = function.output_types
+        @output_shapes = function.output_shapes
+
+        variant_tensor = RawOps.map_dataset(input_dataset.variant_tensor, other_arguments,
+                                            f: function,
+                                            output_types: @output_types,
+                                            output_shapes: @output_shapes)
+
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/repeat_dataset.rb

@@ -0,0 +1,16 @@
+module Tensorflow
+  module Data
+    class RepeatDataset < Dataset
+      def initialize(dataset, count)
+        @output_types = dataset.output_types
+        @output_shapes = dataset.output_shapes
+
+        variant_tensor = RawOps.repeat_dataset(dataset.variant_tensor, Tensor.new(count, dtype: :int64),
+                                               output_types: @output_types,
+                                               output_shapes: @output_shapes)
+
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/shuffle_dataset.rb

@@ -0,0 +1,23 @@
+module Tensorflow
+  module Data
+    class ShuffleDataset < Dataset
+      def initialize(input_dataset, buffer_size)
+        @input_dataset = input_dataset
+        @output_types = input_dataset.output_types
+        @output_shapes = input_dataset.output_shapes
+
+        buffer_size = Tensor.new(buffer_size, dtype: :int64)
+        seed = Tensor.new(::Random.rand(100_000_000), dtype: :int64)
+        seed2 = Tensor.new(::Random.rand(100_000_000), dtype: :int64)
+
+        variant_tensor = RawOps.shuffle_dataset(input_dataset.variant_tensor,
+                                                buffer_size,
+                                                seed,
+                                                seed2,
+                                                output_types: @output_types,
+                                                output_shapes: @output_shapes)
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/tensor_dataset.rb

@@ -0,0 +1,19 @@
+module Tensorflow
+  module Data
+    class TensorDataset < Dataset
+      def initialize(elements)
+        @tensors = self.class.to_tensor_array(elements)
+        @output_types = @tensors.map(&:dtype)
+        @output_shapes = @tensors.map do |tensor|
+          tensor.shape
+        end
+
+        variant_tensor = RawOps.tensor_dataset(@tensors,
+                                               toutput_types: @output_types,
+                                               output_shapes: @output_shapes)
+
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/tensor_slice_dataset.rb

@@ -0,0 +1,15 @@
+module Tensorflow
+  module Data
+    class TensorSliceDataset < Dataset
+      def initialize(elements)
+        @tensors = self.class.to_tensor_array(elements)
+        @output_types = @tensors.map(&:dtype)
+        @output_shapes = @tensors.map do |tensor|
+          tensor.shape[1..]
+        end
+        variant_tensor = RawOps.tensor_slice_dataset(@tensors, toutput_types: @output_types, output_shapes: @output_shapes)
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/tf_record_dataset.rb

@@ -0,0 +1,18 @@
+module Tensorflow
+  module Data
+    class TfRecordDataset < Dataset
+      DEFAULT_BUFFER_SIZE = 256 * 1_048_576  # 256 MB
+
+      def initialize(filenames, compression_type='', buffer_size=DEFAULT_BUFFER_SIZE)
+        filenames = Array(filenames)
+        @output_types = [:string]
+        @output_shapes = [[]]
+
+        buffer_size = Tensor.new(buffer_size, dtype: :int64) if buffer_size
+        variant_tensor = RawOps.tf_record_dataset(filenames, compression_type, buffer_size)
+
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/data/zip_dataset.rb

@@ -0,0 +1,24 @@
+module Tensorflow
+  module Data
+    class ZipDataset < Dataset
+      def initialize(*datasets)
+        datasets = datasets.flatten(1)
+        tensors = Array.new(datasets.size)
+        output_types = Array.new(datasets.size)
+        output_shapes = Array.new(datasets.size)
+
+        datasets.each_with_index do |dataset, i|
+          tensors[i] = dataset.variant_tensor
+          output_types[i] = dataset.output_types
+          output_shapes[i] = dataset.output_shapes
+        end
+
+        @output_types = output_types.flatten
+        @output_shapes = output_shapes.flatten(1)
+        variant_tensor = RawOps.zip_dataset(tensors, n: tensors.count, output_types: @output_types, output_shapes: @output_shapes)
+
+        super(variant_tensor)
+      end
+    end
+  end
+end

data/lib/tensorflow/decorators.rb

@@ -0,0 +1,53 @@
+module Tensorflow
+  module Decorator
+    class Function
+      attr_reader :input_signatures
+
+      def initialize(input_signatures = [])
+        @input_signatures = input_signatures
+      end
+
+      def wrap(method)
+        Graph::FunctionDef.new(method, self.input_signatures)
+      end
+    end
+
+    def self.extended(klass)
+      @waiting_for_method = false
+      this = self
+      klass.instance_eval do
+        @tf = this
+      end
+
+      if klass.is_a?(Object) && klass.to_s == 'main'
+        klass.class.extend(self)
+      end
+    end
+
+    def self.function(input_signature = [])
+      @current_function = Function.new(input_signature)
+    end
+
+    def self.wrap_method(method)
+      # We do this little dance because when the method is wrapped it will trigger method_added. So first we need
+      # to clear out @current_function before continuing
+      if @current_function
+        current_function = @current_function
+        @current_function = nil
+        current_function&.wrap(method)
+      end
+    end
+
+    def singleton_method_added(method_name)
+      super(method_name)
+      method = self.method(method_name)
+      @tf.wrap_method(method)
+    end
+
+    def method_added(method_name)
+      super(method_name)
+      method = self.instance_method(method_name)
+      @tf.wrap_method(method)
+    end
+  end
+end

data/lib/tensorflow/eager/context.rb

@@ -0,0 +1,120 @@
+module Tensorflow
+  module Eager
+    class Context
+      extend Forwardable
+      def_delegators :@name_scope, :name_scope, :scoped_name, :unique_name
+
+      def self.default
+        @default ||= Context.new
+      end
+
+      def self.finalize(pointer)
+        proc { FFI.TFE_DeleteContext(pointer) }
+      end
+
+      def initialize
+        @name_scope = NameScope.new
+        options = FFI.TFE_NewContextOptions
+        Status.check do |status|
+          @pointer = FFI.TFE_NewContext(options, status)
+        end
+        ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
+        FFI.TFE_DeleteContextOptions(options)
+      end
+
+      def as_default
+        raise(Error::InvalidArgumentError, "Must provide block") unless block_given?
+        ExecutionContext.push(self)
+        begin
+          yield self
+        ensure
+          ExecutionContext.pop
+        end
+      end
+
+      def create_operation(op_type, inputs=[], attrs={})
+        Operation.new(self, op_type, inputs, attrs)
+      end
+
+      def execute(operation)
+        # TODO decide how many retvals to allocate
+        retvals = ::FFI::MemoryPointer.new(:pointer, 10)
+        num_retvals = ::FFI::MemoryPointer.new(:int)
+        num_retvals.write_int(retvals.size)
+
+        Status.check do |status|
+          FFI.TFE_Execute(operation, retvals, num_retvals, status)
+        end
+
+        n = num_retvals.read_int
+        if n > 0
+          handles = retvals.read_array_of_pointer(n).map do |handle|
+            TensorHandle.new(self, handle)
+          end
+
+          # TODO handle case where n = 1 and still want an array for retvals
+          n == 1 ? handles.first : handles
+        end
+      end
+
+      def device_policy
+        FFI::ContextDevicePlacementPolicy[FFI.TFE_ContextGetDevicePlacementPolicy(@pointer)]
+      end
+
+      def enable_run_metadata
+        FFI.TFE_ContextEnableRunMetadata(@pointer)
+      end
+
+      def disable_run_metadata
+        FFI.TFE_ContextDisableRunMetadata(@pointer)
+      end
+
+      def start_step
+        FFI.TFE_ContextStartStep(@pointer)
+      end
+
+      def end_step
+        FFI.TFE_ContextEndStep(@pointer)
+      end
+
+      def to_ptr
+        @pointer
+      end
+
+      def shared_name
+        # hard-coded in Python library
+        "cd2c89b7-88b7-44c8-ad83-06c2a9158347"
+      end
+
+      # Mimic graph api
+      def add_to_collection(name, value)
+      end
+
+      # Mimic graph api
+      def add_to_collections(names, value)
+      end
+
+      def get_collection_ref(name)
+      end
+
+      def add_function(function)
+        Status.check do |status|
+          FFI.TFE_ContextAddFunction(self, function, status)
+        end
+      end
+
+      def remove_function(function)
+        name = function.is_a?(Graph::Function) ? function.name : function
+        Status.check do |status|
+          FFI.TFE_ContextRemoveFunction(self, name, status)
+        end
+      end
+
+      def function?(function)
+        name = function.is_a?(Graph::Function) ? function.name : function
+        # result is uchar
+        FFI.TFE_ContextHasFunction(self, name) != 0
+      end
+    end
+  end
+end