torch-rb 0.1.3

data/lib/torch/optim/asgd.rb

@@ -0,0 +1,65 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/asgd.py
+module Torch
+  module Optim
+    class ASGD < Optimizer
+      def initialize(params, lr: 1e-2, lambd: 1e-4, alpha: 0.75, t0: 1e6, weight_decay: 0)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid weight_decay value: #{weight_decay}" if weight_decay < 0
+
+        defaults = {lr: lr, lambd: lambd, alpha: alpha, t0: t0, 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, "ASGD does not support sparse gradients"
+            end
+            state = @state[p]
+
+            # State initialization
+            if state.size == 0
+              state[:step] = 0
+              state[:eta] = group[:lr]
+              state[:mu] = 1
+              state[:ax] = Torch.zeros_like(p.data)
+            end
+
+            state[:step] += 1
+
+            if group[:weight_decay] != 0
+              grad = grad.add(group[:weight_decay], p.data)
+            end
+
+            # decay term
+            p.data.mul!(1 - group[:lambd] * state[:eta])
+
+            # update parameter
+            p.data.add!(-state[:eta], grad)
+
+            # averaging
+            if state[:mu] != 1
+              state[:ax].add!(p.data.sub(state[:ax]).mul(state[:mu]))
+            else
+              state[:ax].copy!(p.data)
+            end
+
+            # update eta and mu
+            state[:eta] = (group[:lr] / ((1 + group[:lambd] * group[:lr] * state[:step]) ** group[:alpha]))
+            state[:mu] = 1 / [1, state[:step] - group[:t0]].max
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

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

@@ -0,0 +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

@@ -0,0 +1,60 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/sgd.py
+module Torch
+  module Optim
+    class SGD < Optimizer
+      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
+
+        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(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
+end

data/lib/torch/tensor.rb

@@ -0,0 +1,196 @@
+module Torch
+  class Tensor
+    include Comparable
+    include Inspector
+
+    alias_method :requires_grad?, :requires_grad
+
+    def self.new(*size)
+      if size.length == 1 && size.first.is_a?(Tensor)
+        size.first
+      else
+        Torch.empty(*size)
+      end
+    end
+
+    def dtype
+      dtype = ENUM_TO_DTYPE[_dtype]
+      raise Error, "Unknown type: #{_dtype}" unless dtype
+      dtype
+    end
+
+    def layout
+      _layout.downcase.to_sym
+    end
+
+    def to_s
+      inspect
+    end
+
+    def to_a
+      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)
+      else
+        shape
+      end
+    end
+
+    def shape
+      dim.times.map { |i| size(i) }
+    end
+
+    def view(*size)
+      _view(size)
+    end
+
+    def item
+      if numel != 1
+        raise Error, "only one element tensors can be converted to Ruby scalars"
+      end
+      _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)
+      else
+        _backward
+      end
+    end
+
+    # TODO read directly from memory
+    def numo
+      raise Error, "Numo not found" unless defined?(Numo::NArray)
+      cls = Torch._dtype_to_numo[dtype]
+      raise Error, "Cannot convert #{dtype} to Numo" unless cls
+      cls.cast(_data).reshape(*shape)
+    end
+
+    def new_ones(*size, **options)
+      Torch.ones_like(Torch.empty(*size), **options)
+    end
+
+    def requires_grad!(requires_grad = true)
+      _requires_grad!(requires_grad)
+    end
+
+    def type(dtype)
+      enum = DTYPE_TO_ENUM[dtype]
+      raise Error, "Unknown type: #{dtype}" unless enum
+      _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(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)
+        else
+          Torch.send(op, self, *args, &block)
+        end
+      end
+    end
+
+    def +(other)
+      add(other)
+    end
+
+    def -(other)
+      sub(other)
+    end
+
+    def *(other)
+      mul(other)
+    end
+
+    def /(other)
+      div(other)
+    end
+
+    def %(other)
+      remainder(other)
+    end
+
+    def **(other)
+      pow(other)
+    end
+
+    def -@
+      neg
+    end
+
+    def <=>(other)
+      item <=> other
+    end
+
+    # based on python_variable_indexing.cpp
+    def [](*indexes)
+      result = self
+      dim = 0
+      indexes.each do |index|
+        if index.is_a?(Numeric)
+          result = result._select(dim, index)
+        elsif index.is_a?(Range)
+          finish = index.end
+          finish += 1 unless index.exclude_end?
+          result = result._slice(dim, index.begin, finish, 1)
+          dim += 1
+        else
+          raise Error, "Unsupported index type"
+        end
+      end
+      result
+    end
+
+    # TODO
+    # based on python_variable_indexing.cpp
+    # def []=(index, value)
+    # end
+
+    private
+
+    def reshape_arr(arr, dims)
+      if dims.empty?
+        arr
+      else
+        arr = arr.flatten
+        dims[1..-1].reverse.each do |dim|
+          arr = arr.each_slice(dim)
+        end
+        arr.to_a
+      end
+    end
+  end
+end