torch-rb 0.5.2 → 0.5.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: af82b518434a574d47ba1e9d4c785e4d295622b0092f50e6d638a88396601ef1
-  data.tar.gz: dbc5b4ad9e5fd1a907688f12feb9e1a01d8bab852b68a1712cd70d655ed4833f
+  metadata.gz: fa1967c6330c6ad818e5452d38c5f83b2dde901c33d03f2a442bc003e4e662be
+  data.tar.gz: 00d57cb1de4e6cec72986cab301b04b66ce53aa1d7a0e84eadbc25ac924d28eb
 SHA512:
-  metadata.gz: 8e372efecb07c2120e03c04e96e9b9d9fd8e2242b35527e101f545397cbd9649b1fd3ccaddaf42b0cdfb6c6f8aa70d529c828a9a9a8fa248bed1ff1bc83d7d5a
-  data.tar.gz: a8ddbd40b7f0b70ee904098285d1c1c02ca51e2d5be46aec7d423985b04f085e9f92aa5d7f1e9dd5785b22898277ff0739b01ab6bee32235273f70c5fd76c74c
+  metadata.gz: ae015b2ea53033df631d914c7e6fc70bdff216742900d6a9a76d5e4479cc8faf05becf2ab8ee9edc874092e36355afecbe6c08f5205678be6ec01f0f0783a469
+  data.tar.gz: c751eaa8253682bc95653aa45cc61a05c9f120fa6df88c3d69873a43f47ca7da07e45c3a1e1c5f36bc63da69fae0dc9626f9db2f20af9d6ef190dbcce60bfb75

data/CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.5.3 (2021-01-14)
+
+- Added `manual_seed` and `manual_seed_all` for CUDA
+- Improved saving and loading models
+- Fixed error with tensor indexing with beginless ranges in Ruby 3.0
+
 ## 0.5.2 (2020-10-29)
 
 - Fixed `undefined symbol` error with CUDA

data/README.md

@@ -8,7 +8,7 @@ Check out:
 - [TorchText](https://github.com/ankane/torchtext) for text and NLP tasks
 - [TorchAudio](https://github.com/ankane/torchaudio) for audio tasks
 
-[![Build Status](https://travis-ci.org/ankane/torch.rb.svg?branch=master)](https://travis-ci.org/ankane/torch.rb)
+[![Build Status](https://github.com/ankane/torch.rb/workflows/build/badge.svg?branch=master)](https://github.com/ankane/torch.rb/actions)
 
 ## Installation
 
@@ -319,6 +319,13 @@ net.load_state_dict(Torch.load("net.pth"))
 net.eval
 ```
 
+When saving a model in Python to load in Ruby, convert parameters to tensors (due to outstanding bugs in LibTorch)
+
+```python
+state_dict = {k: v.data if isinstance(v, torch.nn.Parameter) else v for k, v in state_dict.items()}
+torch.save(state_dict, "net.pth")
+```
+
 ### Tensor Creation
 
 Here’s a list of functions to create tensors (descriptions from the [C++ docs](https://pytorch.org/cppdocs/notes/tensor_creation.html)):
@@ -416,7 +423,7 @@ Here’s the list of compatible versions.
 
 Torch.rb | LibTorch
 --- | ---
-0.5.0+ | 1.7.0
+0.5.0+ | 1.7.0-1.7.1
 0.3.0+ | 1.6.0
 0.2.0-0.2.7 | 1.5.0-1.5.1
 0.1.8 | 1.4.0

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));
@@ -618,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/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/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.5.2"
+  VERSION = "0.5.3"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.5.2
+  version: 0.5.3
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-10-29 00:00:00.000000000 Z
+date: 2021-01-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -289,7 +289,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !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