Skip to content

Commit

Permalink
[MachineLearning.Train] Change internal logic to resolve a memory leak
Browse files Browse the repository at this point in the history
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 <[email protected]>
  • Loading branch information
SeoHyungjun authored and devusr-sw-kim committed Oct 15, 2024
1 parent da187ae commit 471c05d
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,19 @@ 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;

/// <summary>
/// Constructs the dataset.
/// </summary>
/// <since_tizen> 10 </since_tizen>
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");
}
Expand Down Expand Up @@ -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)
{
Expand All @@ -92,7 +106,6 @@ protected virtual void Dispose(bool disposing)

handle = IntPtr.Zero;
}
disposed = true;
}

/// <summary>
Expand Down Expand Up @@ -124,6 +137,11 @@ internal IntPtr GetHandle()
return handle;
}

internal void RemoveOwnership()
{
this.hasOwnership = false;
}

/// <summary>
/// Sets the neural network dataset property.
/// </summary>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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}");
}
Expand Down Expand Up @@ -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;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
}

/// <summary>
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/// <summary>
/// Creates a neural network optimizer.
/// </summary>
Expand All @@ -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}");
}
Expand Down Expand Up @@ -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)
{
Expand All @@ -93,7 +107,6 @@ protected virtual void Dispose(bool disposing)

handle = IntPtr.Zero;
}
disposed = true;
}

/// <summary>
Expand Down Expand Up @@ -122,5 +135,10 @@ internal IntPtr GetHandle()
{
return handle;
}

internal void RemoveOwnership()
{
this.hasOwnership = false;
}
}
}

0 comments on commit 471c05d

Please sign in to comment.