torch-rb 0.1.3

data/lib/torch/nn/module.rb

@@ -0,0 +1,85 @@
+module Torch
+  module NN
+    class Module
+      def initialize
+        @training = true
+      end
+
+      def inspect
+        str = String.new
+        str << "#{self.class.name}(\n"
+        modules.each do |name, mod|
+          str << "  (#{name}): #{mod.inspect}\n"
+        end
+        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|
+          param = instance_variable_get(name)
+          params << param if param.is_a?(Parameter)
+        end
+        params + modules.flat_map { |_, mod| mod.parameters }
+      end
+
+      def zero_grad
+        parameters.each do |param|
+          if param.grad
+            param.grad.detach!
+            param.grad.zero!
+          end
+        end
+      end
+
+      def method_missing(method, *args, &block)
+        modules[method.to_s] || super
+      end
+
+      def respond_to?(method, include_private = false)
+        modules.key?(method.to_s) || 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
+      end
+    end
+  end
+end

data/lib/torch/nn/mse_loss.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class MSELoss < Module
+      def initialize(reduction: "mean")
+        @reduction = reduction
+      end
+
+      def forward(input, target)
+        F.mse_loss(input, target, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/parameter.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class Parameter < Tensor
+      def self.new(data = nil, requires_grad: true)
+        data = Tensor.new unless data
+        Tensor._make_subclass(data, requires_grad)
+      end
+
+      def grad
+        _grad if _grad_defined
+      end
+    end
+  end
+end

data/lib/torch/nn/relu.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class ReLU < Module
+      def initialize #(inplace: false)
+        # @inplace = inplace
+      end
+
+      def forward(input)
+        F.relu(input) #, inplace: @inplace)
+      end
+    end
+  end
+end

data/lib/torch/nn/sequential.rb

@@ -0,0 +1,29 @@
+module Torch
+  module NN
+    class Sequential < Module
+      def initialize(*args)
+        @modules = {}
+        # TODO support hash arg (named modules)
+        args.each_with_index do |mod, idx|
+          add_module(idx.to_s, mod)
+        end
+      end
+
+      def add_module(name, mod)
+        # TODO add checks
+        @modules[name] = mod
+      end
+
+      def forward(input)
+        @modules.values.each do |mod|
+          input = mod.call(input)
+        end
+        input
+      end
+
+      def parameters
+        @modules.flat_map { |_, mod| mod.parameters }
+      end
+    end
+  end
+end

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