torch-rb 0.2.0 → 0.2.5

data/ext/torch/extconf.rb

@@ -2,28 +2,33 @@ require "mkmf-rice"
 
 abort "Missing stdc++" unless have_library("stdc++")
 
-$CXXFLAGS << " -std=c++14"
+$CXXFLAGS += " -std=c++14"
 
 # change to 0 for Linux pre-cxx11 ABI version
-$CXXFLAGS << " -D_GLIBCXX_USE_CXX11_ABI=1"
+$CXXFLAGS += " -D_GLIBCXX_USE_CXX11_ABI=1"
 
 # TODO check compiler name
 clang = RbConfig::CONFIG["host_os"] =~ /darwin/i
 
+# check omp first
 if have_library("omp") || have_library("gomp")
-  $CXXFLAGS << " -DAT_PARALLEL_OPENMP=1"
-  $CXXFLAGS << " -Xclang" if clang
-  $CXXFLAGS << " -fopenmp"
+  $CXXFLAGS += " -DAT_PARALLEL_OPENMP=1"
+  $CXXFLAGS += " -Xclang" if clang
+  $CXXFLAGS += " -fopenmp"
 end
 
-# silence ruby/intern.h warning
-$CXXFLAGS << " -Wno-deprecated-register"
-
-# silence torch warnings
 if clang
-  $CXXFLAGS << " -Wno-shorten-64-to-32 -Wno-missing-noreturn"
+  # silence ruby/intern.h warning
+  $CXXFLAGS += " -Wno-deprecated-register"
+
+  # silence torch warnings
+  $CXXFLAGS += " -Wno-shorten-64-to-32 -Wno-missing-noreturn"
 else
-  $CXXFLAGS << " -Wno-duplicated-cond -Wno-suggest-attribute=noreturn"
+  # silence rice warnings
+  $CXXFLAGS += " -Wno-noexcept-type"
+
+  # silence torch warnings
+  $CXXFLAGS += " -Wno-duplicated-cond -Wno-suggest-attribute=noreturn"
 end
 
 inc, lib = dir_config("torch")
@@ -34,22 +39,30 @@ cuda_inc, cuda_lib = dir_config("cuda")
 cuda_inc ||= "/usr/local/cuda/include"
 cuda_lib ||= "/usr/local/cuda/lib64"
 
-with_cuda = Dir["#{lib}/*torch_cuda*"].any? && have_library("cuda") && have_library("cudnn")
+$LDFLAGS += " -L#{lib}" if Dir.exist?(lib)
+abort "LibTorch not found" unless have_library("torch")
+
+have_library("mkldnn")
+have_library("nnpack")
+
+with_cuda = false
+if Dir["#{lib}/*torch_cuda*"].any?
+  $LDFLAGS += " -L#{cuda_lib}" if Dir.exist?(cuda_lib)
+  with_cuda = have_library("cuda") && have_library("cudnn")
+end
 
-$INCFLAGS << " -I#{inc}"
-$INCFLAGS << " -I#{inc}/torch/csrc/api/include"
+$INCFLAGS += " -I#{inc}"
+$INCFLAGS += " -I#{inc}/torch/csrc/api/include"
 
-$LDFLAGS << " -Wl,-rpath,#{lib}"
-$LDFLAGS << ":#{cuda_lib}/stubs:#{cuda_lib}" if with_cuda
-$LDFLAGS << " -L#{lib}"
-$LDFLAGS << " -L#{cuda_lib}" if with_cuda
+$LDFLAGS += " -Wl,-rpath,#{lib}"
+$LDFLAGS += ":#{cuda_lib}/stubs:#{cuda_lib}" if with_cuda
 
 # https://github.com/pytorch/pytorch/blob/v1.5.0/torch/utils/cpp_extension.py#L1232-L1238
-$LDFLAGS << " -lc10 -ltorch_cpu -ltorch"
+$LDFLAGS += " -lc10 -ltorch_cpu -ltorch"
 if with_cuda
-  $LDFLAGS << " -lcuda -lnvrtc -lnvToolsExt -lcudart -lc10_cuda -ltorch_cuda -lcufft -lcurand -lcublas -lcudnn"
+  $LDFLAGS += " -lcuda -lnvrtc -lnvToolsExt -lcudart -lc10_cuda -ltorch_cuda -lcufft -lcurand -lcublas -lcudnn"
   # TODO figure out why this is needed
-  $LDFLAGS << " -Wl,--no-as-needed,#{lib}/libtorch.so"
+  $LDFLAGS += " -Wl,--no-as-needed,#{lib}/libtorch.so"
 end
 
 # generate C++ functions

data/lib/torch.rb

@@ -1,6 +1,11 @@
 # ext
 require "torch/ext"
 
+# stdlib
+require "fileutils"
+require "net/http"
+require "tmpdir"
+
 # native functions
 require "torch/native/generator"
 require "torch/native/parser"
@@ -174,11 +179,9 @@ require "torch/nn/init"
 
 # utils
 require "torch/utils/data/data_loader"
+require "torch/utils/data/dataset"
 require "torch/utils/data/tensor_dataset"
 
-# random
-require "torch/random"
-
 # hub
 require "torch/hub"
 
@@ -317,12 +320,11 @@ module Torch
     end
 
     def save(obj, f)
-      raise NotImplementedYet unless obj.is_a?(Tensor)
-      File.binwrite(f, _save(obj))
+      File.binwrite(f, _save(to_ivalue(obj)))
     end
 
     def load(f)
-      raise NotImplementedYet
+      to_ruby(_load(File.binread(f)))
     end
 
     # --- begin tensor creation: https://pytorch.org/cppdocs/notes/tensor_creation.html ---
@@ -447,6 +449,100 @@ module Torch
 
     private
 
+    def to_ivalue(obj)
+      case obj
+      when String
+        IValue.from_string(obj)
+      when Integer
+        IValue.from_int(obj)
+      when Tensor
+        IValue.from_tensor(obj)
+      when Float
+        IValue.from_double(obj)
+      when Hash
+        dict = {}
+        obj.each do |k, v|
+          dict[to_ivalue(k)] = to_ivalue(v)
+        end
+        IValue.from_dict(dict)
+      when true, false
+        IValue.from_bool(obj)
+      when nil
+        IValue.new
+      when Array
+        if obj.all? { |v| v.is_a?(Tensor) }
+          IValue.from_list(obj.map { |v| IValue.from_tensor(v) })
+        else
+          raise Error, "Unknown list type"
+        end
+      else
+        raise Error, "Unknown type: #{obj.class.name}"
+      end
+    end
+
+    def to_ruby(ivalue)
+      if ivalue.bool?
+        ivalue.to_bool
+      elsif ivalue.double?
+        ivalue.to_double
+      elsif ivalue.int?
+        ivalue.to_int
+      elsif ivalue.none?
+        nil
+      elsif ivalue.string?
+        ivalue.to_string_ref
+      elsif ivalue.tensor?
+        ivalue.to_tensor
+      elsif ivalue.generic_dict?
+        dict = {}
+        ivalue.to_generic_dict.each do |k, v|
+          dict[to_ruby(k)] = to_ruby(v)
+        end
+        dict
+      elsif ivalue.list?
+        ivalue.to_list.map { |v| to_ruby(v) }
+      else
+        type =
+          if ivalue.capsule?
+            "Capsule"
+          elsif ivalue.custom_class?
+            "CustomClass"
+          elsif ivalue.tuple?
+            "Tuple"
+          elsif ivalue.future?
+            "Future"
+          elsif ivalue.r_ref?
+            "RRef"
+          elsif ivalue.int_list?
+            "IntList"
+          elsif ivalue.double_list?
+            "DoubleList"
+          elsif ivalue.bool_list?
+            "BoolList"
+          elsif ivalue.tensor_list?
+            "TensorList"
+          elsif ivalue.object?
+            "Object"
+          elsif ivalue.module?
+            "Module"
+          elsif ivalue.py_object?
+            "PyObject"
+          elsif ivalue.scalar?
+            "Scalar"
+          elsif ivalue.device?
+            "Device"
+          # elsif ivalue.generator?
+          #   "Generator"
+          elsif ivalue.ptr_type?
+            "PtrType"
+          else
+            "Unknown"
+          end
+
+        raise Error, "Unsupported type: #{type}"
+      end
+    end
+
     def tensor_size(size)
       size.flatten
     end

data/lib/torch/hub.rb

@@ -4,6 +4,58 @@ module Torch
       def list(github, force_reload: false)
         raise NotImplementedYet
       end
+
+      def download_url_to_file(url, dst)
+        uri = URI(url)
+        tmp = "#{Dir.tmpdir}/#{Time.now.to_f}" # TODO better name
+        location = nil
+
+        Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |http|
+          request = Net::HTTP::Get.new(uri)
+
+          puts "Downloading #{url}..."
+          File.open(tmp, "wb") do |f|
+            http.request(request) do |response|
+              case response
+              when Net::HTTPRedirection
+                location = response["location"]
+              when Net::HTTPSuccess
+                response.read_body do |chunk|
+                  f.write(chunk)
+                end
+              else
+                raise Error, "Bad response"
+              end
+            end
+          end
+        end
+
+        if location
+          download_url_to_file(location, dst)
+        else
+          FileUtils.mv(tmp, dst)
+          nil
+        end
+      end
+
+      def load_state_dict_from_url(url, model_dir: nil)
+        unless model_dir
+          torch_home = ENV["TORCH_HOME"] || "#{ENV["XDG_CACHE_HOME"] || "#{ENV["HOME"]}/.cache"}/torch"
+          model_dir = File.join(torch_home, "checkpoints")
+        end
+
+        FileUtils.mkdir_p(model_dir)
+
+        parts = URI(url)
+        filename = File.basename(parts.path)
+        cached_file = File.join(model_dir, filename)
+        unless File.exist?(cached_file)
+          # TODO support hash_prefix
+          download_url_to_file(url, cached_file)
+        end
+
+        Torch.load(cached_file)
+      end
     end
   end
 end

data/lib/torch/inspector.rb

@@ -1,6 +1,6 @@
 module Torch
   module Inspector
-    # TODO make more performance, especially when summarizing
+    # TODO make more performant, especially when summarizing
     # how? only read data that will be displayed
     def inspect
       data =
@@ -14,7 +14,7 @@ module Torch
           if dtype == :bool
             fmt = "%s"
           else
-            values = to_a.flatten
+            values = _flat_data
             abs = values.select { |v| v != 0 }.map(&:abs)
             max = abs.max || 1
             min = abs.min || 1
@@ -25,7 +25,7 @@ module Torch
             end
 
             if floating_point?
-              sci = max / min.to_f > 1000 || max > 1e8 || min < 1e-4
+              sci = max > 1e8 || max < 1e-4
 
               all_int = values.all? { |v| v.finite? && v == v.to_i }
               decimal = all_int ? 1 : 4

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

@@ -20,7 +20,14 @@ module Torch
 
       # 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

@@ -67,8 +67,9 @@ module Torch
         self
       end
 
-      def cuda(device: nil)
-        _apply ->(t) { t.cuda(device) }
+      # TODO add device
+      def cuda
+        _apply ->(t) { t.cuda }
       end
 
       def cpu
@@ -112,8 +113,28 @@ module Torch
         destination
       end
 
+      # TODO add strict option
+      # TODO match PyTorch behavior
       def load_state_dict(state_dict)
-        raise NotImplementedYet
+        state_dict.each do |k, input_param|
+          k1, k2 = k.split(".", 2)
+          mod = named_modules[k1]
+          if mod.is_a?(Module)
+            param = mod.named_parameters[k2]
+            if param.is_a?(Parameter)
+              Torch.no_grad do
+                param.copy!(input_param)
+              end
+            else
+              raise Error, "Unknown parameter: #{k1}"
+            end
+          else
+            raise Error, "Unknown module: #{k1}"
+          end
+        end
+
+        # TODO return missing keys and unexpected keys
+        nil
       end
 
       def parameters
@@ -165,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)
@@ -203,13 +238,15 @@ module Torch
 
       def inspect
         name = self.class.name.split("::").last
-        if children.empty?
+        if named_children.empty?
           "#{name}(#{extra_inspect})"
         else
           str = String.new
           str << "#{name}(\n"
-          children.each do |name, mod|
-            str << "  (#{name}): #{mod.inspect}\n"
+          named_children.each do |name, mod|
+            mod_str = mod.inspect
+            mod_str = mod_str.lines.join("  ")
+            str << "  (#{name}): #{mod_str}\n"
           end
           str << ")"
         end

data/lib/torch/tensor.rb

@@ -4,6 +4,8 @@ module Torch
     include Inspector
 
     alias_method :requires_grad?, :requires_grad
+    alias_method :ndim, :dim
+    alias_method :ndimension, :dim
 
     def self.new(*args)
       FloatTensor.new(*args)
@@ -23,8 +25,17 @@ module Torch
       inspect
     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
@@ -37,6 +48,10 @@ module Torch
       to("cpu")
     end
 
+    def cuda
+      to("cuda")
+    end
+
     def size(dim = nil)
       if dim
         _size_int(dim)
@@ -58,7 +73,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
@@ -74,7 +97,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)
@@ -102,15 +125,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
@@ -139,6 +153,7 @@ module Torch
       neg
     end
 
+    # TODO better compare?
     def <=>(other)
       item <=> other
     end
@@ -186,8 +201,27 @@ module Torch
       end
     end
 
-    def random!(from = 0, to)
-      _random__from_to(from, to)
+    # 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
+      when 1
+        _random__to(*args)
+      when 2
+        _random__from_to(*args)
+      else
+        _random_(*args)
+      end
     end
 
     private
@@ -195,17 +229,5 @@ module Torch
     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
-    end
   end
 end