torch-rb 0.1.1 → 0.1.6

data/lib/torch/optim/adamax.rb

@@ -0,0 +1,68 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/adamax.py
+module Torch
+  module Optim
+    class Adamax < Optimizer
+      def initialize(params, lr: 2e-3, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid epsilon value: #{eps}" if eps < 0
+        raise ArgumentError, "Invalid beta parameter at index 0: #{betas[0]}" if betas[0] < 0 || betas[0] >= 1
+        raise ArgumentError, "Invalid beta parameter at index 1: #{betas[1]}" if betas[1] < 0 || betas[1] >= 1
+        raise ArgumentError, "Invalid weight_decay value: #{weight_decay}" if weight_decay < 0
+
+        defaults = {lr: lr, betas: betas, eps: eps, 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, "Adamax does not support sparse gradients, please consider SparseAdam instead"
+            end
+            state = @state[p]
+
+            # State initialization
+            if state.size == 0
+              state[:step] = 0
+              state[:exp_avg] = Torch.zeros_like(p.data)
+              state[:exp_inf] = Torch.zeros_like(p.data)
+            end
+
+            exp_avg, exp_inf = state[:exp_avg], state[:exp_inf]
+            beta1, beta2 = group[:betas]
+            eps = group[:eps]
+
+            state[:step] += 1
+
+            if group[:weight_decay] != 0
+              grad = grad.add(group[:weight_decay], p.data)
+            end
+
+            # Update biased first moment estimate.
+            exp_avg.mul!(beta1).add!(1 - beta1, grad)
+            # Update the exponentially weighted infinity norm.
+            norm_buf = Torch.cat([
+                exp_inf.mul!(beta2).unsqueeze(0),
+                grad.abs.add!(eps).unsqueeze!(0)
+            ], 0)
+            Torch.max(norm_buf, 0, keepdim: false, out: [exp_inf, exp_inf.new.long])
+
+            bias_correction = 1 - beta1 ** state[:step]
+            clr = group[:lr] / bias_correction
+
+            p.data.addcdiv!(-clr, exp_avg, exp_inf)
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

data/lib/torch/optim/adamw.rb

@@ -0,0 +1,82 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/adamw.py
+module Torch
+  module Optim
+    class AdamW < Optimizer
+      def initialize(params, lr: 1e-3, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 1e-2, amsgrad: false)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid epsilon value: #{eps}" if eps < 0
+        raise ArgumentError, "Invalid beta parameter at index 0: #{betas[0]}" if betas[0] < 0 || betas[0] >= 1
+        raise ArgumentError, "Invalid beta parameter at index 1: #{betas[1]}" if betas[1] < 0 || betas[1] >= 1
+
+        defaults = {lr: lr, betas: betas, eps: eps, weight_decay: weight_decay, amsgrad: amsgrad}
+        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
+
+            # Perform stepweight decay
+            p.data.mul!(1 - group[:lr] * group[:weight_decay])
+
+            # Perform optimization step
+            grad = p.grad.data
+            if grad.sparse?
+              raise Error, "AdamW does not support sparse gradients, please consider SparseAdam instead"
+            end
+            amsgrad = group[:amsgrad]
+
+            state = @state[p]
+
+            # State initialization
+            if state.size == 0
+              state[:step] = 0
+              # Exponential moving average of gradient values
+              state[:exp_avg] = Torch.zeros_like(p.data)
+              # Exponential moving average of squared gradient values
+              state[:exp_avg_sq] = Torch.zeros_like(p.data)
+              if amsgrad
+                # Maintains max of all exp. moving avg. of sq. grad. values
+                state[:max_exp_avg_sq] = Torch.zeros_like(p.data)
+              end
+            end
+
+            exp_avg, exp_avg_sq = state[:exp_avg], state[:exp_avg_sq]
+            if amsgrad
+              max_exp_avg_sq = state[:max_exp_avg_sq]
+            end
+            beta1, beta2 = group[:betas]
+
+            state[:step] += 1
+            bias_correction1 = 1 - beta1 ** state[:step]
+            bias_correction2 = 1 - beta2 ** state[:step]
+
+            # Decay the first and second moment running average coefficient
+            exp_avg.mul!(beta1).add!(1 - beta1, grad)
+            exp_avg_sq.mul!(beta2).addcmul!(1 - beta2, grad, grad)
+            if amsgrad
+              # Maintains the maximum of all 2nd moment running avg. till now
+              Torch.max(max_exp_avg_sq, exp_avg_sq, out: max_exp_avg_sq)
+              # Use the max. for normalizing running avg. of gradient
+              denom = (max_exp_avg_sq.sqrt / Math.sqrt(bias_correction2)).add!(group[:eps])
+            else
+              denom = (exp_avg_sq.sqrt / Math.sqrt(bias_correction2)).add!(group[:eps])
+            end
+
+            step_size = group[:lr] / bias_correction1
+
+            p.data.addcdiv!(-step_size, exp_avg, denom)
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

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