torch-rb 0.2.3 → 0.3.0

data/lib/torch/native/parser.rb

@@ -75,7 +75,7 @@ module Torch
                 v.is_a?(Tensor)
               when "Tensor?"
                 v.nil? || v.is_a?(Tensor)
-              when "Tensor[]"
+              when "Tensor[]", "Tensor?[]"
                 v.is_a?(Array) && v.all? { |v2| v2.is_a?(Tensor) }
               when "int"
                 if k == "reduction"

data/lib/torch/nn/batch_norm.rb

@@ -70,6 +70,11 @@ module Torch
           momentum: exponential_average_factor, eps: @eps
         )
       end
+
+      def extra_inspect
+        s = "%{num_features}, eps: %{eps}, momentum: %{momentum}, affine: %{affine}, track_running_stats: %{track_running_stats}"
+        format(s, **dict)
+      end
     end
   end
 end

data/lib/torch/nn/conv2d.rb

@@ -18,9 +18,15 @@ module Torch
         F.conv2d(input, @weight, @bias, @stride, @padding, @dilation, @groups)
       end
 
-      # TODO add more parameters
       def extra_inspect
-        format("%s, %s, kernel_size: %s, stride: %s", @in_channels, @out_channels, @kernel_size, @stride)
+        s = String.new("%{in_channels}, %{out_channels}, kernel_size: %{kernel_size}, stride: %{stride}")
+        s += ", padding: %{padding}" if @padding != [0] * @padding.size
+        s += ", dilation: %{dilation}" if @dilation != [1] * @dilation.size
+        s += ", output_padding: %{output_padding}" if @output_padding != [0] * @output_padding.size
+        s += ", groups: %{groups}" if @groups != 1
+        s += ", bias: false" unless @bias
+        s += ", padding_mode: %{padding_mode}" if @padding_mode != "zeros"
+        format(s, **dict)
       end
     end
   end

data/lib/torch/nn/convnd.rb

@@ -23,7 +23,7 @@ module Torch
         if bias
           @bias = Parameter.new(Tensor.new(out_channels))
         else
-          raise NotImplementedError
+          register_parameter("bias", nil)
         end
         reset_parameters
       end

data/lib/torch/nn/max_poolnd.rb

@@ -12,7 +12,8 @@ module Torch
       end
 
       def extra_inspect
-        format("kernel_size: %s", @kernel_size)
+        s = "kernel_size: %{kernel_size}, stride: %{stride}, padding: %{padding}, dilation: %{dilation}, ceil_mode: %{ceil_mode}"
+        format(s, **dict)
       end
     end
   end

data/lib/torch/nn/module.rb

@@ -145,7 +145,7 @@ module Torch
         params = {}
         if recurse
           named_children.each do |name, mod|
-            params.merge!(mod.named_parameters(prefix: "#{name}.", recurse: recurse))
+            params.merge!(mod.named_parameters(prefix: "#{prefix}#{name}.", recurse: recurse))
           end
         end
         instance_variables.each do |name|
@@ -186,8 +186,22 @@ module Torch
         named_modules.values
       end
 
-      def named_modules
-        {"" => self}.merge(named_children)
+      # TODO return enumerator?
+      def named_modules(memo: nil, prefix: "")
+        ret = {}
+        memo ||= Set.new
+        unless memo.include?(self)
+          memo << self
+          ret[prefix] = self
+          named_children.each do |name, mod|
+            next unless mod.is_a?(Module)
+            submodule_prefix = prefix + (!prefix.empty? ? "." : "") + name
+            mod.named_modules(memo: memo, prefix: submodule_prefix).each do |m|
+              ret[m[0]] = m[1]
+            end
+          end
+        end
+        ret
       end
 
       def train(mode = true)
@@ -230,7 +244,9 @@ module Torch
           str = String.new
           str << "#{name}(\n"
           named_children.each do |name, mod|
-            str << "  (#{name}): #{mod.inspect}\n"
+            mod_str = mod.inspect
+            mod_str = mod_str.lines.join("  ")
+            str << "  (#{name}): #{mod_str}\n"
           end
           str << ")"
         end
@@ -270,8 +286,11 @@ module Torch
         str % vars
       end
 
+      # used for format
+      # remove tensors for performance
+      # so we can skip call to inspect
       def dict
-        instance_variables.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
+        instance_variables.reject { |k| instance_variable_get(k).is_a?(Tensor) }.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
       end
     end
   end

data/lib/torch/optim/optimizer.rb

@@ -32,9 +32,11 @@ module Torch
       end
 
       def state_dict
+        raise NotImplementedYet
+
         pack_group = lambda do |group|
-          packed = group.select { |k, _| k != :params }.to_h
-          packed[:params] = group[:params].map { |p| p.object_id }
+          packed = group.select { |k, _| k != :params }.map { |k, v| [k.to_s, v] }.to_h
+          packed["params"] = group[:params].map { |p| p.object_id }
           packed
         end
 
@@ -42,8 +44,8 @@ module Torch
         packed_state = @state.map { |k, v| [k.is_a?(Tensor) ? k.object_id : k, v] }.to_h
 
         {
-          state: packed_state,
-          param_groups: param_groups
+          "state" => packed_state,
+          "param_groups" => param_groups
         }
       end
 

data/lib/torch/optim/rprop.rb

@@ -11,9 +11,6 @@ module Torch
       end
 
       def step(closure = nil)
-        # TODO implement []=
-        raise NotImplementedYet
-
         loss = nil
         if closure
           loss = closure.call

data/lib/torch/tensor.rb

@@ -1,6 +1,7 @@
 module Torch
   class Tensor
     include Comparable
+    include Enumerable
     include Inspector
 
     alias_method :requires_grad?, :requires_grad
@@ -25,14 +26,36 @@ module Torch
       inspect
     end
 
+    def each
+      return enum_for(:each) unless block_given?
+
+      size(0).times do |i|
+        yield self[i]
+      end
+    end
+
+    # TODO make more performant
     def to_a
-      reshape_arr(_flat_data, shape)
+      arr = _flat_data
+      if shape.empty?
+        arr
+      else
+        shape[1..-1].reverse.each do |dim|
+          arr = arr.each_slice(dim)
+        end
+        arr.to_a
+      end
     end
 
-    # TODO support dtype
-    def to(device, non_blocking: false, copy: false)
+    def to(device = nil, dtype: nil, non_blocking: false, copy: false)
+      device ||= self.device
       device = Device.new(device) if device.is_a?(String)
-      _to(device, _dtype, non_blocking, copy)
+
+      dtype ||= self.dtype
+      enum = DTYPE_TO_ENUM[dtype]
+      raise Error, "Unknown type: #{dtype}" unless enum
+
+      _to(device, enum, non_blocking, copy)
     end
 
     def cpu
@@ -64,7 +87,15 @@ module Torch
       if numel != 1
         raise Error, "only one element tensors can be converted to Ruby scalars"
       end
-      _flat_data.first
+      to_a.first
+    end
+
+    def to_i
+      item.to_i
+    end
+
+    def to_f
+      item.to_f
     end
 
     # unsure if this is correct
@@ -80,7 +111,7 @@ module Torch
     def numo
       cls = Torch._dtype_to_numo[dtype]
       raise Error, "Cannot convert #{dtype} to Numo" unless cls
-      cls.cast(_flat_data).reshape(*shape)
+      cls.from_string(_data_str).reshape(*shape)
     end
 
     def new_ones(*size, **options)
@@ -108,15 +139,6 @@ module Torch
       _view(size)
     end
 
-    # value and other are swapped for some methods
-    def add!(value = 1, other)
-      if other.is_a?(Numeric)
-        _add__scalar(other, value)
-      else
-        _add__tensor(other, value)
-      end
-    end
-
     def +(other)
       add(other)
     end
@@ -145,12 +167,25 @@ module Torch
       neg
     end
 
+    def &(other)
+      logical_and(other)
+    end
+
+    def |(other)
+      logical_or(other)
+    end
+
+    def ^(other)
+      logical_xor(other)
+    end
+
     # TODO better compare?
     def <=>(other)
       item <=> other
     end
 
-    # based on python_variable_indexing.cpp
+    # based on python_variable_indexing.cpp and
+    # https://pytorch.org/cppdocs/notes/tensor_indexing.html
     def [](*indexes)
       result = self
       dim = 0
@@ -162,6 +197,8 @@ module Torch
           finish += 1 unless index.exclude_end?
           result = result._slice_tensor(dim, index.begin, finish, 1)
           dim += 1
+        elsif index.is_a?(Tensor)
+          result = result.index([index])
         elsif index.nil?
           result = result.unsqueeze(dim)
           dim += 1
@@ -175,24 +212,37 @@ module Torch
       result
     end
 
-    # TODO
-    # based on python_variable_indexing.cpp
+    # based on python_variable_indexing.cpp and
+    # https://pytorch.org/cppdocs/notes/tensor_indexing.html
     def []=(index, value)
       raise ArgumentError, "Tensor does not support deleting items" if value.nil?
 
-      value = Torch.tensor(value) unless value.is_a?(Tensor)
+      value = Torch.tensor(value, dtype: dtype) unless value.is_a?(Tensor)
 
       if index.is_a?(Numeric)
-        copy_to(_select_int(0, index), value)
+        index_put!([Torch.tensor(index)], value)
       elsif index.is_a?(Range)
         finish = index.end
         finish += 1 unless index.exclude_end?
-        copy_to(_slice_tensor(0, index.begin, finish, 1), value)
+        _slice_tensor(0, index.begin, finish, 1).copy!(value)
+      elsif index.is_a?(Tensor)
+        index_put!([index], value)
       else
         raise Error, "Unsupported index type: #{index.class.name}"
       end
     end
 
+    # native functions that need manually defined
+
+    # value and other are swapped for some methods
+    def add!(value = 1, other)
+      if other.is_a?(Numeric)
+        _add__scalar(other, value)
+      else
+        _add__tensor(other, value)
+      end
+    end
+
     # native functions overlap, so need to handle manually
     def random!(*args)
       case args.size
@@ -205,22 +255,9 @@ module Torch
       end
     end
 
-    private
-
-    def copy_to(dst, src)
-      dst.copy!(src)
-    end
-
-    def reshape_arr(arr, dims)
-      if dims.empty?
-        arr
-      else
-        arr = arr.flatten
-        dims[1..-1].reverse.each do |dim|
-          arr = arr.each_slice(dim)
-        end
-        arr.to_a
-      end
+    def clamp!(min, max)
+      _clamp_min_(min)
+      _clamp_max_(max)
     end
   end
 end

data/lib/torch/utils/data/data_loader.rb

@@ -6,9 +6,10 @@ module Torch
 
         attr_reader :dataset
 
-        def initialize(dataset, batch_size: 1)
+        def initialize(dataset, batch_size: 1, shuffle: false)
           @dataset = dataset
           @batch_size = batch_size
+          @shuffle = shuffle
         end
 
         def each
@@ -16,11 +17,15 @@ module Torch
           # this makes it easy to compare results
           base_seed = Torch.empty([], dtype: :int64).random!.item
 
-          max_size = @dataset.size
-          size.times do |i|
-            start_index = i * @batch_size
-            end_index = [start_index + @batch_size, max_size].min
-            batch = (end_index - start_index).times.map { |j| @dataset[start_index + j] }
+          indexes =
+            if @shuffle
+              Torch.randperm(@dataset.size).to_a
+            else
+              @dataset.size.times
+            end
+
+          indexes.each_slice(@batch_size) do |idx|
+            batch = idx.map { |i| @dataset[i] }
             yield collate(batch)
           end
         end

data/lib/torch/utils/data/dataset.rb

@@ -0,0 +1,8 @@
+module Torch
+  module Utils
+    module Data
+      class Dataset
+      end
+    end
+  end
+end

data/lib/torch/utils/data/tensor_dataset.rb

@@ -1,7 +1,7 @@
 module Torch
   module Utils
     module Data
-      class TensorDataset
+      class TensorDataset < Dataset
         def initialize(*tensors)
           unless tensors.all? { |t| t.size(0) == tensors[0].size(0) }
             raise Error, "Tensors must all have same dim 0 size"

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.2.3"
+  VERSION = "0.3.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.2.3
+  version: 0.3.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2020-04-28 00:00:00.000000000 Z
+date: 2020-07-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -95,19 +95,19 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '0'
 - !ruby/object:Gem::Dependency
-  name: npy
+  name: torchvision
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 0.1.1
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 0.1.1
 description: 
 email: andrew@chartkick.com
 executables: []
@@ -260,6 +260,7 @@ files:
 - lib/torch/optim/sgd.rb
 - lib/torch/tensor.rb
 - lib/torch/utils/data/data_loader.rb
+- lib/torch/utils/data/dataset.rb
 - lib/torch/utils/data/tensor_dataset.rb
 - lib/torch/version.rb
 homepage: https://github.com/ankane/torch.rb