torch-rb 0.1.2 → 0.1.3

data/lib/torch/optim/lr_scheduler/lr_scheduler.rb

@@ -0,0 +1,33 @@
+module Torch
+  module Optim
+    module LRScheduler
+      class LRScheduler
+        def initialize(optimizer, last_epoch)
+          @optimizer = optimizer
+          if last_epoch == -1
+            optimizer.param_groups.each do |group|
+              group[:initial_lr] ||= group[:lr]
+            end
+            last_epoch = 0
+          else
+            raise NotImplementedYet
+          end
+          @base_lrs = optimizer.param_groups.map { |group| group[:initial_lr] }
+          @last_epoch = last_epoch
+
+          @step_count = 0
+          step(last_epoch)
+        end
+
+        def step(epoch = nil)
+          @step_count += 1
+          epoch ||= @last_epoch + 1
+          @last_epoch = epoch
+          @optimizer.param_groups.zip(get_lr).each do |param_group, lr|
+            param_group[:lr] = lr
+          end
+        end
+      end
+    end
+  end
+end

data/lib/torch/optim/lr_scheduler/step_lr.rb

@@ -0,0 +1,17 @@
+module Torch
+  module Optim
+    module LRScheduler
+      class StepLR < LRScheduler
+        def initialize(optimizer, step_size:, gamma: 0.1, last_epoch: -1)
+          @step_size = step_size
+          @gamma = gamma
+          super(optimizer, last_epoch)
+        end
+
+        def get_lr
+          @base_lrs.map { |base_lr| base_lr * @gamma ** (@last_epoch / @step_size).floor }
+        end
+      end
+    end
+  end
+end

data/lib/torch/optim/optimizer.rb

@@ -1,6 +1,62 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/optimizer.py
 module Torch
   module Optim
     class Optimizer
+      attr_reader :param_groups
+
+      def initialize(params, defaults)
+        @defaults = defaults
+        @state = Hash.new { |hash, key| hash[key] = {} }
+        @param_groups = []
+
+        param_groups = params
+        if param_groups.empty?
+          raise ArgumentError, "optimizer got an empty parameter list"
+        end
+        if !param_groups[0].is_a?(Hash)
+          param_groups = [{params: param_groups}]
+        end
+
+        param_groups.each do |param_group|
+          add_param_group(param_group)
+        end
+      end
+
+      def add_param_group(param_group)
+        # TODO more advanced logic
+        @param_groups << @defaults.merge(param_group)
+      end
+
+      def load_state_dict(state_dict)
+        raise NotImplementedYet
+      end
+
+      def state_dict
+        pack_group = lambda do |group|
+          packed = group.select { |k, _| k != :params }.to_h
+          packed[:params] = group[:params].map { |p| p.object_id }
+          packed
+        end
+
+        param_groups = @param_groups.map { |g| pack_group.call(g) }
+        packed_state = @state.map { |k, v| [k.is_a?(Tensor) ? k.object_id : k, v] }.to_h
+
+        {
+          state: packed_state,
+          param_groups: param_groups
+        }
+      end
+
+      def zero_grad
+        @param_groups.each do |group|
+          group[:params].each do |p|
+            if p.grad
+              p.grad.detach!
+              p.grad.zero!
+            end
+          end
+        end
+      end
     end
   end
 end

data/lib/torch/optim/rmsprop.rb

@@ -0,0 +1,76 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/rmsprop.py
+module Torch
+  module Optim
+    class RMSprop < Optimizer
+      def initialize(params, lr: 1e-2, alpha: 0.99, eps: 1e-8, weight_decay: 0, momentum: 0, centered: false)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid epsilon value: #{eps}" if eps < 0
+        raise ArgumentError, "Invalid momentum value: #{momentum}" if momentum < 0
+        raise ArgumentError, "Invalid weight_decay value: #{weight_decay}" if weight_decay < 0
+        raise ArgumentError, "Invalid momentum alpha: #{alpha}" if alpha < 0
+
+        defaults = {lr: lr, momentum: momentum, alpha: alpha, eps: eps, centered: centered, weight_decay: weight_decay}
+        super(params, defaults)
+      end
+
+      def step(closure = nil)
+        loss = nil
+        if closure
+          loss = closure.call
+        end
+
+        @param_groups.each do |group|
+          group[:params].each do |p|
+            next unless p.grad
+            grad = p.grad.data
+            if grad.sparse?
+              raise Error, "RMSprop does not support sparse gradients"
+            end
+            state = @state[p]
+
+            # State initialization
+            if state.size == 0
+              state[:step] = 0
+              state[:square_avg] = Torch.zeros_like(p.data)
+              if group[:momentum] > 0
+                state[:momentum_buffer] = Torch.zeros_like(p.data)
+              end
+              if group[:centered]
+                state[:grad_avg] = Torch.zeros_like(p.data)
+              end
+            end
+
+            square_avg = state[:square_avg]
+            alpha = group[:alpha]
+
+            state[:step] += 1
+
+            if group[:weight_decay] != 0
+              grad = grad.add(group[:weight_decay], p.data)
+            end
+
+            square_avg.mul!(alpha).addcmul!(1 - alpha, grad, grad)
+
+            if group[:centered]
+              grad_avg = state[:grad_avg]
+              grad_avg.mul!(alpha).add!(1 - alpha, grad)
+              avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt!.add!(group[:eps])
+            else
+              avg = square_avg.sqrt.add!(group[:eps])
+            end
+
+            if group[:momentum] > 0
+              buf = state[:momentum_buffer]
+              buf.mul!(group[:momentum]).addcdiv!(grad, avg)
+              p.data.add!(-group[:lr], buf)
+            else
+              p.data.addcdiv!(-group[:lr], grad, avg)
+            end
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

data/lib/torch/optim/rprop.rb

@@ -0,0 +1,68 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/rprop.py
+module Torch
+  module Optim
+    class Rprop < Optimizer
+      def initialize(params, lr: 1e-2, etas: [0.5, 1.2], step_sizes: [1e-6, 50])
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid eta values: #{etas[0]}, #{etas[1]}" if etas[0] < 0 || etas[0] >= 1 || etas[1] < 1
+
+        defaults = {lr: lr, etas: etas, step_sizes: step_sizes}
+        super(params, defaults)
+      end
+
+      def step(closure = nil)
+        # TODO implement []=
+        raise NotImplementedYet
+
+        loss = nil
+        if closure
+          loss = closure.call
+        end
+
+        @param_groups.each do |group|
+          group[:params].each do |p|
+            next unless p.grad
+            grad = p.grad.data
+            if grad.sparse?
+              raise Error, "Rprop does not support sparse gradients"
+            end
+            state = @state[p]
+
+            # State initialization
+            if state.size == 0
+              state[:step] = 0
+              state[:prev] = Torch.zeros_like(p.data)
+              state[:step_size] = grad.new.resize_as!(grad).fill!(group[:lr])
+            end
+
+            etaminus, etaplus = group[:etas]
+            step_size_min, step_size_max = group[:step_sizes]
+            step_size = state[:step_size]
+
+            state[:step] += 1
+
+            sign = grad.mul(state[:prev]).sign
+            sign[sign.gt(0)] = etaplus
+            sign[sign.lt(0)] = etaminus
+            sign[sign.eq(0)] = 1
+
+            # update stepsizes with step size updates
+            step_size.mul!(sign).clamp!(step_size_min, step_size_max)
+
+            # for dir<0, dfdx=0
+            # for dir>=0 dfdx=dfdx
+            grad = grad.clone
+            grad[sign.eq(etaminus)] = 0
+
+            # update parameters
+            p.data.addcmul!(-1, grad.sign, step_size)
+
+            state[:prev].copy!(grad)
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

data/lib/torch/optim/sgd.rb

@@ -1,27 +1,59 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/sgd.py
 module Torch
   module Optim
     class SGD < Optimizer
-      def initialize(params, lr:)
-        @params = params
-        @lr = lr
-      end
+      def initialize(params, lr:, momentum: 0, dampening: 0, weight_decay: 0, nesterov: false)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0.0
+        raise ArgumentError, "Invalid momentum value: #{momentum}" if momentum < 0.0
+        raise ArgumentError, "Invalid weight_decay value: #{weight_decay}" if weight_decay < 0.0
 
-      def zero_grad
-        @params.each do |param|
-          if param.grad
-            param.grad.detach!
-            param.grad.zero!
-          end
+        defaults = {lr: lr, momentum: momentum, dampening: dampening, weight_decay: weight_decay, nesterov: nesterov}
+
+        if nesterov && (momentum <= 0 || dampening != 0)
+          raise ArgumentError, "Nesterov momentum requires a momentum and zero dampening"
         end
+
+        super(params, defaults)
       end
 
-      def step
-        @params.each do |param|
-          next unless param.grad
-          d_p = param.grad.data
-          # same as param.data.add!(-@lr, d_p)
-          param.data.sub!(d_p * @lr)
+      def step(closure = nil)
+        loss = nil
+        if closure
+          loss = closure.call
+        end
+
+        @param_groups.each do |group|
+          weight_decay = group[:weight_decay]
+          momentum = group[:momentum]
+          dampening = group[:dampening]
+          nesterov = group[:nesterov]
+
+          group[:params].each do |p|
+            next unless p.grad
+            d_p = p.grad.data
+            if weight_decay != 0
+              d_p.add!(weight_decay, p.data)
+            end
+            if momentum != 0
+              param_state = @state[p]
+              if !param_state.key(:momentum_buffer)
+                buf = param_state[:momentum_buffer] = Torch.clone(d_p).detach
+              else
+                buf = param_state[:momentum_buffer]
+                buf.mul!(momentum).add!(1 - dampening, d_p)
+              end
+              if nesterov
+                d_p = d_p.add(momentum, buf)
+              else
+                d_p = buf
+              end
+            end
+
+            p.data.add!(-group[:lr], d_p)
+          end
         end
+
+        loss
       end
     end
   end

data/lib/torch/tensor.rb

@@ -6,10 +6,10 @@ module Torch
     alias_method :requires_grad?, :requires_grad
 
     def self.new(*size)
-      if size.first.is_a?(Tensor)
+      if size.length == 1 && size.first.is_a?(Tensor)
         size.first
       else
-        Torch.rand(*size)
+        Torch.empty(*size)
       end
     end
 
@@ -31,6 +31,12 @@ module Torch
       reshape_arr(_data, shape)
     end
 
+    # TODO support dtype
+    def to(device, non_blocking: false, copy: false)
+      device = Device.new(device) if device.is_a?(String)
+      _to(device, _dtype, non_blocking, copy)
+    end
+
     def size(dim = nil)
       if dim
         _size(dim)
@@ -54,6 +60,11 @@ module Torch
       _data.first
     end
 
+    # unsure if this is correct
+    def new
+      Torch.empty(0, dtype: dtype)
+    end
+
     def backward(gradient = nil)
       if gradient
         _backward_gradient(gradient)
@@ -84,8 +95,25 @@ module Torch
       _type(enum)
     end
 
+    def add!(value = 1, other)
+      if other.is_a?(Numeric)
+        _add_scalar!(other * value)
+      else
+        # need to use alpha for sparse tensors instead of multiplying
+        _add_alpha!(other, value)
+      end
+    end
+
+    def mul!(other)
+      if other.is_a?(Numeric)
+        _mul_scalar!(other)
+      else
+        _mul!(other)
+      end
+    end
+
     # operations
-    %w(add sub mul div remainder pow neg sum mean num norm min max dot matmul exp log unsqueeze reshape argmax eq).each do |op|
+    %w(abs add argmax div dot eq exp gt log lt matmul max mean min mul neg norm num numel pow remainder reshape sign sqrt sub sum unsqueeze).each do |op|
       define_method(op) do |*args, **options, &block|
         if options.any?
           Torch.send(op, self, *args, **options, &block)
@@ -127,10 +155,11 @@ module Torch
       item <=> other
     end
 
+    # based on python_variable_indexing.cpp
     def [](*indexes)
       result = self
       dim = 0
-      indexes.each_with_index do |index|
+      indexes.each do |index|
         if index.is_a?(Numeric)
           result = result._select(dim, index)
         elsif index.is_a?(Range)
@@ -145,6 +174,11 @@ module Torch
       result
     end
 
+    # TODO
+    # based on python_variable_indexing.cpp
+    # def []=(index, value)
+    # end
+
     private
 
     def reshape_arr(arr, dims)

data/lib/torch/utils/data/data_loader.rb

@@ -2,19 +2,25 @@ module Torch
   module Utils
     module Data
       class DataLoader
+        include Enumerable
+
+        attr_reader :dataset
+
         def initialize(dataset, batch_size: 1)
           @dataset = dataset
           @batch_size = batch_size
         end
 
         def each
-          size = @dataset.size
-          start_index = 0
-          while start_index < size
+          size.times do |i|
+            start_index = i * @batch_size
             yield @dataset[start_index...(start_index + @batch_size)]
-            start_index += @batch_size
           end
         end
+
+        def size
+          (@dataset.size / @batch_size.to_f).ceil
+        end
       end
     end
   end