[NeoML] BaseLayer more clean-up

Signed-off-by: Kirill Golikov <kirill.golikov@abbyy.com>

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

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() override { blob = nullptr; }
+
 protected:
 	CPtr<CDnnBlob> blob;
 

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

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 );
 	}
 }