[NeoML] total clean-up for BaseLayer

Signed-off-by: Kirill Golikov <kirill.golikov@abbyy.com>
neoml-lib · Apr 23, 2024 · 2e1fad5 · 2e1fad5
1 parent f1e59a2
commit 2e1fad5
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 8 deletions.
diff --git a/NeoML/include/NeoML/Dnn/Dnn.h b/NeoML/include/NeoML/Dnn/Dnn.h
@@ -204,7 +204,7 @@ class NEOML_API CBaseLayer : public virtual IObject {
 	virtual size_t GetOutputBlobsSize() const;
 
 	// Releases all temporary resources allocated for the layer
-	virtual void CleanUp();
+	virtual void CleanUp( bool totalCleanUp = false );
 
 	// Returns the total size of trainable parameters in this layer
 	// Returns the total size of trainable parameters of its internal layers, if layer is composite or recurrent
@@ -529,7 +529,7 @@ class NEOML_API CDnn : public CDnnLayerGraph {
 	void RunAndLearnOnce();
 
 	// Releases all temporary resources allocated for RunAndBackwardOnce()
-	void CleanUp();
+	void CleanUp( bool totalCleanUp = false );
 
 	// Gets the maximum sequence length
 	int GetMaxSequenceLength() const { return maxSequenceLength; }

diff --git a/NeoML/include/NeoML/Dnn/Layers/CompositeLayer.h b/NeoML/include/NeoML/Dnn/Layers/CompositeLayer.h
@@ -67,7 +67,7 @@ class NEOML_API CCompositeLayer : public CBaseLayer, public CDnnLayerGraph {
 	size_t GetOutputBlobsSize() const override;
 
 	// Releases all temporary resources allocated for the layer
-	void CleanUp() override;
+	void CleanUp( bool totalCleanUp = false ) override;
 
 	// Returns the total size of trainable parameters
 	size_t GetTrainableParametersSize() const override;

diff --git a/NeoML/include/NeoML/Dnn/Layers/SinkLayer.h b/NeoML/include/NeoML/Dnn/Layers/SinkLayer.h
@@ -33,6 +33,8 @@ class NEOML_API CSinkLayer : public CBaseLayer {
 	// After each call to RunOnce this blob contains the results
 	const CPtr<CDnnBlob>& GetBlob() const;
 
+	void CleanUp( bool totalCleanUp = false ) override { CBaseLayer::CleanUp( totalCleanUp ); blob = nullptr; }
+
 protected:
 	CPtr<CDnnBlob> blob;
 

diff --git a/NeoML/src/Dnn/BaseLayer.cpp b/NeoML/src/Dnn/BaseLayer.cpp
@@ -272,12 +272,22 @@ size_t CBaseLayer::GetOutputBlobsSize() const
 	return result;
 }
 
-void CBaseLayer::CleanUp()
+void CBaseLayer::CleanUp( bool totalCleanUp )
 {
 	inputBlobs.DeleteAll();
 	inputBlobs.SetSize(inputDescs.Size());
 	outputBlobs.DeleteAll();
 	outputBlobs.SetSize(outputDescs.Size());
+
+	if ( totalCleanUp ) {
+		inputDiffBlobs.DeleteAll();
+		outputDiffBlobs.DeleteAll();
+		paramDiffBlobs.DeleteAll();
+		readyOutputDiffs.DeleteAll();
+		clearAllRuntimeBlobs();
+
+		ForceReshape();
+	}
 }
 
 size_t CBaseLayer::GetTrainableParametersSize() const

diff --git a/NeoML/src/Dnn/Dnn.cpp b/NeoML/src/Dnn/Dnn.cpp
@@ -678,10 +678,10 @@ void CDnn::RunAndLearnOnce()
 	solver->Train();
 }
 
-void CDnn::CleanUp()
+void CDnn::CleanUp( bool totalCleanUp )
 {
 	for( int i = 0; i < layers.Size(); i++ ) {
-		layers[i]->CleanUp();
+		layers[i]->CleanUp( totalCleanUp );
 	}
 }
 

diff --git a/NeoML/src/Dnn/Layers/CompositeLayer.cpp b/NeoML/src/Dnn/Layers/CompositeLayer.cpp
@@ -499,10 +499,11 @@ size_t CCompositeLayer::GetOutputBlobsSize() const
 	return result;
 }
 
-void CCompositeLayer::CleanUp()
+void CCompositeLayer::CleanUp( bool totalCleanUp )
 {
+	CBaseLayer::CleanUp( totalCleanUp );
 	for( int i = 0; i < internalDnn->layers.Size(); i++ ) {
-		internalDnn->layers[i]->CleanUp();
+		internalDnn->layers[i]->CleanUp( totalCleanUp );
 	}
 }