torch-rb 0.1.0 → 0.1.5

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

@@ -1,55 +1,220 @@
 module Torch
   module NN
     class Module
-      def inspect
-        str = String.new
-        str << "#{self.class.name}(\n"
-        modules.each do |name, mod|
-          str << "  (#{name}): #{mod.inspect}\n"
+      def initialize
+        @training = true
+        @parameters = {}
+        @buffers = {}
+        @modules = {}
+      end
+
+      def forward
+        raise NotImplementedError
+      end
+
+      def register_buffer(name, tensor)
+        # TODO add checks
+        @buffers[name] = tensor
+      end
+
+      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 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 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)
+        convert = lambda do |t|
+          t.to(device)
         end
-        str << ")"
+
+        _apply(convert)
       end
 
       def call(*input)
         forward(*input)
       end
 
+      def 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
         parameters.each do |param|
           if param.grad
-            raise Error, "Not supported yet"
             param.grad.detach!
             param.grad.zero!
           end
         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
     end
   end

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