torch-rb 0.1.3 → 0.1.4

data/lib/torch.rb

@@ -22,31 +22,90 @@ require "torch/optim/sgd"
 require "torch/optim/lr_scheduler/lr_scheduler"
 require "torch/optim/lr_scheduler/step_lr"
 
-# nn base classes
+# nn parameters
+require "torch/nn/parameter"
+
+# nn containers
 require "torch/nn/module"
+require "torch/nn/sequential"
+
+# nn convolution layers
 require "torch/nn/convnd"
-require "torch/nn/dropoutnd"
+require "torch/nn/conv2d"
+
+# nn pooling layers
+require "torch/nn/max_poolnd"
+require "torch/nn/max_pool2d"
+require "torch/nn/avg_poolnd"
+require "torch/nn/avg_pool2d"
+
+# nn linear layers
+require "torch/nn/bilinear"
+require "torch/nn/identity"
+require "torch/nn/linear"
 
-# nn
+# nn dropout layers
+require "torch/nn/dropoutnd"
 require "torch/nn/alpha_dropout"
-require "torch/nn/conv2d"
 require "torch/nn/dropout"
 require "torch/nn/dropout2d"
 require "torch/nn/dropout3d"
-require "torch/nn/embedding"
 require "torch/nn/feature_alpha_dropout"
+
+# nn activations
+require "torch/nn/leaky_relu"
+require "torch/nn/prelu"
+require "torch/nn/relu"
+require "torch/nn/sigmoid"
+require "torch/nn/softplus"
+
+# nn activations other
+require "torch/nn/log_softmax"
+require "torch/nn/softmax"
+require "torch/nn/softmax2d"
+require "torch/nn/softmin"
+
+# nn sparse layers
+require "torch/nn/embedding"
+require "torch/nn/embedding_bag"
+
+# nn distance functions
+require "torch/nn/cosine_similarity"
+require "torch/nn/pairwise_distance"
+
+# nn loss functions
+require "torch/nn/loss"
+require "torch/nn/weighted_loss"
+require "torch/nn/bce_loss"
+# require "torch/nn/bce_with_logits_loss"
+# require "torch/nn/cosine_embedding_loss"
+require "torch/nn/cross_entropy_loss"
+require "torch/nn/ctc_loss"
+# require "torch/nn/hinge_embedding_loss"
+require "torch/nn/kl_div_loss"
+require "torch/nn/l1_loss"
+# require "torch/nn/margin_ranking_loss"
+require "torch/nn/mse_loss"
+# require "torch/nn/multi_label_margin_loss"
+# require "torch/nn/multi_label_soft_margin_loss"
+# require "torch/nn/multi_margin_loss"
+require "torch/nn/nll_loss"
+require "torch/nn/poisson_nll_loss"
+# require "torch/nn/smooth_l1_loss"
+# require "torch/nn/soft_margin_loss"
+# require "torch/nn/triplet_margin_loss"
+
+# nn other
 require "torch/nn/functional"
 require "torch/nn/init"
-require "torch/nn/linear"
-require "torch/nn/mse_loss"
-require "torch/nn/parameter"
-require "torch/nn/relu"
-require "torch/nn/sequential"
 
 # utils
 require "torch/utils/data/data_loader"
 require "torch/utils/data/tensor_dataset"
 
+# random
+require "torch/random"
+
 module Torch
   class Error < StandardError; end
   class NotImplementedYet < StandardError
@@ -57,7 +116,6 @@ module Torch
 
   # 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,
@@ -73,14 +131,14 @@ module Torch
     float32: 6,
     double: 7,
     float64: 7,
-    # complex_half: 8,
-    # complex_float: 9,
-    # complex_double: 10,
+    complex_half: 8,
+    complex_float: 9,
+    complex_double: 10,
     bool: 11,
-    # qint8: 12,
-    # quint8: 13,
-    # qint32: 14,
-    # bfloat16: 15
+    qint8: 12,
+    quint8: 13,
+    qint32: 14,
+    bfloat16: 15
   }
   ENUM_TO_DTYPE = DTYPE_TO_ENUM.map(&:reverse).to_h
 
@@ -120,6 +178,8 @@ module Torch
     # use method for cases when Numo not available
     # or available after Torch loaded
     def _dtype_to_numo
+      raise Error, "Numo not found" unless defined?(Numo::NArray)
+
       {
         uint8: Numo::UInt8,
         int8: Numo::Int8,
@@ -200,8 +260,12 @@ module Torch
         data = [data].compact
       end
 
-      if options[:dtype].nil? && data.all? { |v| v.is_a?(Integer) }
-        options[:dtype] = :int64
+      if options[:dtype].nil?
+        if data.all? { |v| v.is_a?(Integer) }
+          options[:dtype] = :int64
+        elsif data.all? { |v| v == true || v == false }
+          options[:dtype] = :bool
+        end
       end
 
       _tensor(data, size, tensor_options(**options))
@@ -302,6 +366,10 @@ module Torch
       _pow(input, exponent)
     end
 
+    def topk(input, k)
+      _topk(input, k)
+    end
+
     def min(input)
       _min(input)
     end
@@ -327,6 +395,10 @@ module Torch
       _sign(input)
     end
 
+    def sigmoid(input)
+      _sigmoid(input)
+    end
+
     def gt(input, other)
       _gt(input, other)
     end
@@ -363,6 +435,15 @@ module Torch
       _sqrt(input)
     end
 
+    # TODO make dim keyword argument
+    def log_softmax(input, dim)
+      _log_softmax(input, dim)
+    end
+
+    def softmax(input, dim: nil)
+      _softmax(input, dim)
+    end
+
     def abs(input)
       _abs(input)
     end

data/lib/torch/inspector.rb

@@ -11,32 +11,36 @@ module Torch
         else
           summarize = numel > 1000
 
-          values = to_a.flatten
-          abs = values.select { |v| v != 0 }.map(&:abs)
-          max = abs.max || 1
-          min = abs.min || 1
+          if dtype == :bool
+            fmt = "%s"
+          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
+            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
+            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
+              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
+              total += sci ? 10 : decimal + 1 + max.to_i.to_s.size
 
-            if sci
-              fmt = "%#{total}.4e"
+              if sci
+                fmt = "%#{total}.4e"
+              else
+                fmt = "%#{total}.#{decimal}f"
+              end
             else
-              fmt = "%#{total}.#{decimal}f"
+              total += max.to_s.size
+              fmt = "%#{total}d"
             end
-          else
-            total += max.to_s.size
-            fmt = "%#{total}d"
           end
 
           inspect_level(to_a, fmt, dim - 1, 0, summarize)

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/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

@@ -19,8 +19,8 @@ module Torch
       end
 
       # TODO add more parameters
-      def inspect
-        "Conv2d(#{@in_channels}, #{@out_channels}, kernel_size: #{@kernel_size.inspect}, stride: #{@stride.inspect})"
+      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

@@ -29,11 +29,11 @@ module Torch
       end
 
       def reset_parameters
-        Init.kaiming_uniform!(@weight, Math.sqrt(5))
+        Init.kaiming_uniform!(@weight, a: Math.sqrt(5))
         if @bias
-          fan_in, _ = Init.calculate_fan_in_and_fan_out(@weight)
+          fan_in, _ = Init._calculate_fan_in_and_fan_out(@weight)
           bound = 1 / Math.sqrt(fan_in)
-          Init.uniform!(@bias, -bound, bound)
+          Init.uniform!(@bias, a: -bound, b: bound)
         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/dropoutnd.rb

@@ -7,8 +7,8 @@ module Torch
         @inplace = inplace
       end
 
-      def inspect
-        "#{self.class.name.split("::").last}(p: #{@p.inspect}, inplace: #{@inplace.inspect})"
+      def extra_inspect
+        format("p: %s, inplace: %s", @p, @inplace)
       end
     end
   end

data/lib/torch/nn/embedding_bag.rb

@@ -0,0 +1,34 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/sparse.py
+module Torch
+  module NN
+    class EmbeddingBag < Module
+      def initialize(num_embeddings, embedding_dim, max_norm: nil, norm_type: 2.0,
+        scale_grad_by_freq: false, mode: "mean", sparse: false, _weight: nil)
+
+        super()
+        @num_embeddings = num_embeddings
+        @embedding_dim = embedding_dim
+        @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
+        @mode = mode
+        @sparse = sparse
+      end
+
+      def reset_parameters
+        Init.normal!(@weight)
+      end
+
+      def forward(input, offsets: nil, per_sample_weights: nil)
+        F.embedding_bag(input, @weight, offsets: offsets, max_norm: @max_norm, norm_type: @norm_type, scale_grad_by_freq: @scale_grad_by_freq, mode: @mode, sparse: @sparse, per_sample_weights: per_sample_weights)
+      end
+    end
+  end
+end