torch-rb 0.1.3 → 0.1.4

data/lib/torch/nn/functional.rb

@@ -2,8 +2,12 @@ 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, stride: 1, padding: 0, dilation: 1, groups: 1)
@@ -11,6 +15,14 @@ module Torch
           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)
           kernel_size = [kernel_size, kernel_size] if kernel_size.is_a?(Integer)
           Torch.max_pool2d(input, kernel_size)
@@ -21,24 +33,102 @@ module Torch
           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
+          Torch._embedding(input, weight, 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")
+          raise NotImplementedYet if weight
+          Torch.binary_cross_entropy(input, target, reduction)
+        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(log_probs, targets, input_lengths.to_a, target_lengths.to_a, blank, reduction, zero_infinity)
+        end
+
+        def kl_div(input, target, reduction: "mean")
+          Torch.kl_div(input, target, reduction)
+        end
+
+        def l1_loss(input, target, reduction: "mean")
+          Torch.l1_loss(input, target, reduction)
+        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)
+        def nll_loss(input, target, weight: nil, ignore_index: -100, reduction: "mean")
+          raise NotImplementedYet if weight
+          Torch.nll_loss(input, target, reduction, ignore_index)
+        end
+
+        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 nll_loss(input, target, reduction: "mean")
-          # TODO fix for non-1d
-          Torch.nll_loss(input, target, reduction)
+        # 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)
+          Torch._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
 
@@ -84,12 +174,10 @@ module Torch
           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
+        private
 
-          padding_idx ||= -1
-          Torch._embedding(input, weight, padding_idx, scale_grad_by_freq, sparse)
+        def softmax_dim(ndim)
+          ndim == 0 || ndim == 1 || ndim == 3 ? 0 : 1
         end
       end
     end

data/lib/torch/nn/identity.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class Identity < Module
+      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

data/lib/torch/nn/log_softmax.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class LogSoftmax < Module
+      def initialize(dim: nil)
+        super()
+        @dim = dim
+      end
+
+      def forward(input)
+        F.log_softmax(input, @dim)
+      end
+    end
+  end
+end

data/lib/torch/nn/loss.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class Loss < Module
+      def initialize(reduction)
+        super()
+        @reduction = reduction
+      end
+    end
+  end
+end

data/lib/torch/nn/max_pool2d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class MaxPool2d < MaxPoolNd
+      def forward(input)
+        F.max_pool2d(input, @kernel_size) # TODO other parameters
+      end
+    end
+  end
+end

data/lib/torch/nn/max_poolnd.rb

@@ -0,0 +1,19 @@
+module Torch
+  module NN
+    class MaxPoolNd < Module
+      def initialize(kernel_size) #, stride: nil, padding: 0, dilation: 1, return_indices: false, ceil_mode: false)
+        super()
+        @kernel_size = kernel_size
+        # @stride = stride || kernel_size
+        # @padding = padding
+        # @dilation = dilation
+        # @return_indices = return_indices
+        # @ceil_mode = ceil_mode
+      end
+
+      def extra_inspect
+        format("kernel_size: %s", @kernel_size)
+      end
+    end
+  end
+end

data/lib/torch/nn/module.rb

@@ -3,45 +3,85 @@ module Torch
     class Module
       def initialize
         @training = true
+        @parameters = {}
+        @buffers = {}
+        @modules = {}
       end
 
-      def inspect
-        str = String.new
-        str << "#{self.class.name}(\n"
-        modules.each do |name, mod|
-          str << "  (#{name}): #{mod.inspect}\n"
-        end
-        str << ")"
+      def forward
+        raise NotImplementedError
       end
 
-      def train(mode = true)
-        @training = mode
+      def register_buffer(name, tensor)
+        # TODO add checks
+        @buffers[name] = tensor
+      end
 
-        modules.each do |_, mod|
-          mod.train(mode)
+      def register_parameter(name, param)
+        # TODO add checks
+        @parameters[name] = param
+      end
+
+      def add_module(name, mod)
+        # TODO add checks
+        @modules[name] = mod
+      end
+
+      def _apply(fn)
+        children.each do |mod|
+          mod._apply(fn)
         end
+        # TODO apply to more objects
+        self
       end
 
-      def eval
-        train(false)
+      def apply(fn)
+        children.each do |mod|
+          mod.apply(fn)
+        end
+        fn.call(self)
+        self
       end
 
-      def call(*input)
-        forward(*input)
+      def cuda(device: nil)
+        _apply ->(t) { t.cuda(device) }
+      end
+
+      def cpu
+        _apply ->(t) { t.cpu }
+      end
+
+      def type(dst_type)
+        _apply ->(t) { t.type(dst_type) }
+      end
+
+      def float
+        _apply ->(t) { t.floating_point? ? t.float : t }
+      end
+
+      def double
+        _apply ->(t) { t.floating_point? ? t.double : t }
+      end
+
+      def half
+        _apply ->(t) { t.floating_point? ? t.half : t }
       end
 
       # modifies in-place
       def to(device)
-        instance_variables.each do |name|
-          param = instance_variable_get(name)
-          if param.is_a?(Parameter)
-            instance_variable_set(name, Parameter.new(param.to(device)))
-          end
+        convert = lambda do |t|
+          t.to(device)
         end
-        modules.each do |_, mod|
-          mod.to(device)
-        end
-        self
+
+        _apply(convert)
+      end
+
+      def call(*input)
+        forward(*input)
+      end
+
+      def state_dict
+        raise NotImplementedYet
       end
 
       def parameters
@@ -53,6 +93,38 @@ module Torch
         params + modules.flat_map { |_, mod| mod.parameters }
       end
 
+      def children
+        @modules.values
+      end
+
+      def modules
+        modules = {}
+        instance_variables.each do |name|
+          mod = instance_variable_get(name)
+          modules[name[1..-1]] = mod if mod.is_a?(Module)
+        end
+        @modules.merge(modules)
+      end
+
+      def train(mode = true)
+        @training = mode
+        children.each do |mod|
+          mod.train(mode)
+        end
+        self
+      end
+
+      def eval
+        train(false)
+      end
+
+      def requires_grad!(requires_grad: true)
+        parameters.each do |p|
+          p.requires_grad!(requires_grad)
+        end
+        self
+      end
+
       def zero_grad
         parameters.each do |param|
           if param.grad
@@ -62,6 +134,24 @@ module Torch
         end
       end
 
+      def share_memory
+        _apply ->(t) { t.share_memory! }
+      end
+
+      def inspect
+        name = self.class.name.split("::").last
+        if modules.empty?
+          "#{name}(#{extra_inspect})"
+        else
+          str = String.new
+          str << "#{name}(\n"
+          modules.each do |name, mod|
+            str << "  (#{name}): #{mod.inspect}\n"
+          end
+          str << ")"
+        end
+      end
+
       def method_missing(method, *args, &block)
         modules[method.to_s] || super
       end
@@ -72,13 +162,12 @@ module Torch
 
       private
 
-      def modules
-        modules = {}
-        instance_variables.each do |name|
-          mod = instance_variable_get(name)
-          modules[name[1..-1]] = mod if mod.is_a?(Module)
-        end
-        modules
+      def extra_inspect
+        nil
+      end
+
+      def format(str, *vars)
+        str % vars.map(&:inspect)
       end
     end
   end