torch-rb 0.1.3 → 0.1.8

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/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/margin_ranking_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class MarginRankingLoss < Loss
+      def initialize(margin: 1.0, reduction: "mean")
+        super(reduction)
+        @margin = margin
+      end
+
+      def forward(input1, input2, target)
+        F.margin_ranking_loss(input1, input2, target, margin: @margin, reduction: @reduction)
+      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

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class MaxPool2d < MaxPoolNd
+      def forward(input)
+        F.max_pool2d(input, @kernel_size, @stride, @padding, @dilation, @ceil_mode, @return_indices)
+      end
+    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

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

data/lib/torch/nn/max_unpool3d.rb

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

data/lib/torch/nn/max_unpoolnd.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class MaxUnpoolNd < Module
+      def extra_inspect
+        format("kernel_size: %s, stride: %s, padding: %s", @kernel_size, @stride, @padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/module.rb

@@ -1,56 +1,170 @@
 module Torch
   module NN
     class Module
+      include Utils
+
       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
+        instance_variable_set("@#{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 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 call(*input)
-        forward(*input)
+      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)
+
+        _apply(convert)
+      end
+
+      def call(*input, **kwargs)
+        forward(*input, **kwargs)
+      end
+
+      def state_dict(destination: nil)
+        destination ||= {}
+        named_parameters.each do |k, v|
+          destination[k] = v
         end
-        self
+        destination
+      end
+
+      def load_state_dict(state_dict)
+        raise NotImplementedYet
       end
 
       def parameters
-        params = []
+        named_parameters.values
+      end
+
+      def named_parameters(prefix: "", recurse: true)
+        params = {}
+        if recurse
+          named_children.each do |name, mod|
+            params.merge!(mod.named_parameters(prefix: "#{name}.", recurse: recurse))
+          end
+        end
         instance_variables.each do |name|
           param = instance_variable_get(name)
-          params << param if param.is_a?(Parameter)
+          params[[prefix, name[1..-1]].join] = param if param.is_a?(Parameter)
+        end
+        @parameters.each do |name, param|
+          params[[prefix, name].join] = param
         end
-        params + modules.flat_map { |_, mod| mod.parameters }
+        params
+      end
+
+      def buffers
+        named_buffers.values
+      end
+
+      def named_buffers
+        @buffers || {}
+      end
+
+      def children
+        named_children.values
+      end
+
+      def named_children
+        modules = {}
+        instance_variables.each do |name|
+          mod = instance_variable_get(name)
+          modules[name[1..-1]] = mod if mod.is_a?(Module)
+        end
+        @modules.each do |name, mod|
+          modules[name] = mod
+        end
+        modules
+      end
+
+      def modules
+        named_modules.values
+      end
+
+      def named_modules
+        {"" => self}.merge(named_children)
+      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
@@ -62,23 +176,60 @@ module Torch
         end
       end
 
+      def share_memory
+        _apply ->(t) { t.share_memory! }
+      end
+
+      def inspect
+        name = self.class.name.split("::").last
+        if children.empty?
+          "#{name}(#{extra_inspect})"
+        else
+          str = String.new
+          str << "#{name}(\n"
+          children.each do |name, mod|
+            str << "  (#{name}): #{mod.inspect}\n"
+          end
+          str << ")"
+        end
+      end
+
       def method_missing(method, *args, &block)
-        modules[method.to_s] || super
+        name = method.to_s
+        if named_parameters.key?(name)
+          named_parameters[name]
+        elsif named_buffers.key?(name)
+          named_buffers[name]
+        elsif named_modules.key?(name)
+          named_modules[name]
+        else
+          super
+        end
       end
 
       def respond_to?(method, include_private = false)
-        modules.key?(method.to_s) || super
+        name = method.to_s
+        named_parameters.key?(name) || named_buffers.key?(name) || named_modules.key?(name) || super
       end
 
       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, **options)
+        vars =
+          if vars.any?
+            vars.map(&:inspect)
+          else
+            options.map { |k, v| [k, v.inspect] }.to_h
+          end
+        str % vars
+      end
+
+      def dict
+        instance_variables.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
       end
     end
   end