torch-rb 0.1.1 → 0.1.6

data/lib/torch/nn/mse_loss.rb

@@ -1,8 +1,8 @@
 module Torch
   module NN
-    class MSELoss < Module
+    class MSELoss < Loss
       def initialize(reduction: "mean")
-        @reduction = reduction
+        super(reduction)
       end
 
       def forward(input, target)

data/lib/torch/nn/multi_label_margin_loss.rb

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

data/lib/torch/nn/multi_label_soft_margin_loss.rb

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

data/lib/torch/nn/multi_margin_loss.rb

@@ -0,0 +1,17 @@
+module Torch
+  module NN
+    class MultiMarginLoss < WeightedLoss
+      def initialize(p: 1, margin: 1.0, weight: nil, reduction: "mean")
+        super(weight, reduction)
+        raise ArgumentError, "only p == 1 and p == 2 supported" if p != 1 && p != 2
+        raise ArgumentError, "weight must be nil or have one dimension" unless weight.nil? || weight.dim == 1
+        @p = p
+        @margin = margin
+      end
+
+      def forward(input, target)
+        F.multi_margin_loss(input, target, p: @p, margin: @margin, weight: @weight, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/nll_loss.rb

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

data/lib/torch/nn/pairwise_distance.rb

@@ -0,0 +1,16 @@
+module Torch
+  module NN
+    class PairwiseDistance < Module
+      def initialize(p: 2.0, eps: 1e-6, keepdim: false)
+        super()
+        @norm = p
+        @eps = eps
+        @keepdim = keepdim
+      end
+
+      def forward(x1, x2)
+        F.pairwise_distance(x1, x2, p: @norm, eps: @eps, keepdim: @keepdim)
+      end
+    end
+  end
+end

data/lib/torch/nn/parameter.rb

@@ -5,6 +5,10 @@ module Torch
         data = Tensor.new unless data
         Tensor._make_subclass(data, requires_grad)
       end
+
+      def inspect
+        "Parameter containing:\n#{super}"
+      end
     end
   end
 end

data/lib/torch/nn/poisson_nll_loss.rb

@@ -0,0 +1,16 @@
+module Torch
+  module NN
+    class PoissonNLLLoss < Loss
+      def initialize(log_input: true, full: false, eps: 1e-8, reduction: "mean")
+        super(reduction)
+        @log_input = log_input
+        @full = full
+        @eps = eps
+      end
+
+      def forward(log_input, target)
+        F.poisson_nll_loss(log_input, target, log_input: @log_input, full: @full, eps: @eps, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/prelu.rb

@@ -0,0 +1,19 @@
+module Torch
+  module NN
+    class PReLU < Module
+      def initialize(num_parameters: 1, init: 0.25)
+        @num_parameters = num_parameters
+        super()
+        @weight = Parameter.new(Tensor.new(num_parameters).fill!(init))
+      end
+
+      def forward(input)
+        F.prelu(input, @weight)
+      end
+
+      def extra_inspect
+        format("num_parameters: %s", @num_parameters)
+      end
+    end
+  end
+end

data/lib/torch/nn/reflection_pad1d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ReflectionPad1d < ReflectionPadNd
+      def initialize(padding)
+        super()
+        @padding = _pair(padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/reflection_pad2d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ReflectionPad2d < ReflectionPadNd
+      def initialize(padding)
+        super()
+        @padding = _quadrupal(padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/reflection_padnd.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class ReflectionPadNd < Module
+      def forward(input)
+        F.pad(input, @padding, mode: "reflect")
+      end
+
+      def extra_inspect
+        @padding.inspect
+      end
+    end
+  end
+end

data/lib/torch/nn/relu.rb

@@ -1,12 +1,17 @@
 module Torch
   module NN
     class ReLU < Module
-      def initialize #(inplace: false)
-        # @inplace = inplace
+      def initialize(inplace: false)
+        super()
+        @inplace = inplace
       end
 
       def forward(input)
-        F.relu(input) #, inplace: @inplace)
+        F.relu(input, inplace: @inplace)
+      end
+
+      def extra_inspect
+        @inplace ? "inplace: true" : ""
       end
     end
   end

data/lib/torch/nn/replication_pad1d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ReplicationPad1d < ReplicationPadNd
+      def initialize(padding)
+        super()
+        @padding = _pair(padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/replication_pad2d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ReplicationPad2d < ReplicationPadNd
+      def initialize(padding)
+        super()
+        @padding = _quadrupal(padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/replication_pad3d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ReplicationPad3d < ReplicationPadNd
+      def initialize(padding)
+        super()
+        @padding = _ntuple(6, padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/replication_padnd.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class ReplicationPadNd < Module
+      def forward(input)
+        F.pad(input, @padding, mode: "replicate")
+      end
+
+      def extra_inspect
+        @padding.inspect
+      end
+    end
+  end
+end

data/lib/torch/nn/rnn.rb

@@ -0,0 +1,22 @@
+module Torch
+  module NN
+    class RNN < RNNBase
+      def initialize(*args, **options)
+        if options.key?(:nonlinearity)
+          if options[:nonlinearity] == "tanh"
+            mode = "RNN_TANH"
+          elsif options[:nonlinearity] == "relu"
+            mode = "RNN_RELU"
+          else
+            raise ArgumentError, "Unknown nonlinearity: #{options[:nonlinearity]}"
+          end
+          options.delete(:nonlinearity)
+        else
+          mode = "RNN_TANH"
+        end
+
+        super(mode, *args, **options)
+      end
+    end
+  end
+end

data/lib/torch/nn/rnn_base.rb

@@ -0,0 +1,198 @@
+module Torch
+  module NN
+    class RNNBase < Module
+      def initialize(mode, input_size, hidden_size, num_layers: 1, bias: true,
+        batch_first: false, dropout: 0.0, bidirectional: false)
+
+        super()
+        @mode = mode
+        @input_size = input_size
+        @hidden_size = hidden_size
+        @num_layers = num_layers
+        @bias = bias
+        @batch_first = batch_first
+        @dropout = dropout.to_f
+        @bidirectional = bidirectional
+        num_directions = bidirectional ? 2 : 1
+
+        if !dropout.is_a?(Numeric) || !(dropout >= 0 && dropout <= 1)
+          raise ArgumentError, "dropout should be a number in range [0, 1] " +
+                               "representing the probability of an element being " +
+                               "zeroed"
+        end
+        if dropout > 0 && num_layers == 1
+          warn "dropout option adds dropout after all but last " +
+               "recurrent layer, so non-zero dropout expects " +
+               "num_layers greater than 1, but got dropout=#{dropout} and " +
+               "num_layers=#{num_layers}"
+        end
+
+        gate_size =
+          case mode
+          when "LSTM"
+            4 * hidden_size
+          when "GRU"
+            3 * hidden_size
+          when "RNN_TANH"
+            hidden_size
+          when "RNN_RELU"
+            hidden_size
+          else
+            raise ArgumentError, "Unrecognized RNN mode: #{mode}"
+          end
+
+        @all_weights = []
+        num_layers.times do |layer|
+          num_directions.times do |direction|
+            layer_input_size = layer == 0 ? input_size : hidden_size * num_directions
+
+            w_ih = Parameter.new(Torch::Tensor.new(gate_size, layer_input_size))
+            w_hh = Parameter.new(Torch::Tensor.new(gate_size, hidden_size))
+            b_ih = Parameter.new(Torch::Tensor.new(gate_size))
+            # Second bias vector included for CuDNN compatibility. Only one
+            # bias vector is needed in standard definition.
+            b_hh = Parameter.new(Torch::Tensor.new(gate_size))
+            layer_params = [w_ih, w_hh, b_ih, b_hh]
+
+            suffix = direction == 1 ? "_reverse" : ""
+            param_names = ["weight_ih_l%s%s", "weight_hh_l%s%s"]
+            if bias
+              param_names += ["bias_ih_l%s%s", "bias_hh_l%s%s"]
+            end
+            param_names.map! { |x| x % [layer, suffix] }
+
+            param_names.zip(layer_params) do |name, param|
+              instance_variable_set("@#{name}", param)
+            end
+            @all_weights << param_names
+          end
+        end
+
+        flatten_parameters
+        reset_parameters
+      end
+
+      def flatten_parameters
+        # no-op unless module is on the GPU and cuDNN is enabled
+      end
+
+      def _apply(fn)
+        ret = super
+        flatten_parameters
+        ret
+      end
+
+      def reset_parameters
+        stdv = 1.0 / Math.sqrt(@hidden_size)
+        parameters.each do |weight|
+          Init.uniform!(weight, a: -stdv, b: stdv)
+        end
+      end
+
+      def permute_hidden(hx, permutation)
+        if permutation.nil?
+          return hx
+        end
+        raise NotImplementedYet
+      end
+
+      def forward(input, hx: nil)
+        is_packed = false # TODO isinstance(input, PackedSequence)
+        if is_packed
+          input, batch_sizes, sorted_indices, unsorted_indices = input
+          max_batch_size = batch_sizes[0]
+          max_batch_size = max_batch_size.to_i
+        else
+          batch_sizes = nil
+          max_batch_size = @batch_first ? input.size(0) : input.size(1)
+          sorted_indices = nil
+          unsorted_indices = nil
+        end
+
+        if hx.nil?
+          num_directions = @bidirectional ? 2 : 1
+          hx = Torch.zeros(@num_layers * num_directions, max_batch_size,
+            @hidden_size, dtype: input.dtype, device: input.device)
+        else
+          # Each batch of the hidden state should match the input sequence that
+          # the user believes he/she is passing in.
+          hx = permute_hidden(hx, sorted_indices)
+        end
+
+        check_forward_args(input, hx, batch_sizes)
+        _rnn_impls = {
+          "RNN_TANH" => Torch.method(:rnn_tanh),
+          "RNN_RELU" => Torch.method(:rnn_relu)
+        }
+        _impl = _rnn_impls[@mode]
+        if batch_sizes.nil?
+          result = _impl.call(input, hx, _get_flat_weights, @bias, @num_layers,
+                           @dropout, @training, @bidirectional, @batch_first)
+        else
+          result = _impl.call(input, batch_sizes, hx, _get_flat_weights, @bias,
+                           @num_layers, @dropout, @training, @bidirectional)
+        end
+        output = result[0]
+        hidden = result[1]
+
+        if is_packed
+          raise NotImplementedYet
+          # output = PackedSequence(output, batch_sizes, sorted_indices, unsorted_indices)
+        end
+        [output, permute_hidden(hidden, unsorted_indices)]
+      end
+
+      # TODO add more parameters
+      def extra_inspect
+        s = String.new("%{input_size}, %{hidden_size}")
+        if @num_layers != 1
+          s += ", num_layers: %{num_layers}"
+        end
+        format(s, input_size: @input_size, hidden_size: @hidden_size, num_layers: @num_layers)
+      end
+
+      private
+
+      def _flat_weights
+        @all_weights.flatten.map { |v| instance_variable_get("@#{v}") }.compact
+      end
+
+      def _get_flat_weights
+        _flat_weights
+      end
+
+      def check_input(input, batch_sizes)
+        expected_input_dim = !batch_sizes.nil? ? 2 : 3
+        if input.dim != expected_input_dim
+          raise ArgumentError, "input must have #{expected_input_dim} dimensions, got #{input.dim}"
+        end
+        if @input_size != input.size(-1)
+          raise ArgumentError, "input.size(-1) must be equal to input_size. Expected #{@input_size}, got #{input.size(-1)}"
+        end
+      end
+
+      def get_expected_hidden_size(input, batch_sizes)
+        if !batch_sizes.nil?
+          mini_batch = batch_sizes[0]
+          mini_batch = mini_batch.to_i
+        else
+          mini_batch = @batch_first ? input.size(0) : input.size(1)
+        end
+        num_directions = @bidirectional ? 2 : 1
+        [@num_layers * num_directions, mini_batch, @hidden_size]
+      end
+
+      def check_hidden_size(hx, expected_hidden_size)
+        if hx.size != expected_hidden_size
+          raise ArgumentError, "Expected hidden size #{expected_hidden_size.inspect}, got #{hx.size.inspect}"
+        end
+      end
+
+      def check_forward_args(input, hidden, batch_sizes)
+        check_input(input, batch_sizes)
+        expected_hidden_size = get_expected_hidden_size(input, batch_sizes)
+        check_hidden_size(hidden, expected_hidden_size)
+      end
+    end
+  end
+end