torch-rb 0.1.2 → 0.1.3

data/lib/torch/nn/linear.rb

@@ -20,11 +20,11 @@ module Torch
       end
 
       def reset_parameters
-        Init.kaiming_uniform_(@weight, Math.sqrt(5))
+        Init.kaiming_uniform!(@weight, Math.sqrt(5))
         if @bias
           fan_in, _ = Init.calculate_fan_in_and_fan_out(@weight)
           bound = 1 / Math.sqrt(fan_in)
-          Init.uniform_(@bias, -bound, bound)
+          Init.uniform!(@bias, -bound, bound)
         end
       end
 

data/lib/torch/nn/module.rb

@@ -1,6 +1,10 @@
 module Torch
   module NN
     class Module
+      def initialize
+        @training = true
+      end
+
       def inspect
         str = String.new
         str << "#{self.class.name}(\n"
@@ -10,10 +14,36 @@ module Torch
         str << ")"
       end
 
+      def train(mode = true)
+        @training = mode
+
+        modules.each do |_, mod|
+          mod.train(mode)
+        end
+      end
+
+      def eval
+        train(false)
+      end
+
       def call(*input)
         forward(*input)
       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
+        end
+        modules.each do |_, mod|
+          mod.to(device)
+        end
+        self
+      end
+
       def parameters
         params = []
         instance_variables.each do |name|

data/lib/torch/optim/adadelta.rb

@@ -0,0 +1,57 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/adadelta.py
+module Torch
+  module Optim
+    class Adadelta < Optimizer
+      def initialize(params, lr: 1.0, rho: 0.9, eps: 1e-6, weight_decay: 0)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid rho value: #{rho}" if rho < 0 || rho > 1
+        raise ArgumentError, "Invalid epsilon value: #{eps}" if eps < 0
+        raise ArgumentError, "Invalid weight_decay value: #{weight_decay}" if weight_decay < 0
+
+        defaults = {lr: lr, rho: rho, 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, "Adadelta does not support sparse gradients"
+            end
+            state = @state[p]
+
+            if state.size == 0
+              state[:step] = 0
+              state[:square_avg] = Torch.zeros_like(p.data)
+              state[:acc_delta] = Torch.zeros_like(p.data)
+            end
+
+            square_avg, acc_delta = state[:square_avg], state[:acc_delta]
+            rho, eps = group[:rho], group[:eps]
+
+            state[:step] += 1
+
+            if group[:weight_decay] != 0
+              grad = grad.add(group[:weight_decay], p.data)
+            end
+
+            square_avg.mul!(rho).addcmul!(1 - rho, grad, grad)
+            std = square_avg.add(eps).sqrt!
+            delta = acc_delta.add(eps).sqrt!.div!(std).mul!(grad)
+            p.data.add!(-group[:lr], delta)
+            acc_delta.mul!(rho).addcmul!(1 - rho, delta, delta)
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

data/lib/torch/optim/adagrad.rb

@@ -0,0 +1,71 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/adagrad.py
+module Torch
+  module Optim
+    class Adagrad < Optimizer
+      def initialize(params, lr: 1e-2, lr_decay: 0, weight_decay: 0, initial_accumulator_value: 0, eps: 1e-10)
+        raise ArgumentError, "Invalid learning rate: #{lr}" if lr < 0
+        raise ArgumentError, "Invalid lr_decay value: #{lr_decay}" if lr_decay < 0
+        raise ArgumentError, "Invalid initial_accumulator_value value: #{initial_accumulator_value}" if initial_accumulator_value < 0
+        raise ArgumentError, "Invalid weight_decay value: #{weight_decay}" if weight_decay < 0
+        raise ArgumentError, "Invalid epsilon value: #{eps}" if eps < 0
+
+        defaults = {lr: lr, lr_decay: lr_decay, eps: eps, weight_decay: weight_decay, initial_accumulator_value: initial_accumulator_value}
+        super(params, defaults)
+
+        @param_groups.each do |group|
+          group[:params].each do |p|
+            state = @state[p]
+            state[:step] = 0
+            state[:sum] = Torch.full_like(p.data, initial_accumulator_value)
+          end
+        end
+      end
+
+      def share_memory
+        @param_groups.each do |group|
+          group[:params].each do |p|
+            state = @state[p]
+            state[:sum].share_memory!
+          end
+        end
+      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
+            state = @state[p]
+
+            state[:step] += 1
+
+            if group[:weight_decay] != 0
+              if p.grad.data.sparse?
+                raise Error, "weight_decay option is not compatible with sparse gradients"
+              end
+              grad = grad.add(group[:weight_decay], p.data)
+            end
+
+            clr = group[:lr] / (1 + (state[:step] - 1) * group[:lr_decay])
+
+            if grad.sparse?
+              raise NotImplementedYet
+            else
+              state[:sum].addcmul!(1, grad, grad)
+              std = state[:sum].sqrt.add!(group[:eps])
+              p.data.addcdiv!(-clr, grad, std)
+            end
+          end
+        end
+
+        loss
+      end
+    end
+  end
+end

data/lib/torch/optim/adam.rb

@@ -0,0 +1,81 @@
+# ported from https://github.com/pytorch/pytorch/blob/master/torch/optim/adam.py
+module Torch
+  module Optim
+    class Adam < Optimizer
+      def initialize(params, lr: 1e-3, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0, 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
+            grad = p.grad.data
+            if grad.sparse?
+              raise Error, "Adam 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]
+
+            if group[:weight_decay] != 0
+              grad.add!(group[:weight_decay], p.data)
+            end
+
+            # 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/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