torch-rb 0.1.5 → 0.1.6

data/lib/torch/nn/group_norm.rb

@@ -0,0 +1,36 @@
+module Torch
+  module NN
+    class GroupNorm < Module
+      def initialize(num_groups, num_channels, eps: 1e-5, affine: true)
+        super()
+        @num_groups = num_groups
+        @num_channels = num_channels
+        @eps = eps
+        @affine = affine
+        if @affine
+          @weight = Parameter.new(Torch::Tensor.new(num_channels))
+          @bias = Parameter.new(Torch::Tensor.new(num_channels))
+        else
+          register_parameter("weight", nil)
+          register_parameter("bias", nil)
+        end
+        reset_parameters
+      end
+
+      def reset_parameters
+        if @affine
+          Init.ones!(@weight)
+          Init.zeros!(@bias)
+        end
+      end
+
+      def forward(input)
+        F.group_norm(input, @num_groups, weight: @weight, bias: @bias, eps: @eps)
+      end
+
+      def extra_inspect
+        format("%{num_groups}, %{num_channels}, eps: %{eps}, affine: %{affine}", **dict)
+      end
+    end
+  end
+end

data/lib/torch/nn/gru.rb

@@ -0,0 +1,49 @@
+module Torch
+  module NN
+    class GRU < RNNBase
+      def initialize(*args, **options)
+        super("GRU", *args, **options)
+      end
+
+      def run_impl(input, hx, batch_sizes)
+        if batch_sizes.nil?
+          Torch.gru(input, hx, _get_flat_weights, @bias, @num_layers,
+                             @dropout, @training, @bidirectional, @batch_first)
+        else
+          Torch.gru(input, batch_sizes, hx, _get_flat_weights, @bias,
+                             @num_layers, @dropout, @training, @bidirectional)
+        end
+      end
+
+      def forward_impl(input, hx, batch_sizes, max_batch_size, sorted_indices)
+        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)
+        result = run_impl(input, hx, batch_sizes)
+        output = result[0]
+        hidden = result[1]
+        [output, hidden]
+      end
+
+      def forward_tensor(input, hx: nil)
+        batch_sizes = nil
+        max_batch_size = @batch_first ? input.size(0) : input.size(1)
+        sorted_indices = nil
+        unsorted_indices = nil
+        output, hidden = forward_impl(input, hx, batch_sizes, max_batch_size, sorted_indices)
+        [output, permute_hidden(hidden, unsorted_indices)]
+      end
+
+      def forward(input, hx: nil)
+        forward_tensor(input, hx: hx)
+      end
+    end
+  end
+end

data/lib/torch/nn/hardshrink.rb

@@ -0,0 +1,18 @@
+module Torch
+  module NN
+    class Hardshrink < Module
+      def initialize(lambd: 0.5)
+        super()
+        @lambd = lambd
+      end
+
+      def forward(input)
+        F.hardshrink(input, @lambd)
+      end
+
+      def extra_inspect
+        @lambd.to_s
+      end
+    end
+  end
+end

data/lib/torch/nn/instance_norm.rb

@@ -0,0 +1,20 @@
+module Torch
+  module NN
+    class InstanceNorm < BatchNorm
+      def initialize(num_features, eps: 1e-5, momentum: 0.1, affine: false, track_running_stats: false)
+        super(num_features, eps: eps, momentum: momentum, affine: affine, track_running_stats: track_running_stats)
+      end
+
+      def forward(input)
+        _check_input_dim(input)
+
+        F.instance_norm(
+          input, running_mean: @running_mean, running_var: @running_var,
+          weight: @weight, bias: @bias,
+          use_input_stats: @training || !@track_running_stats,
+          momentum: @momentum, eps: @eps
+        )
+      end
+    end
+  end
+end

data/lib/torch/nn/instance_norm1d.rb

@@ -0,0 +1,18 @@
+module Torch
+  module NN
+    class InstanceNorm1d < InstanceNorm
+      def _check_input_dim(input)
+        if input.dim == 2
+          raise ArgumentError,
+            "InstanceNorm1d returns 0-filled tensor to 2D tensor." +
+            "This is because InstanceNorm1d reshapes inputs to" +
+            "(1, N * C, ...) from (N, C,...) and this makes" +
+            "variances 0."
+        end
+        if input.dim != 3
+          raise "expected 3D input (got #{input.dim}D input)"
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/instance_norm2d.rb

@@ -0,0 +1,11 @@
+module Torch
+  module NN
+    class InstanceNorm2d < InstanceNorm
+      def _check_input_dim(input)
+        if input.dim != 4
+          raise ArgumentError, "expected 4D input (got #{input.dim}D input)"
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/instance_norm3d.rb

@@ -0,0 +1,11 @@
+module Torch
+  module NN
+    class InstanceNorm3d < InstanceNorm
+      def _check_input_dim(input)
+        if input.dim != 5
+          raise ArgumentError, "expected 5D input (got #{input.dim}D input)"
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/layer_norm.rb

@@ -0,0 +1,35 @@
+module Torch
+  module NN
+    class LayerNorm < Module
+      def initialize(normalized_shape, eps: 1e-5, elementwise_affine: true)
+        super()
+        @normalized_shape = Array(normalized_shape)
+        @eps = eps
+        @elementwise_affine = elementwise_affine
+        if @elementwise_affine
+          @weight = Parameter.new(Torch::Tensor.new(*normalized_shape))
+          @bias = Parameter.new(Torch::Tensor.new(*normalized_shape))
+        else
+          register_parameter("weight", nil)
+          register_parameter("bias", nil)
+        end
+        reset_parameters
+      end
+
+      def reset_parameters
+        if @elementwise_affine
+          Init.ones!(@weight)
+          Init.zeros!(@bias)
+        end
+      end
+
+      def forward(input)
+        F.layer_norm(input, @normalized_shape, weight: @weight, bias: @bias, eps: @eps)
+      end
+
+      def extra_inspect
+        format("%{normalized_shape}, eps: %{eps}, elementwise_affine: %{elementwise_affine}", **dict)
+      end
+    end
+  end
+end

data/lib/torch/nn/local_response_norm.rb

@@ -0,0 +1,21 @@
+module Torch
+  module NN
+    class LocalResponseNorm < Module
+      def initialize(size, alpha: 1e-4, beta: 0.75, k: 1.0)
+        super()
+        @size = size
+        @alpha = alpha
+        @beta = beta
+        @k = k
+      end
+
+      def forward(input)
+        F.local_response_norm(input, @size, alpha: @alpha, beta: @beta, k: @k)
+      end
+
+      def extra_inspect
+        format("%{size}, alpha: %{alpha}, beta: %{beta}, k: %{k}", **dict)
+      end
+    end
+  end
+end

data/lib/torch/nn/log_sigmoid.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class LogSigmoid < Module
+      def forward(input)
+        F.log_sigmoid(input)
+      end
+    end
+  end
+end

data/lib/torch/nn/lp_pool1d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class LPPool1d < LPPoolNd
+      def forward(input)
+        F.lp_pool1d(input, @norm_type.to_f, @kernel_size, @stride, @ceil_mode)
+      end
+    end
+  end
+end

data/lib/torch/nn/lp_pool2d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class LPPool2d < LPPoolNd
+      def forward(input)
+        F.lp_pool2d(input, @norm_type.to_f, @kernel_size, @stride, @ceil_mode)
+      end
+    end
+  end
+end

data/lib/torch/nn/lp_poolnd.rb

@@ -0,0 +1,22 @@
+module Torch
+  module NN
+    class LPPoolNd < Module
+      def initialize(norm_type, kernel_size, stride: nil, ceil_mode: false)
+        super()
+        @norm_type = norm_type
+        @kernel_size = kernel_size
+        @stride = stride
+        @ceil_mode = ceil_mode
+      end
+
+      def extra_inspect
+        format("norm_type: %{norm_type}, kernel_size: %{kernel_size}, stride: %{stride}, ceil_mode: %{ceil_mode}",
+          norm_type: @norm_type,
+          kernel_size: @kernel_size,
+          stride: @stride,
+          ceil_mode: @ceil_mode
+        )
+      end
+    end
+  end
+end

data/lib/torch/nn/lstm.rb

@@ -0,0 +1,66 @@
+module Torch
+  module NN
+    class LSTM < RNNBase
+      def initialize(*args, **options)
+        super("LSTM", *args, **options)
+      end
+
+      def check_forward_args(input, hidden, batch_sizes)
+        check_input(input, batch_sizes)
+        expected_hidden_size = get_expected_hidden_size(input, batch_sizes)
+
+        # TODO pass message
+        check_hidden_size(hidden[0], expected_hidden_size)
+        check_hidden_size(hidden[1], expected_hidden_size)
+      end
+
+      def permute_hidden(hx, permutation)
+        if permutation.nil?
+          return hx
+        end
+        raise NotImplementedYet
+      end
+
+      def forward_impl(input, hx, batch_sizes, max_batch_size, sorted_indices)
+        if hx.nil?
+          num_directions = @bidirectional ? 2 : 1
+          zeros = Torch.zeros(@num_layers * num_directions, max_batch_size, @hidden_size, dtype: input.dtype, device: input.device)
+          hx = [zeros, zeros]
+        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)
+        if batch_sizes.nil?
+          result = Torch.lstm(input, hx, _get_flat_weights, @bias, @num_layers,
+                              @dropout, @training, @bidirectional, @batch_first)
+        else
+          result = Torch.lstm(input, batch_sizes, hx, _get_flat_weights, @bias,
+                              @num_layers, @dropout, @training, @bidirectional)
+        end
+        output = result[0]
+        hidden = result[1..-1]
+
+        [output, hidden]
+      end
+
+      def forward_tensor(input, hx: nil)
+        batch_sizes = nil
+        max_batch_size = @batch_first ? input.size(0) : input.size(1)
+        sorted_indices = nil
+        unsorted_indices = nil
+
+        output, hidden = forward_impl(input, hx, batch_sizes, max_batch_size, sorted_indices)
+
+        [output, permute_hidden(hidden, unsorted_indices)]
+      end
+
+      def forward(input, hx: nil)
+        # TODO PackedSequence
+        forward_tensor(input, hx: hx)
+      end
+    end
+  end
+end

data/lib/torch/nn/max_pool1d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class MaxPool1d < MaxPoolNd
+      def forward(input)
+        F.max_pool1d(input, @kernel_size, @stride, @padding, @dilation, @ceil_mode, @return_indices)
+      end
+    end
+  end
+end

data/lib/torch/nn/max_pool2d.rb

@@ -2,7 +2,7 @@ module Torch
   module NN
     class MaxPool2d < MaxPoolNd
       def forward(input)
-        F.max_pool2d(input, @kernel_size) # TODO other parameters
+        F.max_pool2d(input, @kernel_size, @stride, @padding, @dilation, @ceil_mode, @return_indices)
       end
     end
   end

data/lib/torch/nn/max_pool3d.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class MaxPool3d < MaxPoolNd
+      def forward(input)
+        F.max_pool3d(input, @kernel_size, @stride, @padding, @dilation, @ceil_mode, @return_indices)
+      end
+    end
+  end
+end

data/lib/torch/nn/max_poolnd.rb

@@ -1,14 +1,14 @@
 module Torch
   module NN
     class MaxPoolNd < Module
-      def initialize(kernel_size) #, stride: nil, padding: 0, dilation: 1, return_indices: false, ceil_mode: false)
+      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
+        @stride = stride || kernel_size
+        @padding = padding
+        @dilation = dilation
+        @return_indices = return_indices
+        @ceil_mode = ceil_mode
       end
 
       def extra_inspect

data/lib/torch/nn/max_unpool1d.rb

@@ -0,0 +1,16 @@
+module Torch
+  module NN
+    class MaxUnpool1d < MaxUnpoolNd
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super()
+        @kernel_size = _single(kernel_size)
+        @stride = _single(stride || kernel_size)
+        @padding = _single(padding)
+      end
+
+      def forward(input, indices, output_size: nil)
+        F.max_unpool1d(input, indices, @kernel_size, stride: @stride, padding: @padding, output_size: output_size)
+      end
+    end
+  end
+end

data/lib/torch/nn/max_unpool2d.rb

@@ -0,0 +1,16 @@
+module Torch
+  module NN
+    class MaxUnpool2d < MaxUnpoolNd
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super()
+        @kernel_size = _pair(kernel_size)
+        @stride = _pair(stride || kernel_size)
+        @padding = _pair(padding)
+      end
+
+      def forward(input, indices, output_size: nil)
+        F.max_unpool2d(input, indices, @kernel_size, @stride, @padding, output_size)
+      end
+    end
+  end
+end