torch-rb 0.1.0 → 0.1.5

data/lib/torch/native/parser.rb

@@ -0,0 +1,97 @@
+module Torch
+  module Native
+    class Parser
+      def initialize(functions)
+        @functions = functions
+        @name = @functions.first.ruby_name
+        @min_args = @functions.map { |f| f.args.count { |a| a[:pos] && a[:default].nil? } }.min
+        @max_args = @functions.map { |f| f.args.count { |a| a[:pos] } }.max
+      end
+
+      def parse(args, options)
+        candidates = @functions.dup
+
+        if args.size < @min_args || args.size > @max_args
+          expected = String.new(@min_args.to_s)
+          expected += "..#{@max_args}" if @max_args != @min_args
+          return {error: "wrong number of arguments (given #{args.size}, expected #{expected})"}
+        end
+
+        # exclude functions where options don't match
+        options.each do |k, v|
+          candidates.select! do |func|
+            func.args.any? { |a| a[:name] == k.to_s }
+          end
+          # TODO show all bad keywords at once like Ruby?
+          return {error: "unknown keyword: #{k}"} if candidates.empty?
+        end
+
+        # exclude functions missing required options
+        candidates.reject! do |func|
+          # TODO make more generic
+          func.out? && !options[:out]
+        end
+
+        final_values = {}
+
+        # check args
+        candidates.select! do |func|
+          good = true
+
+          values = args.zip(func.args).map { |a, fa| [fa[:name], a] }.to_h
+          values.merge!(options.map { |k, v| [k.to_s, v] }.to_h)
+          func.args.each do |fa|
+            values[fa[:name]] ||= fa[:default]
+          end
+
+          arg_types = func.args.map { |a| [a[:name], a[:type]] }.to_h
+
+          values.each do |k, v|
+            t = arg_types[k].split("(").first
+            good =
+              case t
+              when "Tensor"
+                v.is_a?(Tensor)
+              when "Tensor[]"
+                v.all? { |v2| v2.is_a?(Tensor) }
+              when "int"
+                v.is_a?(Integer)
+              when "int[]"
+                v.all? { |v2| v2.is_a?(Integer) }
+              when "Scalar"
+                v.is_a?(Numeric)
+              when "bool"
+                v == true || v == false
+              else
+                raise Error, "Unknown argument type: #{arg_types[k]}. Please report a bug with #{@name}"
+              end
+
+            if !good
+              if candidates.size == 1
+                k = "input" if k == "self"
+                return {error: "#{@name}(): argument '#{k}' must be #{t}"}
+              end
+              break
+            end
+          end
+
+          if good
+            final_values = values
+          end
+
+          good
+        end
+
+        if candidates.size != 1
+          raise Error, "This should never happen. Please report a bug with #{@name}."
+        end
+
+        func = candidates.first
+        {
+          name: func.cpp_name,
+          args: func.args.map { |a| final_values[a[:name]] }
+        }
+      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/avg_pool2d.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class AvgPool2d < AvgPoolNd
+      def initialize(kernel_size)
+        super()
+        @kernel_size = kernel_size
+      end
+
+      def forward(input)
+        F.avg_pool2d(input, @kernel_size)
+      end
+    end
+  end
+end

data/lib/torch/nn/avg_poolnd.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class AvgPoolNd < Module
+      def extra_inspect
+        format("kernel_size: %s", @kernel_size)
+      end
+    end
+  end
+end

data/lib/torch/nn/bce_loss.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class BCELoss < WeightedLoss
+      def initialize(weight: nil, reduction: "mean")
+        super(weight, reduction)
+      end
+
+      def forward(input, target)
+        F.binary_cross_entropy(input, target, weight: @weight, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/bce_with_logits_loss.rb

@@ -0,0 +1,15 @@
+module Torch
+  module NN
+    class BCEWithLogitsLoss < Loss
+      def initialize(weight: nil, reduction: "mean", pos_weight: nil)
+        super(reduction)
+        register_buffer("weight", weight)
+        register_buffer("pos_weight", pos_weight)
+      end
+
+      def forward(input, target)
+        F.binary_cross_entropy_with_logits(input, target, weight: weight, pos_weight: pos_weight, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/bilinear.rb

@@ -0,0 +1,38 @@
+module Torch
+  module NN
+    class Bilinear < Module
+      def initialize(in1_features, in2_features, out_features, bias: true)
+        super()
+
+        @in1_features = in1_features
+        @in2_features = in2_features
+        @out_features = out_features
+        @weight = Parameter.new(Tensor.new(out_features, in1_features, in2_features))
+
+        if bias
+          @bias = Parameter.new(Tensor.new(out_features))
+        else
+          raise NotImplementedYet
+        end
+
+        reset_parameters
+      end
+
+      def reset_parameters
+        bound = 1 / Math.sqrt(@weight.size(1))
+        Init.uniform!(@weight, a: -bound, b: bound)
+        if @bias
+          Init.uniform!(@bias, a: -bound, b: bound)
+        end
+      end
+
+      def forward(input1, input2)
+        F.bilinear(input1, input2, @weight, @bias)
+      end
+
+      def extra_inspect
+        format("in1_features: %s, in2_features: %s, out_features: %s, bias: %s", @in1_features, @in2_features, @out_features, !@bias.nil?)
+      end
+    end
+  end
+end

data/lib/torch/nn/conv2d.rb

@@ -1,39 +1,24 @@
 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
+    class Conv2d < ConvNd
+      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 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)
+      def forward(input)
+        if @padding_mode == "circular"
+          raise NotImplementedError
         end
+        F.conv2d(input, @weight, @bias, stride: @stride, padding: @padding, dilation: @dilation, groups: @groups)
       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})"
+      # TODO add more parameters
+      def extra_inspect
+        format("%s, %s, kernel_size: %s, stride: %s", @in_channels, @out_channels, @kernel_size, @stride)
       end
 
       private

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, a: Math.sqrt(5))
+        if @bias
+          fan_in, _ = Init._calculate_fan_in_and_fan_out(@weight)
+          bound = 1 / Math.sqrt(fan_in)
+          Init.uniform!(@bias, a: -bound, b: bound)
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/cosine_embedding_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class CosineEmbeddingLoss < Loss
+      def initialize(margin: 0, reduction: "mean")
+        super(reduction)
+        @margin = margin
+      end
+
+      def forward(input1, input2, target)
+        F.cosine_embedding_loss(input1, input2, target, margin: @margin, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/cosine_similarity.rb

@@ -0,0 +1,15 @@
+module Torch
+  module NN
+    class CosineSimilarity < Module
+      def initialize(dim: 1, eps: 1e-8)
+        super()
+        @dim = dim
+        @eps = eps
+      end
+
+      def forward(x1, x2)
+        F.cosine_similarity(x1, x2, dim: @dim, eps: @eps)
+      end
+    end
+  end
+end

data/lib/torch/nn/cross_entropy_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class CrossEntropyLoss < WeightedLoss
+      def initialize(weight: nil, ignore_index: -100, reduction: "mean")
+        super(weight, reduction)
+        @ignore_index = ignore_index
+      end
+
+      def forward(input, target)
+        F.cross_entropy(input, target, weight: @weight, ignore_index: @ignore_index, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/ctc_loss.rb

@@ -0,0 +1,15 @@
+module Torch
+  module NN
+    class CTCLoss < Loss
+      def initialize(blank: 0, reduction: "mean", zero_infinity: false)
+        super(reduction)
+        @blank = blank
+        @zero_infinity = zero_infinity
+      end
+
+      def forward(log_probs, targets, input_lengths, target_lengths)
+        F.ctc_loss(log_probs, targets, input_lengths, target_lengths, blank: @blank, reduction: @reduction, zero_infinity: @zero_infinity)
+      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 extra_inspect
+        format("p: %s, inplace: %s", @p, @inplace)
+      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