From 471c05d5e87a84ae63125c79647a00dcc2b3a5c2 Mon Sep 17 00:00:00 2001 From: SeoHyungjun Date: Fri, 4 Oct 2024 17:17:00 +0900 Subject: [PATCH] [MachineLearning.Train] Change internal logic to resolve a memory leak There were issues with heap use after free and heap overflow occurring (TNINE-4152, 4153, 4154). In order to solve this problem, modified the internal logic of the API. Signed-off-by: SeoHyungjun --- .../Tizen.MachineLearning.Train/Dataset.cs | 20 ++++++++++++++++++- .../Tizen.MachineLearning.Train/Layer.cs | 5 ++++- .../Tizen.MachineLearning.Train/Model.cs | 7 ++++++- .../Tizen.MachineLearning.Train/Optimizer.cs | 20 ++++++++++++++++++- 4 files changed, 48 insertions(+), 4 deletions(-) diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Dataset.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Dataset.cs index cfaf094c2a0..5658d4faa8e 100644 --- a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Dataset.cs +++ b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Dataset.cs @@ -37,6 +37,9 @@ public class Dataset: IDisposable private IntPtr handle = IntPtr.Zero; private bool disposed = false; + /// if false, model will be destroy dataset handle + private bool hasOwnership = true; + /// /// Constructs the dataset. /// @@ -44,6 +47,9 @@ public class Dataset: IDisposable public Dataset() { NNTrainerError ret = Interop.Dataset.Create(out handle); + if (ret != NNTrainerError.None) { + handle = IntPtr.Zero; + } NNTrainer.CheckException(ret, "Failed to create dataset instance"); Log.Info(NNTrainer.Tag, "Create Dataset"); } @@ -82,6 +88,14 @@ protected virtual void Dispose(bool disposing) { // release managed object } + + disposed = true; + + if (!hasOwnership){ + Log.Info(NNTrainer.Tag, "Cannot destroy dataset already added in a Model. Model will destroy this dataset"); + return; + } + // release unmanaged object if (handle != IntPtr.Zero) { @@ -92,7 +106,6 @@ protected virtual void Dispose(bool disposing) handle = IntPtr.Zero; } - disposed = true; } /// @@ -124,6 +137,11 @@ internal IntPtr GetHandle() return handle; } + internal void RemoveOwnership() + { + this.hasOwnership = false; + } + /// /// Sets the neural network dataset property. /// diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Layer.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Layer.cs index 8f66059aa7d..db19839f520 100644 --- a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Layer.cs +++ b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Layer.cs @@ -49,6 +49,9 @@ public class Layer: IDisposable public Layer(NNTrainerLayerType type) { NNTrainerError ret = Interop.Layer.Create(out handle, type); + if (ret != NNTrainerError.None) { + handle = IntPtr.Zero; + } NNTrainer.CheckException(ret, "Failed to create model instance"); Log.Info(NNTrainer.Tag, $"Create layer with type:{type}"); } @@ -98,7 +101,7 @@ protected virtual void Dispose(bool disposing) disposed = true; if (!hasOwnership){ - Log.Error(NNTrainer.Tag, "Cannot destroy layer already added in a Model. Model will destroy this layer"); + Log.Info(NNTrainer.Tag, "Cannot destroy layer already added in a Model. Model will destroy this layer"); return; } diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Model.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Model.cs index d92636e0f57..fa6a86fa44e 100644 --- a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Model.cs +++ b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Model.cs @@ -65,6 +65,9 @@ public Model(string modelConf) NNTrainer.CheckException(NNTrainerError.InvalidParameter, "modelConf is null"); NNTrainerError ret = Interop.Model.ConstructWithConf(modelConf, out handle); + if (ret != NNTrainerError.None) { + handle = IntPtr.Zero; + } NNTrainer.CheckException(ret, "Failed to create model instance with modelConf"); Log.Info(NNTrainer.Tag, "Created Model with Conf path: "+ modelConf); } @@ -281,7 +284,7 @@ public void AddLayer(Layer layer) public Layer GetLayer(string layerName) { IntPtr layerHandle = IntPtr.Zero; - if (string.IsNullOrEmpty(layerName)) + if (string.IsNullOrEmpty(layerName)) NNTrainer.CheckException(NNTrainerError.InvalidParameter, "layerName is null"); NNTrainerError ret = Interop.Model.GetLayer(handle, layerName, out layerHandle); @@ -311,6 +314,7 @@ public void SetOptimizer(Optimizer optimizer) NNTrainerError ret = Interop.Model.SetOptimizer(handle, optimizer.GetHandle()); NNTrainer.CheckException(ret, "Failed to set optimizer"); + optimizer.RemoveOwnership(); } /// @@ -335,6 +339,7 @@ public void SetDataset(Dataset dataset) NNTrainerError ret = Interop.Model.SetDataset(handle, dataset.GetHandle()); NNTrainer.CheckException(ret, "Failed to set dataset"); + dataset.RemoveOwnership(); } internal static TensorsInfo CreateTensorsInfoFormHandle(IntPtr handle) diff --git a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Optimizer.cs b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Optimizer.cs index de7bdf389c3..e10174e3d18 100644 --- a/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Optimizer.cs +++ b/src/Tizen.MachineLearning.Train/Tizen.MachineLearning.Train/Optimizer.cs @@ -37,6 +37,9 @@ public class Optimizer: IDisposable private IntPtr handle = IntPtr.Zero; private bool disposed = false; + /// if false, model will be destroy optimizer handle + private bool hasOwnership = true; + /// /// Creates a neural network optimizer. /// @@ -45,6 +48,9 @@ public class Optimizer: IDisposable public Optimizer(NNTrainerOptimizerType type) { NNTrainerError ret = Interop.Optimizer.Create(out handle, type); + if (ret != NNTrainerError.None) { + handle = IntPtr.Zero; + } NNTrainer.CheckException(ret, "Failed to create optimizer instance"); Log.Info(NNTrainer.Tag, $"Create optimizer with type:{type}"); } @@ -83,6 +89,14 @@ protected virtual void Dispose(bool disposing) { // release managed object } + + disposed = true; + + if (!hasOwnership){ + Log.Info(NNTrainer.Tag, "Cannot destroy optimizer already added in a Model. Model will destroy this optimizer"); + return; + } + // release unmanaged object if (handle != IntPtr.Zero) { @@ -93,7 +107,6 @@ protected virtual void Dispose(bool disposing) handle = IntPtr.Zero; } - disposed = true; } /// @@ -122,5 +135,10 @@ internal IntPtr GetHandle() { return handle; } + + internal void RemoveOwnership() + { + this.hasOwnership = false; + } } }