torch-rb 0.1.0

data/ext/torch/extconf.rb

@@ -0,0 +1,22 @@
+require "mkmf-rice"
+
+abort "Missing stdc++" unless have_library("stdc++")
+
+$CXXFLAGS << " -std=c++11"
+
+# silence ruby/intern.h warning
+$CXXFLAGS << " -Wno-deprecated-register"
+
+inc, lib = dir_config("torch")
+
+inc ||= "/usr/local/include"
+lib ||= "/usr/local/lib"
+
+$INCFLAGS << " -I#{inc}"
+$INCFLAGS << " -I#{inc}/torch/csrc/api/include"
+
+$LDFLAGS << " -Wl,-rpath,#{lib}"
+$LDFLAGS << " -L#{lib}"
+$LDFLAGS << " -ltorch -lc10"
+
+create_makefile("torch/ext")

data/lib/torch-rb.rb

@@ -0,0 +1 @@
+require "torch"

data/lib/torch.rb

@@ -0,0 +1,327 @@
+# ext
+require "torch/ext"
+
+# modules
+require "torch/inspector"
+require "torch/tensor"
+require "torch/version"
+
+# nn
+require "torch/nn/module"
+require "torch/nn/init"
+require "torch/nn/conv2d"
+require "torch/nn/functional"
+require "torch/nn/linear"
+require "torch/nn/parameter"
+require "torch/nn/sequential"
+require "torch/nn/relu"
+require "torch/nn/mse_loss"
+
+# utils
+require "torch/utils/data/data_loader"
+require "torch/utils/data/tensor_dataset"
+
+module Torch
+  class Error < StandardError; end
+
+  # keys: https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.dtype
+  # values: https://github.com/pytorch/pytorch/blob/master/c10/core/ScalarType.h
+  # complex and quantized types not supported by PyTorch yet
+  DTYPE_TO_ENUM = {
+    uint8: 0,
+    int8: 1,
+    short: 2,
+    int16: 2,
+    int: 3,
+    int32: 3,
+    long: 4,
+    int64: 4,
+    half: 5,
+    float16: 5,
+    float: 6,
+    float32: 6,
+    double: 7,
+    float64: 7,
+    # complex_half: 8,
+    # complex_float: 9,
+    # complex_double: 10,
+    bool: 11,
+    # qint8: 12,
+    # quint8: 13,
+    # qint32: 14,
+    # bfloat16: 15
+  }
+  ENUM_TO_DTYPE = DTYPE_TO_ENUM.map(&:reverse).to_h
+
+  class << self
+    # Torch.float, Torch.long, etc
+    DTYPE_TO_ENUM.each_key do |type|
+      define_method(type) do
+        type
+      end
+    end
+
+    # https://pytorch.org/docs/stable/torch.html
+
+    def tensor?(obj)
+      obj.is_a?(Tensor)
+    end
+
+    # TODO don't copy
+    def from_numo(ndarray)
+      dtype = _dtype_to_numo.find { |k, v| ndarray.is_a?(v) }
+      raise Error, "Cannot convert #{ndarray.class.name} to tensor" unless dtype
+      tensor(ndarray.to_a, dtype: dtype[0])
+    end
+
+    # private
+    # use method for cases when Numo not available
+    # or available after Torch loaded
+    def _dtype_to_numo
+      {
+        uint8: Numo::UInt8,
+        int8: Numo::Int8,
+        int16: Numo::Int16,
+        int32: Numo::Int32,
+        int64: Numo::Int64,
+        float32: Numo::SFloat,
+        float64: Numo::DFloat
+      }
+    end
+
+    # --- begin tensor creation: https://pytorch.org/cppdocs/notes/tensor_creation.html ---
+
+    def arange(start, finish = nil, step = 1, **options)
+      # ruby doesn't support start = 0, finish, step = 1, ...
+      if finish.nil?
+        finish = start
+        start = 0
+      end
+      _arange(start, finish, step, tensor_options(**options))
+    end
+
+    def empty(*size, **options)
+      _empty(tensor_size(size), tensor_options(**options))
+    end
+
+    def eye(n, m = nil, **options)
+      _eye(n, m || n, tensor_options(**options))
+    end
+
+    def full(size, fill_value, **options)
+      _full(size, fill_value, tensor_options(**options))
+    end
+
+    def linspace(start, finish, steps = 100, **options)
+      _linspace(start, finish, steps, tensor_options(**options))
+    end
+
+    def logspace(start, finish, steps = 100, base = 10.0, **options)
+      _logspace(start, finish, steps, base, tensor_options(**options))
+    end
+
+    def ones(*size, **options)
+      _ones(tensor_size(size), tensor_options(**options))
+    end
+
+    def rand(*size, **options)
+      _rand(tensor_size(size), tensor_options(**options))
+    end
+
+    def randint(low = 0, high, size, **options)
+      _randint(low, high, size, tensor_options(**options))
+    end
+
+     def randn(*size, **options)
+      _randn(tensor_size(size), tensor_options(**options))
+    end
+
+    def randperm(n, **options)
+      _randperm(n, tensor_options(**options))
+    end
+
+    def zeros(*size, **options)
+      _zeros(tensor_size(size), tensor_options(**options))
+    end
+
+    def tensor(data, **options)
+      size = []
+      if data.respond_to?(:to_a)
+        data = data.to_a
+        d = data
+        while d.is_a?(Array)
+          size << d.size
+          d = d.first
+        end
+        data = data.flatten
+      else
+        data = [data].compact
+      end
+
+      if options[:dtype].nil? && data.all? { |v| v.is_a?(Integer) }
+        options[:dtype] = :int64
+      end
+
+      _tensor(data, size, tensor_options(**options))
+    end
+
+    # --- begin like ---
+
+    def ones_like(input, **options)
+      ones(input.size, like_options(input, options))
+    end
+
+    def empty_like(input, **options)
+      empty(input.size, like_options(input, options))
+    end
+
+    def full_like(input, fill_value, **options)
+      full(input.size, fill_value, like_options(input, options))
+    end
+
+    def rand_like(input, **options)
+      rand(input.size, like_options(input, options))
+    end
+
+    def randint_like(input, low, high = nil, **options)
+      # ruby doesn't support input, low = 0, high, ...
+      if high.nil?
+        high = low
+        low = 0
+      end
+      rand(input.size, like_options(input, options))
+    end
+
+    def randn_like(input, **options)
+      randn(input.size, like_options(input, options))
+    end
+
+    def zeros_like(input, **options)
+      zeros(input.size, like_options(input, options))
+    end
+
+    # --- begin operations ---
+
+    %w(add sub mul div remainder).each do |op|
+      define_method(op) do |input, other, **options|
+        execute_op(op, input, other, **options)
+      end
+    end
+
+    def neg(input)
+      _neg(input)
+    end
+
+    def no_grad
+      previous_value = grad_enabled?
+      begin
+        _set_grad_enabled(false)
+        yield
+      ensure
+        _set_grad_enabled(previous_value)
+      end
+    end
+
+    # TODO support out
+    def mean(input, dim = nil, keepdim: false)
+      if dim
+        _mean_dim(input, dim, keepdim)
+      else
+        _mean(input)
+      end
+    end
+
+    # TODO support dtype
+    def sum(input, dim = nil, keepdim: false)
+      if dim
+        _sum_dim(input, dim, keepdim)
+      else
+        _sum(input)
+      end
+    end
+
+    def norm(input)
+      _norm(input)
+    end
+
+    def pow(input, exponent)
+      _pow(input, exponent)
+    end
+
+    def min(input)
+      _min(input)
+    end
+
+    def max(input)
+      _max(input)
+    end
+
+    def exp(input)
+      _exp(input)
+    end
+
+    def log(input)
+      _log(input)
+    end
+
+    def unsqueeze(input, dim)
+      _unsqueeze(input, dim)
+    end
+
+    def dot(input, tensor)
+      _dot(input, tensor)
+    end
+
+    def matmul(input, other)
+      _matmul(input, other)
+    end
+
+    private
+
+    def execute_op(op, input, other, out: nil)
+      scalar = other.is_a?(Numeric)
+      if out
+        # TODO make work with scalars
+        raise Error, "out not supported with scalar yet" if scalar
+        send("_#{op}_out", out, input, other)
+      else
+        if scalar
+          send("_#{op}_scalar", input, other)
+        else
+          send("_#{op}", input, other)
+        end
+      end
+    end
+
+    def tensor_size(size)
+      size.flatten
+    end
+
+    def tensor_options(dtype: nil, layout: nil, device: nil, requires_grad: nil)
+      options = TensorOptions.new
+      unless dtype.nil?
+        type = DTYPE_TO_ENUM[dtype]
+        raise Error, "Unknown dtype: #{dtype.inspect}" unless type
+        options = options.dtype(type)
+      end
+      unless device.nil?
+        options = options.device(device.to_s)
+      end
+      unless layout.nil?
+        options = options.layout(layout.to_s)
+      end
+      unless requires_grad.nil?
+        options = options.requires_grad(requires_grad)
+      end
+      options
+    end
+
+    def like_options(input, options)
+      options = options.dup
+      options[:dtype] ||= input.dtype
+      options[:layout] ||= input.layout
+      options[:device] ||= input.device
+      options
+    end
+  end
+end

data/lib/torch/inspector.rb

@@ -0,0 +1,62 @@
+module Torch
+  module Inspector
+    def inspect
+      data =
+        if numel == 0
+          "[]"
+        elsif dim == 0
+          to_a.first
+        else
+          values = to_a.flatten
+          abs = values.select { |v| v != 0 }.map(&:abs)
+          max = abs.max || 1
+          min = abs.min || 1
+
+          total = 0
+          if values.any? { |v| v < 0 }
+            total += 1
+          end
+
+          if floating_point?
+            sci = max / min.to_f > 1000 || max > 1e8 || min < 1e-4
+
+            all_int = values.all? { |v| v == v.to_i }
+            decimal = all_int ? 1 : 4
+
+            total += sci ? 10 : decimal + 1 + max.to_i.to_s.size
+
+            if sci
+              fmt = "%#{total}.4e"
+            else
+              fmt = "%#{total}.#{decimal}f"
+            end
+          else
+            total += max.to_s.size
+            fmt = "%#{total}d"
+          end
+
+          inspect_level(to_a, fmt, dim - 1)
+        end
+
+      attributes = []
+      if requires_grad
+        attributes << "requires_grad: true"
+      end
+      if ![:float32, :int64, :bool].include?(dtype)
+        attributes << "dtype: #{dtype.inspect}"
+      end
+
+      "tensor(#{data}#{attributes.map { |a| ", #{a}" }.join("")})"
+    end
+
+    private
+
+    def inspect_level(arr, fmt, total, level = 0)
+      if level == total
+        "[#{arr.map { |v| fmt % v }.join(", ")}]"
+      else
+        "[#{arr.map { |row| inspect_level(row, fmt, total, level + 1) }.join(",#{"\n" * (total - level)}#{" " * (level + 8)}")}]"
+      end
+    end
+  end
+end

data/lib/torch/nn/conv2d.rb

@@ -0,0 +1,50 @@
+module Torch
+  module NN
+    class Conv2d < Module
+      attr_reader :bias, :weight
+
+      def initialize(in_channels, out_channels, kernel_size) #, stride: 1, padding: 0, dilation: 1, groups: 1)
+        @in_channels = in_channels
+        @out_channels = out_channels
+        @kernel_size = pair(kernel_size)
+        @stride = pair(1)
+        # @stride = pair(stride)
+        # @padding = pair(padding)
+        # @dilation = pair(dilation)
+
+        # TODO divide by groups
+        @weight = Parameter.new(Tensor.new(out_channels, in_channels, *@kernel_size))
+        @bias = Parameter.new(Tensor.new(out_channels))
+
+        reset_parameters
+      end
+
+      def reset_parameters
+        Init.kaiming_uniform_(@weight, Math.sqrt(5))
+        if @bias
+          fan_in, _ = Init.calculate_fan_in_and_fan_out(@weight)
+          bound = 1 / Math.sqrt(fan_in)
+          Init.uniform_(@bias, -bound, bound)
+        end
+      end
+
+      def call(input)
+        F.conv2d(input, @weight, @bias) # @stride, self.padding, self.dilation, self.groups)
+      end
+
+      def inspect
+        "Conv2d(#{@in_channels}, #{@out_channels}, kernel_size: #{@kernel_size.inspect}, stride: #{@stride.inspect})"
+      end
+
+      private
+
+      def pair(value)
+        if value.is_a?(Array)
+          value
+        else
+          [value] * 2
+        end
+      end
+    end
+  end
+end