Add optimization after loading (#1127)

Author: NeoML-maintainer

NeoML/include/NeoML/Dnn/DnnOptimization.h

@@ -116,4 +116,6 @@ struct NEOML_API CDnnOptimizationSettings final {
 CDnnOptimizationReport NEOML_API OptimizeDnn( CDnn& dnn,
 	const CDnnOptimizationSettings& settings = CDnnOptimizationSettings() );
 
+CDnnOptimizationReport NEOML_API OptimizeDnnOnLoad(CDnn& dnn, size_t dnnWeightsTotalSize);
+
 } // namespace NeoML

NeoML/src/Dnn/Dnn.cpp

@@ -104,6 +104,8 @@ limitations under the License.
 #include <NeoML/Dnn/Layers/TransformerLayer.h>
 #include <NeoML/Dnn/Layers/TransformerSourceMaskLayer.h>
 #include <NeoML/Dnn/Layers/Upsampling2DLayer.h>
+#include <NeoML/Dnn/DnnOptimization.h>
+
 #endif //!NEOML_COMPACT
 
 namespace NeoML {
@@ -810,7 +812,7 @@ static constexpr int dnnVersion = 2000;
 void CDnn::Serialize( CArchive& archive )
 {
 	NeoAssertMsg( !IsReferenceDnn(), "For ReferenceDnn serializing is restricted" );
-
+	size_t before = mathEngine.GetCurrentMemoryUsage();
 	int version = dnnVersion;
 	archive.Serialize( version );
 
@@ -853,6 +855,9 @@ void CDnn::Serialize( CArchive& archive )
 		// In order to avoid the CDnnSolver::Reset for the next solver
 		rebuild();
 	}
+	size_t after = mathEngine.GetCurrentMemoryUsage();
+
+	OptimizeDnnOnLoad(*this, after - before);
 }
 
 void CDnn::SerializeCheckpoint( CArchive& archive )

NeoML/src/Dnn/DnnOptimization.cpp

@@ -43,10 +43,30 @@ CDnnOptimizationReport OptimizeDnn( CDnn& dnn, const CDnnOptimizationSettings& s
 		optimization::CMobileNetV3Optimizer( graph ).Apply( report );
 
 		CArray<int> chains;
-		OptimizeRowwiseChains( dnn, chains );
+		OptimizeRowwiseChains(dnn, chains);
 		report.RowwiseChainCount = chains.Size();
 	}
 	return report;
 }
 
+CDnnOptimizationReport OptimizeDnnOnLoad(CDnn& dnn, size_t size)
+{
+	CDnnOptimizationReport report;
+	optimization::CGraph graph(dnn);
+
+	report.UnpackedCompositeLayers = optimization::UnpackComposites(graph);
+	report.RemovedTrivialLayers = optimization::RemoveTrivialLayers(graph);
+	optimization::CBatchNormFusionOptimizer(graph).Apply(report);
+
+	if (size < 1024 * 1024) {
+		optimization::CMobileNetV2Optimizer(graph).Apply(report);
+		
+		CArray<int> chains;
+		OptimizeRowwiseChains(dnn, chains);
+		report.RowwiseChainCount = chains.Size();
+	}
+
+	return report;
+}
+
 } // namespace NeoML