Merge pull request #1173 from shaltielshmid/optimizer-load-clone-tensors

NiklasGustafsson · web-flow · commit c7aab8c6ecb5 · 2023-12-07T08:45:55.000-08:00
Fix optimizer load state dict copy tensor by reference
diff --git a/RELEASENOTES.md b/RELEASENOTES.md
@@ -22,6 +22,8 @@ __Bug Fixes__:
 
 #1154 : `mu_product` was not initialized in `NAdam` optimizer
 #1170 : Calling `torch.nn.rnn.utils.pad_packed_sequence` with a CUDA tensor and unsorted_indices threw an error
+#1172 : `optim.LoadStateDict` from an existing `StateDictionary` updated to make sure to copy value and to the right device.
+#1176 : When specific `Optimizers` load in a conditional tensor, made sure to copy to the right device.
 
 ## NuGet Version 0.101.2
 
diff --git a/src/TorchSharp/Optimizers/ASGD.cs b/src/TorchSharp/Optimizers/ASGD.cs
@@ -241,7 +241,7 @@ public override void LoadStateDict(OptimizerState source)
                     eta = st_state.eta;
                     mu = st_state.mu;
                     ax.Dispose();
-                    ax = st_state.ax;
+                    ax = st_state.ax.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Optimizers/Adadelta.cs b/src/TorchSharp/Optimizers/Adadelta.cs
@@ -237,8 +237,8 @@ public override void LoadStateDict(OptimizerState source)
                     acc_delta.Dispose();
 
                     step = st_state.step;
-                    square_avg = st_state.square_avg;
-                    acc_delta = st_state.acc_delta;
+                    square_avg = st_state.square_avg.to(_parameter.device, copy: true);
+                    acc_delta = st_state.acc_delta.to(_parameter.device, copy: true);
                 }
 
                 public override bool ApproximatelyEquals(OptimizerState other)
diff --git a/src/TorchSharp/Optimizers/Adagrad.cs b/src/TorchSharp/Optimizers/Adagrad.cs
@@ -231,7 +231,7 @@ public override void LoadStateDict(OptimizerState source)
                     var st_state = source as State;
                     sum.Dispose();
                     step = st_state.step;
-                    sum = st_state.sum;
+                    sum = st_state.sum.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Optimizers/Adam.cs b/src/TorchSharp/Optimizers/Adam.cs
@@ -250,13 +250,7 @@ public override void LoadStateDict(BinaryReader reader)
                     step = reader.ReadInt64();
                     exp_avg.Load(reader);
                     exp_avg_sq.Load(reader);
-                    var hasMax = reader.ReadBoolean();
-                    if (hasMax) {
-                        TensorExtensionMethods.Load(ref max_exp_avg_sq, reader);
-                    } else {
-                        max_exp_avg_sq?.Dispose();
-                        max_exp_avg_sq = null;
-                    }
+                    LoadConditionalStateTensor(reader, ref max_exp_avg_sq, _parameter.device);
                 }
 
                 /// <summary>
@@ -285,14 +279,12 @@ public override void LoadStateDict(OptimizerState source)
                     var st_state = source as State;
                     exp_avg.Dispose();
                     exp_avg_sq.Dispose();
-                    if (max_exp_avg_sq is not null) {
-                        max_exp_avg_sq.Dispose();
-                    }
-
+                    max_exp_avg_sq?.Dispose();
+                    
                     step = st_state.step;
-                    exp_avg = st_state.exp_avg;
-                    exp_avg_sq = st_state.exp_avg_sq;
-                    max_exp_avg_sq = st_state.max_exp_avg_sq;
+                    exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);
+                    exp_avg_sq = st_state.exp_avg_sq.to(_parameter.device, copy: true);
+                    max_exp_avg_sq = st_state.max_exp_avg_sq?.to(_parameter.device, copy: true);
                 }
 
                 public override bool ApproximatelyEquals(OptimizerState other)
diff --git a/src/TorchSharp/Optimizers/AdamW.cs b/src/TorchSharp/Optimizers/AdamW.cs
@@ -286,14 +286,12 @@ public override void LoadStateDict(OptimizerState source)
                     var st_state = source as State;
                     exp_avg.Dispose();
                     exp_avg_sq.Dispose();
-                    if (max_exp_avg_sq is not null) {
-                        max_exp_avg_sq.Dispose();
-                    }
-
+                    max_exp_avg_sq?.Dispose();
+                    
                     step = st_state.step;
-                    exp_avg = st_state.exp_avg;
-                    exp_avg_sq = st_state.exp_avg_sq;
-                    max_exp_avg_sq = st_state.max_exp_avg_sq;
+                    exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);
+                    exp_avg_sq = st_state.exp_avg_sq.to(_parameter.device, copy: true);
+                    max_exp_avg_sq = st_state.max_exp_avg_sq?.to(_parameter.device, copy: true);
                 }
 
                 public override bool ApproximatelyEquals(OptimizerState other)
diff --git a/src/TorchSharp/Optimizers/Adamax.cs b/src/TorchSharp/Optimizers/Adamax.cs
@@ -255,8 +255,8 @@ public override void LoadStateDict(OptimizerState source)
                     exp_inf.Dispose();
 
                     step = st_state.step;
-                    exp_avg = st_state.exp_avg;
-                    exp_inf = st_state.exp_inf;
+                    exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);
+                    exp_inf = st_state.exp_inf.to(_parameter.device, copy: true);
                 }
 
                 public override bool ApproximatelyEquals(OptimizerState other)
diff --git a/src/TorchSharp/Optimizers/NAdam.cs b/src/TorchSharp/Optimizers/NAdam.cs
@@ -269,8 +269,8 @@ public override void LoadStateDict(OptimizerState source)
 
                     step = st_state.step;
                     mu_product = st_state.mu_product;
-                    exp_avg = st_state.exp_avg;
-                    exp_avg_sq = st_state.exp_avg_sq;
+                    exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);
+                    exp_avg_sq = st_state.exp_avg_sq.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Optimizers/Optimizer.cs b/src/TorchSharp/Optimizers/Optimizer.cs
@@ -555,12 +555,13 @@ public virtual bool ApproximatelyEquals(OptimizerState other)
             /// <param name="device">The device to move all state to.</param>
             public virtual void to(Device device) { }
 
-            protected static void LoadConditionalStateTensor(BinaryReader reader, ref Tensor result)
+            protected static void LoadConditionalStateTensor(BinaryReader reader, ref Tensor result, Device device)
             {
                 var hasTensor = reader.ReadBoolean();
 
                 if (hasTensor) {
                     TensorExtensionMethods.Load(ref result, reader);
+                    result = result.to(device, disposeAfter: true);
                 } else {
                     if (result is not null)
                         result.Dispose();
diff --git a/src/TorchSharp/Optimizers/RAdam.cs b/src/TorchSharp/Optimizers/RAdam.cs
@@ -262,8 +262,8 @@ public override void LoadStateDict(OptimizerState source)
                     exp_avg_sq.Dispose();
 
                     step = st_state.step;
-                    exp_avg = st_state.exp_avg;
-                    exp_avg_sq = st_state.exp_avg_sq;
+                    exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);
+                    exp_avg_sq = st_state.exp_avg_sq.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Optimizers/RMSprop.cs b/src/TorchSharp/Optimizers/RMSprop.cs
@@ -251,8 +251,8 @@ public override void LoadStateDict(BinaryReader reader)
                 {
                     step = reader.ReadInt64();
                     square_avg.Load(reader);
-                    LoadConditionalStateTensor(reader, ref momentum_buffer);
-                    LoadConditionalStateTensor(reader, ref grad_avg);
+                    LoadConditionalStateTensor(reader, ref momentum_buffer, _parameter.device);
+                    LoadConditionalStateTensor(reader, ref grad_avg, _parameter.device);
                 }
                 /// <summary>
                 /// Save the optimizer parameter state to a stream.
@@ -278,9 +278,9 @@ public override void LoadStateDict(OptimizerState source)
                     momentum_buffer.Dispose();
 
                     step = st_state.step;
-                    square_avg = st_state.square_avg;
-                    grad_avg = st_state.grad_avg;
-                    momentum_buffer = st_state.momentum_buffer;
+                    square_avg = st_state.square_avg.to(_parameter.device, copy: true);
+                    grad_avg = st_state.grad_avg.to(_parameter.device, copy: true);
+                    momentum_buffer = st_state.momentum_buffer.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Optimizers/Rprop.cs b/src/TorchSharp/Optimizers/Rprop.cs
@@ -285,8 +285,8 @@ public override void LoadStateDict(OptimizerState source)
                     prev.Dispose();
                     step_size.Dispose();
                     step = st_state.step;
-                    prev = st_state.prev;
-                    step_size = st_state.step_size;
+                    prev = st_state.prev.to(_parameter.device, copy: true);
+                    step_size = st_state.step_size.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Optimizers/SGD.cs b/src/TorchSharp/Optimizers/SGD.cs
@@ -232,7 +232,7 @@ public override void to(Device device)
                 /// <param name="reader">A binary reader connected to a stream open for reading.</param>
                 public override void LoadStateDict(BinaryReader reader)
                 {
-                    LoadConditionalStateTensor(reader, ref momentum_buffer);
+                    LoadConditionalStateTensor(reader, ref momentum_buffer, _parameter.device);
                 }
 
                 /// <summary>
@@ -251,10 +251,8 @@ public override void SaveStateDict(BinaryWriter writer)
                 public override void LoadStateDict(OptimizerState source)
                 {
                     var st_state = source as State;
-                    if (momentum_buffer is not null) {
-                        momentum_buffer.Dispose();
-                    }
-                    momentum_buffer = st_state.momentum_buffer;
+                    momentum_buffer?.Dispose();
+                    momentum_buffer = st_state.momentum_buffer.to(_parameter.device, copy: true);
                 }
 
                 /// <summary>
diff --git a/src/TorchSharp/Tensor/Tensor.cs b/src/TorchSharp/Tensor/Tensor.cs
@@ -713,11 +713,15 @@ public Tensor cuda(Device? device = null)
             /// </summary>
             /// <param name="type">The target type</param>
             /// <param name="copy">When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.</param>
-            public Tensor to_type(ScalarType type, bool copy = false)
+            /// <param name="disposeAfter">When disposeAfter is set, the current Tensor will be disposed after creating the new one</param>
+            public Tensor to_type(ScalarType type, bool copy = false, bool disposeAfter = false)
             {
                 var res = NativeMethods.THSTensor_to_type(Handle, (sbyte)type, copy);
                 if (res == IntPtr.Zero)
                     CheckForErrors();
+                if (disposeAfter)
+                    this.Dispose();
+
                 return new Tensor(res);
             }
 
@@ -743,12 +747,16 @@ public Tensor set_(Tensor source)
             /// <param name="deviceType">The device type, e.g. 'CPU' or 'CUDA'.</param>
             /// <param name="deviceIndex">The optional device index.</param>
             /// <param name="copy">When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.</param>
-            public Tensor to(DeviceType deviceType, int deviceIndex = -1, bool copy = false)
+            /// <param name="disposeAfter">When disposeAfter is set, the current Tensor will be disposed after creating the new one</param>
+            public Tensor to(DeviceType deviceType, int deviceIndex = -1, bool copy = false, bool disposeAfter = false)
             {
                 torch.InitializeDeviceType(deviceType);
                 var res = NativeMethods.THSTensor_to_device(Handle, (int)deviceType, deviceIndex, copy);
                 if (res == IntPtr.Zero)
                     CheckForErrors();
+                if (disposeAfter)
+                    this.Dispose();
+
                 return new Tensor(res);
             }
 
@@ -758,32 +766,42 @@ public Tensor to(DeviceType deviceType, int deviceIndex = -1, bool copy = false)
             /// <param name="type">The target type</param>
             /// <param name="device">The target device</param>
             /// <param name="copy">When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.</param>
-            public Tensor to(ScalarType type, torch.Device device, bool copy = false)
+            /// <param name="disposeAfter">When disposeAfter is set, the current Tensor will be disposed after creating the new one</param>
+            public Tensor to(ScalarType type, torch.Device device, bool copy = false, bool disposeAfter = false)
             {
                 torch.InitializeDevice(device);
                 var res = NativeMethods.THSTensor_to_type_and_device(Handle, (sbyte)type, (int)device.type, device.index, copy);
                 if (res == IntPtr.Zero)
                     CheckForErrors();
+                if (disposeAfter)
+                    this.Dispose();
                 return new Tensor(res);
             }
 
             /// <summary>
             /// Cast the tensor to the given element type.
             /// </summary>
+            /// <param name="type">The target type</param>
+            /// <param name="copy">When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.</param>
+            /// <param name="disposeAfter">When disposeAfter is set, the current Tensor will be disposed after creating the new one</param>
             /// <remarks>Alias for to_type</remarks>
-            public Tensor to(ScalarType type) => to_type(type);
+            public Tensor to(ScalarType type, bool copy = false, bool disposeAfter = false) => to_type(type, copy, disposeAfter);
 
             /// <summary>
             /// Moves the tensor data.
             /// </summary>
             /// <param name="device">A string denoting the target device.</param>
-            public Tensor to(string device) => to(new torch.Device(device));
+            /// <param name="copy">When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.</param>
+            /// <param name="disposeAfter">When disposeAfter is set, the current Tensor will be disposed after creating the new one</param>
+            public Tensor to(string device, bool copy = false, bool disposeAfter = false) => to(new torch.Device(device), copy, disposeAfter);
 
             /// <summary>
             /// Moves the tensor data.
             /// </summary>
             /// <param name="device">The target device</param>
-            public Tensor to(torch.Device device) => to(device.type, device.index);
+            /// <param name="copy">When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion.</param>
+            /// <param name="disposeAfter">When disposeAfter is set, the current Tensor will be disposed after creating the new one</param>
+            public Tensor to(torch.Device device, bool copy = false, bool disposeAfter = false) => to(device.type, device.index, copy, disposeAfter);
 
             /// <summary>
             /// Moves the tensor data.
diff --git a/test/TorchSharpTest/TestTorchTensor.cs b/test/TorchSharpTest/TestTorchTensor.cs
@@ -3959,6 +3959,59 @@ public void Meta()
             Assert.Equal(x.shape, z.shape);
         }
 
+        [Fact]
+        [TestOf(nameof(Tensor.to))]
+        public void CastMoveAndDisposeAfter()
+        {
+            {
+                // Cast the input on the same device
+                using var input = torch.ones(10, float32, torch.CPU);
+                using var cast = input.to(int32, disposeAfter: true);
+                Assert.True(input.IsInvalid);
+                Assert.False(cast.IsInvalid);
+                // make sure we can access the values
+                Assert.Equal(1, cast[0].ToInt32()); 
+            }
+            if (torch.cuda.is_available()) {
+                {
+                    // Move the input to a different device
+                    using var input = torch.ones(10, float32, torch.CPU);
+                    using var moved = input.to(torch.CUDA, disposeAfter: true);
+                    Assert.True(input.IsInvalid);
+                    Assert.False(moved.IsInvalid);
+                    // make sure we can access the values
+                    Assert.Equal(1, moved[0].ToSingle());
+                }
+                {
+                    // Cast and move the input to a different device
+                    using var input = torch.ones(10, float32, torch.CPU);
+                    using var moved = input.to(int32, torch.CUDA, disposeAfter: true);
+                    Assert.True(input.IsInvalid);
+                    Assert.False(moved.IsInvalid);
+                    // make sure we can access the values
+                    Assert.Equal(1, moved[0].ToInt32());
+                }
+            }
+            {
+                // Sanity: If we cast to the same type, values should still be accessible
+                using var input = torch.ones(10, float32, torch.CPU);
+                using var cast = input.to(float32, disposeAfter: true);
+                Assert.True(input.IsInvalid);
+                Assert.False(cast.IsInvalid);
+                // make sure we can access the values
+                Assert.Equal(1, cast[0].ToSingle());
+            }
+            {
+                // Sanity: If we move to the same device, values should still be accessible
+                using var input = torch.ones(10, float32, torch.CPU);
+                using var moved = input.to(torch.CPU, disposeAfter: true);
+                Assert.True(input.IsInvalid);
+                Assert.False(moved.IsInvalid);
+                // make sure we can access the values
+                Assert.Equal(1, moved[0].ToSingle());
+            }
+        }
+
         [Fact]
         [TestOf(nameof(Tensor.masked_scatter))]
         [TestOf(nameof(Tensor.masked_scatter_))]
diff --git a/test/TorchSharpTest/TestTorchTensorBugs.cs b/test/TorchSharpTest/TestTorchTensorBugs.cs
@@ -17,6 +17,7 @@
 using System.Numerics;
 using System.Reflection.Metadata;
 using System.Runtime.InteropServices;
+using TorchSharp.Modules;
 
 #nullable enable
 
@@ -1423,5 +1424,37 @@ public void Validate1170()
                 var error = torch.nn.utils.rnn.pad_packed_sequence(packed);
             }
         }
+
+        [Fact]
+        public void Validate1172_Clone()
+        {
+            var lin1 = torch.nn.Linear(10, 10);
+
+            var optim1 = torch.optim.Adam(lin1.parameters());
+            var optim2 = torch.optim.Adam(lin1.parameters());
+            optim2.load_state_dict(optim1.state_dict());
+            optim1.Dispose();
+
+            var state = (optim2.state_dict().State[0] as Adam.State)!;
+            Assert.False(state.exp_avg.IsInvalid);
+            Assert.False(state.exp_avg_sq.IsInvalid);
+        }
+
+        [Fact]
+        public void Validate1172_Device()
+        {
+            if (torch.cuda.is_available()) {
+                var lin1 = torch.nn.Linear(10, 10);
+                var optim1 = torch.optim.Adam(lin1.parameters());
+                var sd = optim1.state_dict();
+
+                lin1.cuda();
+                var optim2 = torch.optim.Adam(lin1.parameters());
+                Assert.Equal(DeviceType.CUDA, (optim2.state_dict().State[0] as Adam.State)!.exp_avg.device.type);
+
+                optim2.load_state_dict(sd);
+                Assert.Equal(DeviceType.CUDA, (optim2.state_dict().State[0] as Adam.State)!.exp_avg.device.type);
+            }
+        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -241,7 +241,7 @@ public override void LoadStateDict(OptimizerState source)`
`241`	`241`	`eta = st_state.eta;`
`242`	`242`	`mu = st_state.mu;`
`243`	`243`	`ax.Dispose();`
`244`		`- ax = st_state.ax;`
	`244`	`+ ax = st_state.ax.to(_parameter.device, copy: true);`
`245`	`245`	`}`
`246`	`246`
`247`	`247`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -237,8 +237,8 @@ public override void LoadStateDict(OptimizerState source)`
`237`	`237`	`acc_delta.Dispose();`
`238`	`238`
`239`	`239`	`step = st_state.step;`
`240`		`- square_avg = st_state.square_avg;`
`241`		`- acc_delta = st_state.acc_delta;`
	`240`	`+ square_avg = st_state.square_avg.to(_parameter.device, copy: true);`
	`241`	`+ acc_delta = st_state.acc_delta.to(_parameter.device, copy: true);`
`242`	`242`	`}`
`243`	`243`
`244`	`244`	`public override bool ApproximatelyEquals(OptimizerState other)`
Original file line number	Diff line number	Diff line change
`@@ -231,7 +231,7 @@ public override void LoadStateDict(OptimizerState source)`
`231`	`231`	`var st_state = source as State;`
`232`	`232`	`sum.Dispose();`
`233`	`233`	`step = st_state.step;`
`234`		`- sum = st_state.sum;`
	`234`	`+ sum = st_state.sum.to(_parameter.device, copy: true);`
`235`	`235`	`}`
`236`	`236`
`237`	`237`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -255,8 +255,8 @@ public override void LoadStateDict(OptimizerState source)`
`255`	`255`	`exp_inf.Dispose();`
`256`	`256`
`257`	`257`	`step = st_state.step;`
`258`		`- exp_avg = st_state.exp_avg;`
`259`		`- exp_inf = st_state.exp_inf;`
	`258`	`+ exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);`
	`259`	`+ exp_inf = st_state.exp_inf.to(_parameter.device, copy: true);`
`260`	`260`	`}`
`261`	`261`
`262`	`262`	`public override bool ApproximatelyEquals(OptimizerState other)`
Original file line number	Diff line number	Diff line change
`@@ -269,8 +269,8 @@ public override void LoadStateDict(OptimizerState source)`
`269`	`269`
`270`	`270`	`step = st_state.step;`
`271`	`271`	`mu_product = st_state.mu_product;`
`272`		`- exp_avg = st_state.exp_avg;`
`273`		`- exp_avg_sq = st_state.exp_avg_sq;`
	`272`	`+ exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);`
	`273`	`+ exp_avg_sq = st_state.exp_avg_sq.to(_parameter.device, copy: true);`
`274`	`274`	`}`
`275`	`275`
`276`	`276`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -262,8 +262,8 @@ public override void LoadStateDict(OptimizerState source)`
`262`	`262`	`exp_avg_sq.Dispose();`
`263`	`263`
`264`	`264`	`step = st_state.step;`
`265`		`- exp_avg = st_state.exp_avg;`
`266`		`- exp_avg_sq = st_state.exp_avg_sq;`
	`265`	`+ exp_avg = st_state.exp_avg.to(_parameter.device, copy: true);`
	`266`	`+ exp_avg_sq = st_state.exp_avg_sq.to(_parameter.device, copy: true);`
`267`	`267`	`}`
`268`	`268`
`269`	`269`	`/// <summary>`