torch-rb 0.1.3

data/lib/torch/inspector.rb

@@ -0,0 +1,85 @@
+module Torch
+  module Inspector
+    # TODO make more performance, especially when summarizing
+    # how? only read data that will be displayed
+    def inspect
+      data =
+        if numel == 0
+          "[]"
+        elsif dim == 0
+          item
+        else
+          summarize = numel > 1000
+
+          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.finite? && 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, 0, summarize)
+        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
+
+    # TODO DRY code
+    def inspect_level(arr, fmt, total, level, summarize)
+      if level == total
+        cols =
+          if summarize && arr.size > 7
+            arr[0..2].map { |v| fmt % v } +
+            ["..."] +
+            arr[-3..-1].map { |v| fmt % v }
+          else
+            arr.map { |v| fmt % v }
+          end
+
+        "[#{cols.join(", ")}]"
+      else
+        rows =
+          if summarize && arr.size > 7
+            arr[0..2].map { |row| inspect_level(row, fmt, total, level + 1, summarize) } +
+            ["..."] +
+            arr[-3..-1].map { |row| inspect_level(row, fmt, total, level + 1, summarize) }
+          else
+            arr.map { |row| inspect_level(row, fmt, total, level + 1, summarize) }
+          end
+
+        "[#{rows.join(",#{"\n" * (total - level)}#{" " * (level + 8)}")}]"
+      end
+    end
+  end
+end

data/lib/torch/nn/alpha_dropout.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class AlphaDropout < DropoutNd
+      def forward(input)
+        F.alpha_dropout(input, p: @p, training: @training, inplace: @inplace)
+      end
+    end
+  end
+end

data/lib/torch/nn/conv2d.rb

@@ -0,0 +1,37 @@
+module Torch
+  module NN
+    class Conv2d < ConvNd
+      attr_reader :bias, :weight
+
+      def initialize(in_channels, out_channels, kernel_size, stride: 1, padding: 0, dilation: 1, groups: 1, bias: true, padding_mode: "zeros")
+        kernel_size = pair(kernel_size)
+        stride = pair(stride)
+        padding = pair(padding)
+        dilation = pair(dilation)
+        super(in_channels, out_channels, kernel_size, stride, padding, dilation, false, pair(0), groups, bias, padding_mode)
+      end
+
+      def forward(input)
+        if @padding_mode == "circular"
+          raise NotImplementedError
+        end
+        F.conv2d(input, @weight, @bias, stride: @stride, padding: @padding, dilation: @dilation, groups: @groups)
+      end
+
+      # TODO add more parameters
+      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

data/lib/torch/nn/convnd.rb

@@ -0,0 +1,41 @@
+module Torch
+  module NN
+    class ConvNd < Module
+      def initialize(in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding, groups, bias, padding_mode)
+        super()
+        raise ArgumentError, "in_channels must be divisible by groups" if in_channels % groups != 0
+        raise ArgumentError, "out_channels must be divisible by groups" if out_channels % groups != 0
+        @in_channels = in_channels
+        @out_channels = out_channels
+        @kernel_size = kernel_size
+        @stride = stride
+        @padding = padding
+        @dilation = dilation
+        @transposed = transposed
+        @output_padding = output_padding
+        @groups = groups
+        @padding_mode = padding_mode
+        if transposed
+          @weight = Parameter.new(Tensor.new(in_channels, out_channels / groups, *kernel_size))
+        else
+          @weight = Parameter.new(Tensor.new(out_channels, in_channels / groups, *kernel_size))
+        end
+        if bias
+          @bias = Parameter.new(Tensor.new(out_channels))
+        else
+          raise NotImplementedError
+        end
+        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
+    end
+  end
+end

data/lib/torch/nn/dropout.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class Dropout < DropoutNd
+      def forward(input)
+        F.dropout(input, p: @p, training: @training, inplace: @inplace)
+      end
+    end
+  end
+end

data/lib/torch/nn/dropout2d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class Dropout2d < DropoutNd
+      def forward(input)
+        F.dropout2d(input, p: @p, training: @training, inplace: @inplace)
+      end
+    end
+  end
+end

data/lib/torch/nn/dropout3d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class Dropout3d < DropoutNd
+      def forward(input)
+        F.dropout3d(input, p: @p, training: @training, inplace: @inplace)
+      end
+    end
+  end
+end

data/lib/torch/nn/dropoutnd.rb

@@ -0,0 +1,15 @@
+module Torch
+  module NN
+    class DropoutNd < Module
+      def initialize(p: 0.5, inplace: false)
+        super()
+        @p = p
+        @inplace = inplace
+      end
+
+      def inspect
+        "#{self.class.name.split("::").last}(p: #{@p.inspect}, inplace: #{@inplace.inspect})"
+      end
+    end
+  end
+end

data/lib/torch/nn/embedding.rb

@@ -0,0 +1,52 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/sparse.py
+module Torch
+  module NN
+    class Embedding < Module
+      def initialize(num_embeddings, embedding_dim, padding_idx: nil, max_norm: nil,
+        norm_type: 2.0, scale_grad_by_freq: false, sparse: false, _weight: nil)
+
+        super()
+        @num_embeddings = num_embeddings
+        @embedding_dim = embedding_dim
+
+        if padding_idx
+          if padding_idx > 0
+            raise ArgumentError, "Padding_idx must be within num_embeddings" unless padding_idx < @num_embeddings
+          elsif padding_idx < 0
+            raise ArgumentError, "Padding_idx must be within num_embeddings" unless padding_idx >= -@num_embeddings
+            padding_idx = @num_embeddings + padding_idx
+          end
+        end
+        @padding_idx = padding_idx
+        @max_norm = max_norm
+        @norm_type = norm_type
+        @scale_grad_by_freq = scale_grad_by_freq
+        if _weight.nil?
+          @weight = Parameter.new(Tensor.new(num_embeddings, embedding_dim))
+          reset_parameters
+        else
+          raise ArgumentError, "Shape of weight does not match num_embeddings and embedding_dim" unless _weight.shape == [num_embeddings, embedding_dim]
+          @weight = Parameter.new(_weight)
+        end
+        @sparse = sparse
+      end
+
+      def reset_parameters
+        Init.normal!(@weight)
+        if @padding_idx
+          Torch.no_grad do
+            @weight[@padding_idx].fill!(0)
+          end
+        end
+      end
+
+      def forward(input)
+        F.embedding(input, @weight, padding_idx: @padding_idx, max_norm: @max_norm, norm_type: @norm_type, scale_grad_by_freq: @scale_grad_by_freq, sparse: @sparse)
+      end
+
+      def inspect
+        "Embedding(#{@num_embeddings}, #{@embedding_dim})"
+      end
+    end
+  end
+end

data/lib/torch/nn/feature_alpha_dropout.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class FeatureAlphaDropout < DropoutNd
+      def forward(input)
+        F.feature_alpha_dropout(input, p: @p, training: @training, inplace: @inplace)
+      end
+    end
+  end
+end

data/lib/torch/nn/functional.rb

@@ -0,0 +1,100 @@
+module Torch
+  module NN
+    class Functional
+      class << self
+        def relu(input)
+          Torch.relu(input)
+        end
+
+        def conv2d(input, weight, bias, stride: 1, padding: 0, dilation: 1, groups: 1)
+          # TODO pair stride and padding when needed
+          Torch.conv2d(input, weight, bias, stride, padding, dilation, groups)
+        end
+
+        def max_pool2d(input, kernel_size)
+          kernel_size = [kernel_size, kernel_size] if kernel_size.is_a?(Integer)
+          Torch.max_pool2d(input, kernel_size)
+        end
+
+        def avg_pool2d(input, kernel_size)
+          kernel_size = [kernel_size, kernel_size] if kernel_size.is_a?(Integer)
+          Torch.avg_pool2d(input, kernel_size)
+        end
+
+        def linear(input, weight, bias)
+          Torch.linear(input, weight, bias)
+        end
+
+        def mse_loss(input, target, reduction: "mean")
+          Torch.mse_loss(input, target, reduction)
+        end
+
+        def cross_entropy(input, target)
+          nll_loss(log_softmax(input, 1), target)
+        end
+
+        def nll_loss(input, target, reduction: "mean")
+          # TODO fix for non-1d
+          Torch.nll_loss(input, target, reduction)
+        end
+
+        def log_softmax(input, dim)
+          input.log_softmax(dim)
+        end
+
+        def dropout(input, p: 0.5, training: true, inplace: false)
+          if inplace
+            Torch._dropout!(input, p, training)
+          else
+            Torch._dropout(input, p, training)
+          end
+        end
+
+        def dropout2d(input, p: 0.5, training: true, inplace: false)
+          raise ArgumentError, "dropout probability has to be between 0 and 1, but got #{p}" if p < 0 || p > 1
+
+          if inplace
+            Torch._feature_dropout!(input, p, training)
+          else
+            Torch._feature_dropout(input, p, training)
+          end
+        end
+
+        def dropout3d(input, p: 0.5, training: true, inplace: false)
+          if inplace
+            Torch._feature_dropout!(input, p, training)
+          else
+            Torch._feature_dropout(input, p, training)
+          end
+        end
+
+        def alpha_dropout(input, p: 0.5, training: true, inplace: false)
+          if inplace
+            Torch._alpha_dropout!(input, p, training)
+          else
+            Torch._alpha_dropout(input, p, training)
+          end
+        end
+
+        def feature_alpha_dropout(input, p: 0.5, training: true, inplace: false)
+          if inplace
+            Torch._feature_alpha_dropout!(input, p, training)
+          else
+            Torch._feature_alpha_dropout(input, p, training)
+          end
+        end
+
+        def embedding(input, weight, padding_idx: nil, max_norm: nil, norm_type: 2.0, scale_grad_by_freq: false, sparse: false)
+          # TODO handle max_norm and norm_type
+          raise NotImplementedYet unless max_norm.nil? && norm_type == 2.0
+
+          padding_idx ||= -1
+          Torch._embedding(input, weight, padding_idx, scale_grad_by_freq, sparse)
+        end
+      end
+    end
+
+    # shortcut
+    F = Functional
+  end
+end

data/lib/torch/nn/init.rb

@@ -0,0 +1,30 @@
+module Torch
+  module NN
+    module Init
+      class << self
+        def calculate_fan_in_and_fan_out(tensor)
+          dimensions = tensor.dim
+          if dimensions < 2
+            raise Error, "Fan in and fan out can not be computed for tensor with fewer than 2 dimensions"
+          end
+
+          if dimensions == 2
+            fan_in = tensor.size(1)
+            fan_out = tensor.size(0)
+          else
+            num_input_fmaps = tensor.size(1)
+            num_output_fmaps = tensor.size(0)
+            receptive_field_size = 1
+            if tensor.dim > 2
+              receptive_field_size = tensor[0][0].numel
+            end
+            fan_in = num_input_fmaps * receptive_field_size
+            fan_out = num_output_fmaps * receptive_field_size
+          end
+
+          [fan_in, fan_out]
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/linear.rb

@@ -0,0 +1,36 @@
+module Torch
+  module NN
+    class Linear < Module
+      attr_reader :bias, :weight
+
+      def initialize(in_features, out_features, bias: true)
+        @in_features = in_features
+        @out_features = out_features
+
+        @weight = Parameter.new(Tensor.new(out_features, in_features))
+        if bias
+          @bias = Parameter.new(Tensor.new(out_features))
+        end
+
+        reset_parameters
+      end
+
+      def call(input)
+        F.linear(input, @weight, @bias)
+      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 inspect
+        "Linear(in_features: #{@in_features.inspect}, out_features: #{@out_features.inspect}, bias: #{(!@bias.nil?).inspect})"
+      end
+    end
+  end
+end