torch-rb 0.1.0 → 0.1.5

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

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

@@ -2,12 +2,25 @@ module Torch
   module NN
     class Functional
       class << self
-        def relu(input)
-          Torch.relu(input)
+        def relu(input, inplace: false)
+          if inplace
+            input.relu!
+          else
+            input.relu
+          end
         end
 
-        def conv2d(input, weight, bias)
-          Torch.conv2d(input, weight, bias)
+        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 prelu(input, weight)
+          Torch.prelu(input, weight)
+        end
+
+        def leaky_relu(input, negative_slope = 0.01)
+          Torch.leaky_relu(input, negative_slope)
         end
 
         def max_pool2d(input, kernel_size)
@@ -15,26 +28,196 @@ module Torch
           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
+
+        # linear layers
+
+        def bilinear(input1, input2, weight, bias)
+          Torch.bilinear(input1, input2, weight, bias)
+        end
+
         def linear(input, weight, bias)
           Torch.linear(input, weight, bias)
         end
 
+        # sparse layers
+
+        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
+          # weight and indices are swapped from Python interface
+          Torch._embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
+        end
+
+        def embedding_bag(input, weight, offsets: nil, max_norm: nil, norm_type: 2, scale_grad_by_freq: false, mode: "mean", sparse: false, per_sample_weights: nil)
+          # need to handle nils
+          raise NotImplementedYet
+
+          # TODO handle max_norm and norm_type
+          raise NotImplementedYet unless max_norm.nil? && norm_type == 2.0
+
+          Torch._embedding_bag(input, weight, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights)
+        end
+
+        # distance functions
+
+        def cosine_similarity(x1, x2, dim: 1, eps: 1e-8)
+          Torch._cosine_similarity(x1, x2, dim, eps)
+        end
+
+        def pairwise_distance(x1, x2, p: 2.0, eps: 1e-6, keepdim: false)
+          Torch._pairwise_distance(x1, x2, p, eps, keepdim)
+        end
+
+        # loss functions
+
+        def binary_cross_entropy(input, target, weight: nil, reduction: "mean")
+          NN._binary_cross_entropy(input, target, weight, reduction)
+        end
+
+        def binary_cross_entropy_with_logits(input, target, weight: nil, reduction: "mean", pos_weight: nil)
+          Torch._binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction)
+        end
+
+        def cosine_embedding_loss(input1, input2, target, margin: 0, reduction: "mean")
+          raise NotImplementedYet
+        end
+
+        def cross_entropy(input, target, weight: nil, ignore_index: -100, reduction: "mean")
+          nll_loss(log_softmax(input, 1), target, weight: weight, ignore_index: ignore_index, reduction: reduction)
+        end
+
+        def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank: 0, reduction: "mean", zero_infinity: false)
+          # call to_a on input_lengths and target_lengths for C++
+          Torch._ctc_loss_intlist(log_probs, targets, input_lengths.to_a, target_lengths.to_a, blank, reduction, zero_infinity)
+        end
+
+        def hinge_embedding_loss(input, target, margin: 1.0, reduction: "mean")
+          Torch._hinge_embedding_loss(input, target, margin, reduction)
+        end
+
+        def kl_div(input, target, reduction: "mean")
+          Torch._kl_div(input, target, reduction)
+        end
+
+        def l1_loss(input, target, reduction: "mean")
+          NN._l1_loss(input, target, reduction)
+        end
+
+        def margin_ranking_loss(input1, input2, target, margin: 0, reduction: "mean")
+          raise NotImplementedYet
+        end
+
         def mse_loss(input, target, reduction: "mean")
-          Torch.mse_loss(input, target, reduction)
+          NN._mse_loss(input, target, reduction)
+        end
+
+        def multilabel_margin_loss(input, target, reduction: "mean")
+          NN._multilabel_margin_loss(input, target, reduction)
+        end
+
+        def multilabel_soft_margin_loss(input, target, weight: nil)
+          raise NotImplementedYet
+        end
+
+        def multi_margin_loss(input, target, p: 1, margin: 1.0, weight: nil, reduction: "mean")
+          NN._multi_margin_loss(input, target, p, margin, weight, reduction)
         end
 
-        def cross_entropy(input, target)
-          nll_loss(log_softmax(input, 1), target)
+        def nll_loss(input, target, weight: nil, ignore_index: -100, reduction: "mean")
+          NN._nll_loss(input, target, weight, reduction, ignore_index)
         end
 
-        def nll_loss(input, target)
-          # TODO fix for non-1d
-          Torch.nll_loss(input, target)
+        def poisson_nll_loss(input, target, log_input: true, full: false, eps: 1e-8, reduction: "mean")
+          Torch._poisson_nll_loss(input, target, log_input, full, eps, reduction)
+        end
+
+        def soft_margin_loss(input, target, reduction: "mean")
+          NN._soft_margin_loss(input, target, reduction)
+        end
+
+        def smooth_l1_loss(input, target, reduction: "mean")
+          NN._smooth_l1_loss(input, target, reduction)
+        end
+
+        def triplet_margin_loss(anchor, positive, negative, margin: 1.0, p: 2, eps: 1e-06, swap: false, reduction: "mean")
+          Torch._triplet_margin_loss(anchor, positive, negative, margin, p, eps, swap, reduction)
+        end
+
+        # end loss
+
+        def softmax(input, dim: nil)
+          dim ||= softmax_dim(input.dim)
+          input.softmax(dim: dim)
         end
 
-        def log_softmax(input, dim)
+        def softmin(input, dim: nil)
+          dim ||= softmax_dim(input.dim)
+          (-input).softmax(dim: dim)
+        end
+
+        def softplus(input, beta: 1, threshold: 20)
+          NN._softplus(input, beta, threshold)
+        end
+
+        # TODO make dim keyword argument and update examples
+        def log_softmax(input, dim = nil)
+          dim ||= softmax_dim(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
+
+        private
+
+        def softmax_dim(ndim)
+          ndim == 0 || ndim == 1 || ndim == 3 ? 0 : 1
+        end
       end
     end
 

data/lib/torch/nn/hinge_embedding_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class HingeEmbeddingLoss < Loss
+      def initialize(margin: 1.0, reduction: "mean")
+        super(reduction)
+        @margin = margin
+      end
+
+      def forward(input, target)
+        F.hinge_embedding_loss(input, target, margin: @margin, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/identity.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class Identity < Module
+      # written this way to support unused arguments
+      def initialize(*args, **options)
+        super()
+      end
+
+      def forward(input)
+        input
+      end
+    end
+  end
+end

data/lib/torch/nn/init.rb

@@ -2,7 +2,64 @@ module Torch
   module NN
     module Init
       class << self
-        def calculate_fan_in_and_fan_out(tensor)
+        def calculate_gain(nonlinearity, param: 0.01)
+          _calculate_gain(nonlinearity, param)
+        end
+
+        def uniform!(tensor, a: 0.0, b: 1.0)
+          _uniform!(tensor, a, b)
+        end
+
+        def normal!(tensor, mean: 0.0, std: 1.0)
+          _normal!(tensor, mean, std)
+        end
+
+        def constant!(tensor, val)
+          _constant!(tensor, val)
+        end
+
+        def ones!(tensor)
+          _ones!(tensor)
+        end
+
+        def zeros!(tensor)
+          _zeros!(tensor)
+        end
+
+        def eye!(tensor)
+          _eye!(tensor)
+        end
+
+        def dirac!(tensor)
+          _dirac!(tensor)
+        end
+
+        def xavier_uniform!(tensor, gain: 1.0)
+          _xavier_uniform!(tensor, gain)
+        end
+
+        def xavier_normal!(tensor, gain: 1.0)
+          _xavier_normal!(tensor, gain)
+        end
+
+        def kaiming_uniform!(tensor, a: 0, mode: "fan_in", nonlinearity: "leaky_relu")
+          _kaiming_uniform!(tensor, a, mode, nonlinearity)
+        end
+
+        def kaiming_normal!(tensor, a: 0, mode: "fan_in", nonlinearity: "leaky_relu")
+          _kaiming_normal!(tensor, a, mode, nonlinearity)
+        end
+
+        def orthogonal!(tensor, gain: 1)
+          _orthogonal!(tensor, gain)
+        end
+
+        def sparse!(tensor, sparsity, std: 0.01)
+          _sparse!(tensor, sparsity, std)
+        end
+
+        # TODO move to C++ when released
+        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"

data/lib/torch/nn/kl_div_loss.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class KLDivLoss < Loss
+      def initialize(reduction: "mean")
+        super(reduction)
+      end
+
+      def forward(input, target)
+        F.kl_div(input, target, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/l1_loss.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class L1Loss < Loss
+      def initialize(reduction: "mean")
+        super(reduction)
+      end
+
+      def forward(input, target)
+        F.l1_loss(input, target, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/leaky_relu.rb

@@ -0,0 +1,20 @@
+module Torch
+  module NN
+    class LeakyReLU < Module
+      def initialize(negative_slope: 1e-2) #, inplace: false)
+        super()
+        @negative_slope = negative_slope
+        # @inplace = inplace
+      end
+
+      def forward(input)
+        F.leaky_relu(input, @negative_slope) #, inplace: @inplace)
+      end
+
+      def extra_inspect
+        inplace_str = @inplace ? ", inplace: true" : ""
+        format("negative_slope: %s%s", @negative_slope, inplace_str)
+      end
+    end
+  end
+end

data/lib/torch/nn/linear.rb

@@ -1,35 +1,36 @@
 module Torch
   module NN
     class Linear < Module
-      attr_reader :bias, :weight
-
       def initialize(in_features, out_features, bias: true)
+        super()
         @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))
+        else
+          register_parameter("bias", nil)
         end
 
         reset_parameters
       end
 
-      def call(input)
-        F.linear(input, @weight, @bias)
-      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
 
-      def inspect
-        "Linear(in_features: #{@in_features.inspect}, out_features: #{@out_features.inspect}, bias: #{(!@bias.nil?).inspect})"
+      def forward(input)
+        F.linear(input, @weight, @bias)
+      end
+
+      def extra_inspect
+        format("in_features: %s, out_features: %s, bias: %s", @in_features, @out_features, !@bias.nil?)
       end
     end
   end