torch-rb 0.4.1 → 0.5.3

data/ext/torch/ext.cpp

@@ -44,8 +44,13 @@ std::vector<TensorIndex> index_vector(Array a) {
     if (obj.is_instance_of(rb_cInteger)) {
       indices.push_back(from_ruby<int64_t>(obj));
     } else if (obj.is_instance_of(rb_cRange)) {
-      torch::optional<int64_t> start_index = from_ruby<int64_t>(obj.call("begin"));
-      torch::optional<int64_t> stop_index = -1;
+      torch::optional<int64_t> start_index = torch::nullopt;
+      torch::optional<int64_t> stop_index = torch::nullopt;
+
+      Object begin = obj.call("begin");
+      if (!begin.is_nil()) {
+        start_index = from_ruby<int64_t>(begin);
+      }
 
       Object end = obj.call("end");
       if (!end.is_nil()) {
@@ -53,12 +58,14 @@ std::vector<TensorIndex> index_vector(Array a) {
       }
 
       Object exclude_end = obj.call("exclude_end?");
-      if (!exclude_end) {
+      if (stop_index.has_value() && !exclude_end) {
         if (stop_index.value() == -1) {
           stop_index = torch::nullopt;
         } else {
           stop_index = stop_index.value() + 1;
         }
+      } else if (!stop_index.has_value() && exclude_end) {
+        stop_index = -1;
       }
 
       indices.push_back(torch::indexing::Slice(start_index, stop_index));
@@ -348,16 +355,6 @@ void Init_ext()
       *[](Tensor& self) {
         return self.is_contiguous();
       })
-    .define_method(
-      "addcmul!",
-      *[](Tensor& self, Scalar value, const Tensor & tensor1, const Tensor & tensor2) {
-        return self.addcmul_(tensor1, tensor2, value);
-      })
-    .define_method(
-      "addcdiv!",
-      *[](Tensor& self, Scalar value, const Tensor & tensor1, const Tensor & tensor2) {
-        return self.addcdiv_(tensor1, tensor2, value);
-      })
     .define_method(
       "_requires_grad!",
       *[](Tensor& self, bool requires_grad) {
@@ -372,7 +369,7 @@ void Init_ext()
     .define_method(
       "grad=",
       *[](Tensor& self, torch::Tensor& grad) {
-        self.grad() = grad;
+        self.mutable_grad() = grad;
       })
     .define_method(
       "_dtype",
@@ -609,7 +606,7 @@ void Init_ext()
     .define_method(
       "grad=",
       *[](Parameter& self, torch::Tensor& grad) {
-        self.grad() = grad;
+        self.mutable_grad() = grad;
       });
 
   Class rb_cDevice = define_class_under<torch::Device>(rb_mTorch, "Device")
@@ -628,5 +625,7 @@ void Init_ext()
   Module rb_mCUDA = define_module_under(rb_mTorch, "CUDA")
     .add_handler<torch::Error>(handle_error)
     .define_singleton_method("available?", &torch::cuda::is_available)
-    .define_singleton_method("device_count", &torch::cuda::device_count);
+    .define_singleton_method("device_count", &torch::cuda::device_count)
+    .define_singleton_method("manual_seed", &torch::cuda::manual_seed)
+    .define_singleton_method("manual_seed_all", &torch::cuda::manual_seed_all);
 }

data/ext/torch/ruby_arg_parser.h

@@ -91,7 +91,7 @@ struct RubyArgs {
   inline c10::optional<int64_t> toInt64Optional(int i);
   inline c10::optional<bool> toBoolOptional(int i);
   inline c10::optional<double> toDoubleOptional(int i);
-  // inline c10::OptionalArray<double> doublelistOptional(int i);
+  inline c10::OptionalArray<double> doublelistOptional(int i);
   // inline at::Layout layout(int i);
   // inline at::Layout layoutWithDefault(int i, at::Layout default_layout);
   inline c10::optional<at::Layout> layoutOptional(int i);
@@ -105,7 +105,7 @@ struct RubyArgs {
   inline c10::optional<at::MemoryFormat> memoryformatOptional(int i);
   // inline at::QScheme toQScheme(int i);
   inline std::string string(int i);
-  // inline c10::optional<std::string> stringOptional(int i);
+  inline c10::optional<std::string> stringOptional(int i);
   // inline PyObject* pyobject(int i);
   inline int64_t toInt64(int i);
   // inline int64_t toInt64WithDefault(int i, int64_t default_int);
@@ -249,6 +249,25 @@ inline c10::optional<double> RubyArgs::toDoubleOptional(int i) {
   return toDouble(i);
 }
 
+inline c10::OptionalArray<double> RubyArgs::doublelistOptional(int i) {
+  if (NIL_P(args[i])) return {};
+
+  VALUE arg = args[i];
+  auto size = RARRAY_LEN(arg);
+  std::vector<double> res(size);
+  for (idx = 0; idx < size; idx++) {
+    VALUE obj = rb_ary_entry(arg, idx);
+    if (FIXNUM_P(obj) || RB_FLOAT_TYPE_P(obj)) {
+      res[idx] = from_ruby<double>(obj);
+    } else {
+      rb_raise(rb_eArgError, "%s(): argument '%s' must be %s, but found element of type %s at pos %d",
+          signature.name.c_str(), signature.params[i].name.c_str(),
+          signature.params[i].type_name().c_str(), rb_obj_classname(obj), idx + 1);
+    }
+  }
+  return res;
+}
+
 inline c10::optional<at::Layout> RubyArgs::layoutOptional(int i) {
   if (NIL_P(args[i])) return c10::nullopt;
 
@@ -285,6 +304,11 @@ inline std::string RubyArgs::string(int i) {
   return from_ruby<std::string>(args[i]);
 }
 
+inline c10::optional<std::string> RubyArgs::stringOptional(int i) {
+  if (!args[i]) return c10::nullopt;
+  return from_ruby<std::string>(args[i]);
+}
+
 inline int64_t RubyArgs::toInt64(int i) {
   if (NIL_P(args[i])) return signature.params[i].default_int;
   return from_ruby<int64_t>(args[i]);

data/ext/torch/templates.h

@@ -19,6 +19,7 @@ using torch::TensorOptions;
 using torch::Layout;
 using torch::MemoryFormat;
 using torch::IntArrayRef;
+using torch::ArrayRef;
 using torch::TensorList;
 using torch::Storage;
 

data/ext/torch/wrap_outputs.h

@@ -90,3 +90,10 @@ inline Object wrap(torch::TensorList x) {
   }
   return Object(a);
 }
+
+inline Object wrap(std::tuple<double, double> x) {
+  Array a;
+  a.push(to_ruby<double>(std::get<0>(x)));
+  a.push(to_ruby<double>(std::get<1>(x)));
+  return Object(a);
+}

data/lib/torch.rb

@@ -261,6 +261,8 @@ module Torch
         Torch._from_blob(bytes, [bytes.bytesize], TensorOptions.new.dtype(DTYPE_TO_ENUM[dtype]))
       elsif args.size == 1 && args.first.is_a?(Array)
         Torch.tensor(args.first, dtype: dtype, device: device)
+      elsif args.size == 0
+        Torch.empty(0, dtype: dtype, device: device)
       else
         Torch.empty(*args, dtype: dtype, device: device)
       end
@@ -434,7 +436,8 @@ module Torch
       zeros(input.size, **like_options(input, options))
     end
 
-    def stft(input, n_fft, hop_length: nil, win_length: nil, window: nil, center: true, pad_mode: "reflect", normalized: false, onesided: true)
+    # center option
+    def stft(input, n_fft, hop_length: nil, win_length: nil, window: nil, center: true, pad_mode: "reflect", normalized: false, onesided: true, return_complex: nil)
       if center
         signal_dim = input.dim
         extended_shape = [1] * (3 - signal_dim) + input.size
@@ -442,12 +445,7 @@ module Torch
         input = NN::F.pad(input.view(extended_shape), [pad, pad], mode: pad_mode)
         input = input.view(input.shape[-signal_dim..-1])
       end
-      _stft(input, n_fft, hop_length, win_length, window, normalized, onesided)
-    end
-
-    def clamp(tensor, min, max)
-      tensor = _clamp_min(tensor, min)
-      _clamp_max(tensor, max)
+      _stft(input, n_fft, hop_length, win_length, window, normalized, onesided, return_complex)
     end
 
     private

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
 
-        # TODO return missing keys and unexpected keys
+        _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
+
+        if error_msgs.any?
+          # just show first error
+          raise Error, error_msgs[0]
+        end
+
         nil
       end
 
@@ -300,6 +318,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/optim/adadelta.rb

@@ -39,14 +39,14 @@ module Torch
             state[:step] += 1
 
             if group[:weight_decay] != 0
-              grad = grad.add(group[:weight_decay], p.data)
+              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

@@ -49,7 +49,7 @@ module Torch
               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)
+              grad = grad.add(p.data, alpha: group[:weight_decay])
             end
 
             clr = group[:lr] / (1 + (state[:step] - 1) * group[:lr_decay])
@@ -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

@@ -42,7 +42,7 @@ module Torch
             state[:step] += 1
 
             if group[:weight_decay] != 0
-              grad = grad.add(group[:weight_decay], p.data)
+              grad = grad.add(p.data, alpha: group[:weight_decay])
             end
 
             # Update biased first moment estimate.
@@ -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/asgd.rb

@@ -36,7 +36,7 @@ module Torch
             state[:step] += 1
 
             if group[:weight_decay] != 0
-              grad = grad.add(group[:weight_decay], p.data)
+              grad = grad.add(p.data, alpha: group[:weight_decay])
             end
 
             # decay term

data/lib/torch/optim/rmsprop.rb

@@ -46,25 +46,25 @@ module Torch
             state[:step] += 1
 
             if group[:weight_decay] != 0
-              grad = grad.add(group[:weight_decay], p.data)
+              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]
-              grad_avg.mul!(alpha).add!(1 - alpha, grad)
-              avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt!.add!(group[:eps])
+              grad_avg.mul!(alpha).add!(grad, alpha: 1 - alpha)
+              avg = square_avg.addcmul(grad_avg, grad_avg, value: -1).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)
+              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