torch-rb 0.4.2 → 0.6.0

data/lib/torch/nn/linear.rb

@@ -1,6 +1,8 @@
 module Torch
   module NN
     class Linear < Module
+      attr_reader :in_features, :out_features
+
       def initialize(in_features, out_features, bias: true)
         super()
         @in_features = in_features

data/lib/torch/nn/module.rb

@@ -113,35 +113,53 @@ module Torch
         forward(*input, **kwargs)
       end
 
-      def state_dict(destination: nil)
+      def state_dict(destination: nil, prefix: "")
         destination ||= {}
-        named_parameters.each do |k, v|
-          destination[k] = v
+        save_to_state_dict(destination, prefix: prefix)
+
+        named_children.each do |name, mod|
+          next unless mod
+          mod.state_dict(destination: destination, prefix: prefix + name + ".")
         end
         destination
       end
 
-      # TODO add strict option
-      # TODO match PyTorch behavior
-      def load_state_dict(state_dict)
-        state_dict.each do |k, input_param|
-          k1, k2 = k.split(".", 2)
-          mod = named_modules[k1]
-          if mod.is_a?(Module)
-            param = mod.named_parameters[k2]
-            if param.is_a?(Parameter)
-              Torch.no_grad do
-                param.copy!(input_param)
-              end
-            else
-              raise Error, "Unknown parameter: #{k1}"
-            end
-          else
-            raise Error, "Unknown module: #{k1}"
+      def load_state_dict(state_dict, strict: true)
+        # TODO support strict: false
+        raise "strict: false not implemented yet" unless strict
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+
+        # TODO handle metadata
+
+        _load = lambda do |mod, prefix = ""|
+          # TODO handle metadata
+          local_metadata = {}
+          mod.send(:load_from_state_dict, state_dict, prefix, local_metadata, true, missing_keys, unexpected_keys, error_msgs)
+          mod.named_children.each do |name, child|
+            _load.call(child, prefix + name + ".") unless child.nil?
+          end
+        end
+
+        _load.call(self)
+
+        if strict
+          if unexpected_keys.any?
+            error_msgs << "Unexpected key(s) in state_dict: #{unexpected_keys.join(", ")}"
+          end
+
+          if missing_keys.any?
+            error_msgs << "Missing key(s) in state_dict: #{missing_keys.join(", ")}"
           end
         end
 
-        # TODO return missing keys and unexpected keys
+        if error_msgs.any?
+          # just show first error
+          raise Error, error_msgs[0]
+        end
+
         nil
       end
 
@@ -268,6 +286,12 @@ module Torch
           named_buffers[name]
         elsif named_modules.key?(name)
           named_modules[name]
+        elsif method.end_with?("=") && named_modules.key?(method[0..-2])
+          if instance_variable_defined?("@#{method[0..-2]}")
+            instance_variable_set("@#{method[0..-2]}", *args)
+          else
+            raise NotImplementedYet
+          end
         else
           super
         end
@@ -300,6 +324,68 @@ module Torch
       def dict
         instance_variables.reject { |k| instance_variable_get(k).is_a?(Tensor) }.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
       end
+
+      def load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+        # TODO add hooks
+
+        # TODO handle non-persistent buffers
+        persistent_buffers = named_buffers
+        local_name_params = named_parameters(recurse: false).merge(persistent_buffers)
+        local_state = local_name_params.select { |_, v| !v.nil? }
+
+        local_state.each do |name, param|
+          key = prefix + name
+          if state_dict.key?(key)
+            input_param = state_dict[key]
+
+            # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
+            if param.shape.length == 0 && input_param.shape.length == 1
+              input_param = input_param[0]
+            end
+
+            if input_param.shape != param.shape
+              # local shape should match the one in checkpoint
+              error_msgs << "size mismatch for #{key}: copying a param with shape #{input_param.shape} from checkpoint, " +
+                            "the shape in current model is #{param.shape}."
+              next
+            end
+
+            begin
+              Torch.no_grad do
+                param.copy!(input_param)
+              end
+            rescue => e
+              error_msgs << "While copying the parameter named #{key.inspect}, " +
+                            "whose dimensions in the model are #{param.size} and " +
+                            "whose dimensions in the checkpoint are #{input_param.size}, " +
+                            "an exception occurred: #{e.inspect}"
+            end
+          elsif strict
+            missing_keys << key
+          end
+        end
+
+        if strict
+          state_dict.each_key do |key|
+            if key.start_with?(prefix)
+              input_name = key[prefix.length..-1]
+              input_name = input_name.split(".", 2)[0]
+              if !named_children.key?(input_name) && !local_state.key?(input_name)
+                unexpected_keys << key
+              end
+            end
+          end
+        end
+      end
+
+      def save_to_state_dict(destination, prefix: "")
+        named_parameters(recurse: false).each do |k, v|
+          destination[prefix + k] = v
+        end
+        named_buffers.each do |k, v|
+          destination[prefix + k] = v
+        end
+      end
     end
   end
 end

data/lib/torch/nn/parameter.rb

@@ -3,7 +3,7 @@ module Torch
     class Parameter < Tensor
       def self.new(data = nil, requires_grad: true)
         data = Tensor.new unless data
-        Tensor._make_subclass(data, requires_grad)
+        _make_subclass(data, requires_grad)
       end
 
       def inspect

data/lib/torch/optim/adadelta.rb

@@ -42,11 +42,11 @@ module Torch
               grad = grad.add(p.data, alpha: group[:weight_decay])
             end
 
-            square_avg.mul!(rho).addcmul!(1 - rho, grad, grad)
+            square_avg.mul!(rho).addcmul!(grad, grad, value: 1 - rho)
             std = square_avg.add(eps).sqrt!
             delta = acc_delta.add(eps).sqrt!.div!(std).mul!(grad)
             p.data.add!(delta, alpha: -group[:lr])
-            acc_delta.mul!(rho).addcmul!(1 - rho, delta, delta)
+            acc_delta.mul!(rho).addcmul!(delta, delta, value: 1 - rho)
           end
         end
 

data/lib/torch/optim/adagrad.rb

@@ -57,9 +57,9 @@ module Torch
             if grad.sparse?
               raise NotImplementedYet
             else
-              state[:sum].addcmul!(1, grad, grad)
+              state[:sum].addcmul!(grad, grad, value: 1)
               std = state[:sum].sqrt.add!(group[:eps])
-              p.data.addcdiv!(-clr, grad, std)
+              p.data.addcdiv!(grad, std, value: -clr)
             end
           end
         end

data/lib/torch/optim/adam.rb

@@ -58,7 +58,7 @@ module Torch
 
             # Decay the first and second moment running average coefficient
             exp_avg.mul!(beta1).add!(grad, alpha: 1 - beta1)
-            exp_avg_sq.mul!(beta2).addcmul!(1 - beta2, grad, grad)
+            exp_avg_sq.mul!(beta2).addcmul!(grad, grad, value: 1 - beta2)
             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)
@@ -70,7 +70,7 @@ module Torch
 
             step_size = group[:lr] / bias_correction1
 
-            p.data.addcdiv!(-step_size, exp_avg, denom)
+            p.data.addcdiv!(exp_avg, denom, value: -step_size)
           end
         end
 

data/lib/torch/optim/adamax.rb

@@ -57,7 +57,7 @@ module Torch
             bias_correction = 1 - beta1 ** state[:step]
             clr = group[:lr] / bias_correction
 
-            p.data.addcdiv!(-clr, exp_avg, exp_inf)
+            p.data.addcdiv!(exp_avg, exp_inf, value: -clr)
           end
         end
 

data/lib/torch/optim/adamw.rb

@@ -59,7 +59,7 @@ module Torch
 
             # Decay the first and second moment running average coefficient
             exp_avg.mul!(beta1).add!(grad, alpha: 1 - beta1)
-            exp_avg_sq.mul!(beta2).addcmul!(1 - beta2, grad, grad)
+            exp_avg_sq.mul!(beta2).addcmul!(grad, grad, value: 1 - beta2)
             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)
@@ -71,7 +71,7 @@ module Torch
 
             step_size = group[:lr] / bias_correction1
 
-            p.data.addcdiv!(-step_size, exp_avg, denom)
+            p.data.addcdiv!(exp_avg, denom, value: -step_size)
           end
         end
 

data/lib/torch/optim/rmsprop.rb

@@ -49,7 +49,7 @@ module Torch
               grad = grad.add(p.data, alpha: group[:weight_decay])
             end
 
-            square_avg.mul!(alpha).addcmul!(1 - alpha, grad, grad)
+            square_avg.mul!(alpha).addcmul!(grad, grad, value: 1 - alpha)
 
             if group[:centered]
               grad_avg = state[:grad_avg]
@@ -61,10 +61,10 @@ module Torch
 
             if group[:momentum] > 0
               buf = state[:momentum_buffer]
-              buf.mul!(group[:momentum]).addcdiv!(1, grad, avg)
+              buf.mul!(group[:momentum]).addcdiv!(grad, avg, value: 1)
               p.data.add!(buf, alpha: -group[:lr])
             else
-              p.data.addcdiv!(-group[:lr], grad, avg)
+              p.data.addcdiv!(grad, avg, value: -group[:lr])
             end
           end
         end

data/lib/torch/optim/rprop.rb

@@ -52,7 +52,7 @@ module Torch
             grad[sign.eq(etaminus)] = 0
 
             # update parameters
-            p.data.addcmul!(-1, grad.sign, step_size)
+            p.data.addcmul!(grad.sign, step_size, value: -1)
 
             state[:prev].copy!(grad)
           end

data/lib/torch/tensor.rb

@@ -135,6 +135,10 @@ module Torch
       Torch.ones_like(Torch.empty(*size), **options)
     end
 
+    def requires_grad=(requires_grad)
+      _requires_grad!(requires_grad)
+    end
+
     def requires_grad!(requires_grad = true)
       _requires_grad!(requires_grad)
     end
@@ -174,5 +178,10 @@ module Torch
       return _random!(0, *args) if args.size == 1
       _random!(*args)
     end
+
+    # center option
+    def stft(*args)
+      Torch.stft(*args)
+    end
   end
 end

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

@@ -60,7 +60,7 @@ module Torch
           when Array
             batch.transpose.map { |v| default_convert(v) }
           else
-            raise NotImplementedYet
+            batch
           end
         end
 

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.4.2"
+  VERSION = "0.6.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.4.2
+  version: 0.6.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-10-28 00:00:00.000000000 Z
+date: 2021-03-26 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -24,92 +24,8 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '2.2'
-- !ruby/object:Gem::Dependency
-  name: bundler
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: rake
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: rake-compiler
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: minitest
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '5'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '5'
-- !ruby/object:Gem::Dependency
-  name: numo-narray
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: torchvision
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: 0.1.1
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: 0.1.1
 description:
-email: andrew@chartkick.com
+email: andrew@ankane.org
 executables: []
 extensions:
 - ext/torch/extconf.rb
@@ -121,13 +37,20 @@ files:
 - codegen/function.rb
 - codegen/generate_functions.rb
 - codegen/native_functions.yaml
+- ext/torch/cuda.cpp
+- ext/torch/device.cpp
 - ext/torch/ext.cpp
 - ext/torch/extconf.rb
+- ext/torch/ivalue.cpp
+- ext/torch/nn.cpp
 - ext/torch/nn_functions.h
+- ext/torch/random.cpp
 - ext/torch/ruby_arg_parser.cpp
 - ext/torch/ruby_arg_parser.h
 - ext/torch/templates.h
+- ext/torch/tensor.cpp
 - ext/torch/tensor_functions.h
+- ext/torch/torch.cpp
 - ext/torch/torch_functions.h
 - ext/torch/utils.h
 - ext/torch/wrap_outputs.h
@@ -282,14 +205,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '2.4'
+      version: '2.6'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.4
+rubygems_version: 3.2.3
 signing_key:
 specification_version: 4
 summary: Deep learning for Ruby, powered by LibTorch