[NeoML] CleanUp CSourceLayer

NeoML/include/NeoML/Dnn/Layers/ParameterLayer.h

@@ -1,4 +1,4 @@
-/* Copyright © 2024 ABBYY Production LLC
+/* Copyright © 2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -31,12 +31,14 @@ class NEOML_API CParameterLayer : public CBaseLayer {
     void SetBlob(CDnnBlob* _blob);
     void SetBlobDesc(const CBlobDesc& _desc);
     const CPtr<CDnnBlob>& GetBlob() const { return paramBlobs[0]; }
+
 protected:
     void AllocateOutputBlobs() override;
     void Reshape() override;
     void RunOnce() override;
     void BackwardOnce() override;
     void LearnOnce() override;
+
 private:
     CBlobDesc desc;
 };

NeoML/include/NeoML/Dnn/Layers/SourceLayer.h

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -24,7 +24,8 @@ namespace NeoML {
 class NEOML_API CSourceLayer : public CBaseLayer {
 	NEOML_DNN_LAYER( CSourceLayer )
 public:
-	explicit CSourceLayer( IMathEngine& mathEngine ) : CBaseLayer( mathEngine, "CCnnSourceLayer", false ), storeBlob( false ) {}
+	explicit CSourceLayer( IMathEngine& mathEngine ) :
+		CBaseLayer( mathEngine, "CCnnSourceLayer", false ), storeBlob( false ) {}
 
 	// Sets the input data blob
 	void SetBlob( CDnnBlob* blob );
@@ -38,9 +39,12 @@ class NEOML_API CSourceLayer : public CBaseLayer {
 
 	void Serialize( CArchive& archive ) override;
 
+	void CleanUp( bool totalCleanUp = false ) override;
+
 protected:
 	CPtr<CDnnBlob> blob;
 	bool storeBlob;
+
 	// CBaseLayer class methods
 	void Reshape() override;
 	void RunOnce() override;

NeoML/src/Dnn/Layers/ParameterLayer.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2024 ABBYY Production LLC
+/* Copyright © 2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -31,9 +31,10 @@ void CParameterLayer::SetBlob(CDnnBlob* _blob)
 	bool sameBlob = _blob == paramBlobs[0].Ptr();
 	paramBlobs[0] = _blob;
 
-	if (!outputDescs.IsEmpty()) {
-		if (paramBlobs[0]->GetDataType() != outputDescs[0].GetDataType()
-			|| !paramBlobs[0]->GetDesc().HasEqualDimensions(outputDescs[0]))
+	if( !outputDescs.IsEmpty() ) {
+		if( _blob != nullptr
+			&& ( _blob->GetDataType() != outputDescs[0].GetDataType()
+			|| !_blob->GetDesc().HasEqualDimensions( outputDescs[0] ) ) )
 		{
 			outputDescs[0] = paramBlobs[0]->GetDesc();
 			ForceReshape();

NeoML/src/Dnn/Layers/SourceLayer.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -26,8 +26,9 @@ void CSourceLayer::SetBlob( CDnnBlob* _blob )
 	blob = _blob;
 
 	if( !outputDescs.IsEmpty() ) {
-		if( blob->GetDataType() != outputDescs[0].GetDataType()
-			|| !blob->GetDesc().HasEqualDimensions( outputDescs[0] ) )
+		if( blob != nullptr
+			&& ( blob->GetDataType() != outputDescs[0].GetDataType()
+			|| !blob->GetDesc().HasEqualDimensions( outputDescs[0] ) ) )
 		{
 			outputDescs[0] = blob->GetDesc();
 			ForceReshape();
@@ -93,6 +94,12 @@ void CSourceLayer::Serialize( CArchive& archive )
 	}
 }
 
+void CSourceLayer::CleanUp( bool totalCleanUp )
+{
+	CBaseLayer::CleanUp( totalCleanUp );
+	SetBlob( nullptr );
+}
+
 CSourceLayer* Source( CDnn& network, const char* name )
 {
 	CPtr<CSourceLayer> source = new CSourceLayer( network.GetMathEngine() );

NeoML/test/src/DnnSolverTest.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2021-2023 ABBYY
+/* Copyright © 2021-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -21,13 +21,19 @@ limitations under the License.
 using namespace NeoML;
 using namespace NeoMLTest;
 
+namespace NeoMLTest {
+
 // Returns coefficient of neuron with one in and one out.
 static float getFcCoeff( const CFullyConnectedLayer* fcLayer )
 {
 	NeoAssert( fcLayer != 0 );
 	return fcLayer->GetWeightsData()->GetData().GetValue();
 }
 
+} // namespace NeoMLTest
+
+//---------------------------------------------------------------------------------------------------------------------
+
 // Check build/change correctness with gradient accumulation enabled
 TEST( CDnnSolverTest, NetworkModificationOnGradientAccumulation )
 {
@@ -150,51 +156,68 @@ TEST( CDnnSolverTest, NetworkModificationOnGradientAccumulation )
 	EXPECT_EQ( e1, e2 );
 }
 
-// Net for weight check.
-class CWeightCheckNet {
+//---------------------------------------------------------------------------------------------------------------------
+
+namespace NeoMLTest {
+
+// Network for weights check.
+class CWeightCheckNet final {
 public:
-	CWeightCheckNet();
-	void SetSolver( CDnnSolver* solver ) { dnn.SetSolver( solver ); }
+	CWeightCheckNet( CDnnSolver* solver );
+
 	float RunAndLearnOnce();
 	void GetWeights( CArray<float>& weights ) const;
+	void CleanUp();
+
 private:
 	CRandom random;
 	CDnn dnn;
+
+	CPtr<CSourceLayer> data;
+	CPtr<CSourceLayer> label;
 	CPtr<CFullyConnectedLayer> fc;
 	CPtr<CLossLayer> loss;
+
+	void setInputs();
 };
 
-CWeightCheckNet::CWeightCheckNet() :
+CWeightCheckNet::CWeightCheckNet( CDnnSolver* solver ) :
 	random( 0xAAAAAAAA ),
 	dnn( random, MathEngine() )
 {
-	CPtr<CSourceLayer> data = AddLayer<CSourceLayer>( "data", dnn );
+	dnn.SetSolver( solver );
+
+	data = AddLayer<CSourceLayer>( "data", dnn );
+	label = AddLayer<CSourceLayer>( "label", dnn );
+
+	fc = AddLayer<CFullyConnectedLayer>( "fc", { data } );
+	fc->SetNumberOfElements( 2 );
+	fc->SetZeroFreeTerm( true );
+
+	loss = AddLayer<CCrossEntropyLossLayer>( "loss", { fc, label } );
+	setInputs();
+}
+
+void CWeightCheckNet::setInputs()
+{
 	{
 		CPtr<CDnnBlob> dataBlob = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, 1, 1, 2 );
 		CArray<float> buff( { 0.25f, -0.345f } );
 		dataBlob->CopyFrom( buff.GetPtr() );
 		data->SetBlob( dataBlob );
 	}
-
-	CPtr<CSourceLayer> label = AddLayer<CSourceLayer>( "label", dnn );
 	{
 		CPtr<CDnnBlob> labelBlob = CDnnBlob::CreateDataBlob( MathEngine(), CT_Int, 1, 1, 1 );
 		CArray<int> buff( { 0 } );
 		labelBlob->CopyFrom( buff.GetPtr() );
 		label->SetBlob( labelBlob );
 	}
-
-	fc = AddLayer<CFullyConnectedLayer>( "fc", { data } );
-	fc->SetNumberOfElements( 2 );
-	fc->SetZeroFreeTerm( true );
 	{
 		CPtr<CDnnBlob> weightBlob = CDnnBlob::CreateDataBlob( MathEngine(), CT_Float, 1, 2, 2 );
 		CArray<float> buff = { -0.5f, 0.9f, 0.3f, -0.7f };
 		weightBlob->CopyFrom( buff.GetPtr() );
 		fc->SetWeightsData( weightBlob );
 	}
-
-	loss = AddLayer<CCrossEntropyLossLayer>( "loss", { fc, label } );
 }
 
 float CWeightCheckNet::RunAndLearnOnce()
@@ -210,23 +233,44 @@ void CWeightCheckNet::GetWeights( CArray<float>& weights ) const
 	weightBlob->CopyTo( weights.GetPtr() );
 }
 
+void CWeightCheckNet::CleanUp()
+{
+	dnn.CleanUp( /*totalCleanUp*/true );
+	MathEngine().CleanUp();
+	dnn.Random().Reset( 0xAAAAAAAA );
+	dnn.GetSolver()->Reset();
+	setInputs();
+}
+
+//---------------------------------------------------------------------------------------------------------------------
+
 void testSolver( CDnnSolver* solver, const CArray<CArray<float>>& expected )
 {
-	CWeightCheckNet net;
 	CArray<float> weights;
-	net.SetSolver( solver );
+	auto check = [&]( CWeightCheckNet& net )
+	{
 	for( int i = 0; i < expected.Size(); ++i ) {
 		float loss = net.RunAndLearnOnce();
 		loss;
 		net.GetWeights( weights );
-		ASSERT_EQ( expected[i].Size(), weights.Size() );
+		EXPECT_EQ( expected[i].Size(), weights.Size() );
 		for( int j = 0; j < weights.Size(); ++j ) {
-			ASSERT_TRUE( FloatEq( expected[i][j], weights[j] ) );
+			EXPECT_NEAR( expected[i][j], weights[j], 1e-5 );
 		}
 	}
+	};
+
+	CWeightCheckNet net( solver );
+	check( net );
+
+	net.CleanUp();
+	check( net );
 }
 
-// ====================================================================================================================
+} // namespace NeoMLTest
+
+//---------------------------------------------------------------------------------------------------------------------
+
 // Sgd.
 
 TEST( CDnnSolverTest, SgdNoReg )
@@ -339,7 +383,8 @@ TEST( CDnnSolverTest, SgdCompatL2 )
 	testSolver( sgd, expected );
 }
 
-// ====================================================================================================================
+//---------------------------------------------------------------------------------------------------------------------
+
 // Adam.
 
 TEST( CDnnSolverTest, AdamNoReg )
@@ -469,7 +514,8 @@ TEST( CDnnSolverTest, AdamCompatL2 )
 	testSolver( adam, expected );
 }
 
-// ====================================================================================================================
+//---------------------------------------------------------------------------------------------------------------------
+
 // Nadam.
 
 TEST( CDnnSolverTest, NadamNoReg )
@@ -532,6 +578,12 @@ TEST( CDnnSolverTest, NadamL2 )
 	testSolver( adam, expected );
 }
 
+//---------------------------------------------------------------------------------------------------------------------
+
+// Serialization.
+
+namespace NeoMLTest {
+
 static bool checkBlobEquality( CDnnBlob& firstBlob, CDnnBlob& secondBlob )
 {
 	if( !firstBlob.HasEqualDimensions( &secondBlob ) ) {
@@ -683,6 +735,10 @@ static void solverSerializationTestImpl( CPtr<CDnnSolver> firstSolver, bool trai
 	EXPECT_TRUE( checkLstmEquality( reverse, secondReverse ) );
 }
 
+} // namespace NeoMLTest
+
+// sgd.
+
 TEST( CDnnSimpleGradientSolverTest, Serialization1 )
 {
 	CPtr<CDnnSimpleGradientSolver> sgd = new CDnnSimpleGradientSolver( MathEngine() );
@@ -707,6 +763,10 @@ TEST( CDnnSimpleGradientSolverTest, Serialization2 )
 	solverSerializationTestImpl( sgd.Ptr(), false );
 }
 
+//---------------------------------------------------------------------------------------------------------------------
+
+// Adam.
+
 TEST( CDnnAdaptiveGradientSolverTest, Serialization1 )
 {
 	CPtr<CDnnAdaptiveGradientSolver> adam = new CDnnAdaptiveGradientSolver( MathEngine() );
@@ -769,6 +829,10 @@ TEST( CDnnNesterovGradientSolverTest, Serialization2 )
 	solverSerializationTestImpl( nadam.Ptr(), false );
 }
 
+//---------------------------------------------------------------------------------------------------------------------
+
+// Lamb.
+
 TEST( CDnnLambGradientSolverTest, Serialization1 )
 {
 	CPtr<CDnnLambGradientSolver> lamb = new CDnnLambGradientSolver( MathEngine() );
@@ -834,7 +898,7 @@ TEST( CDnnLambGradientSolverTest, Serialization4 )
 	solverSerializationTestImpl( lamb.Ptr(), true );
 }
 
-// ====================================================================================================================
+//---------------------------------------------------------------------------------------------------------------------
 
 TEST( CDnnSolverTest, CompositeLearningRate )
 {