Support meta data in NeoOnnx import (neoml-lib#659)

* Support meta data in NeoOnnx import Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Fix naming and remove default value Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Mention metadata in docs Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Fix Onnx2Neoml Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Fix NeoProxy Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Typo fix Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com> * Fix exported symbols Signed-off-by: Valeriy Fedyunin <valery.fedyunin@abbyy.com>
zimka · Jun 17, 2022 · 2759ccc · 2759ccc
1 parent 662b1a9
commit 2759ccc
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 13 deletions.
diff --git a/NeoML/docs/en/Onnx.md b/NeoML/docs/en/Onnx.md
@@ -18,8 +18,10 @@ The **NeoOnnx** library lets you load third-party neural networks serialized in
 ```c++
 #include <NeoOnnx/NeoOnnx.h>
 
-NEOONNX_API void LoadFromOnnx( const char* fileName, NeoML::CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs );
-NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs );
+NEOONNX_API void LoadFromOnnx( const char* fileName, NeoML::CDnn& dnn, CArray<const char*>& inputs,
+    CArray<const char*>& outputs, CMap<CString, CString>& metadata );
+NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn& dnn, CArray<const char*>& inputs,
+    CArray<const char*>& outputs, CMap<CString, CString>& metadata );
 ```
 
 Loads a network from a file or a buffer.
@@ -28,6 +30,8 @@ For each network input the `dnn` network will have a `CSourceLayer` with the sam
 
 For each network output the `dnn` network will have a `CSinkLayer` with the same name. Also output names will be added to the `outputs` array.
 
+Model's `metadata_props` will be written to the `metadata`.
+
 Both `inputs` and `outputs` name pointers are attached to the names of the corresponding layers in `dnn` and will become invalid after changing name or deleting corresponding layer.
 
 ## Build

diff --git a/NeoML/docs/ru/Onnx.md b/NeoML/docs/ru/Onnx.md
@@ -18,8 +18,10 @@
 ```c++
 #include <NeoOnnx/NeoOnnx.h>
 
-NEOONNX_API void LoadFromOnnx( const char* fileName, NeoML::CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs );
-NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs );
+NEOONNX_API void LoadFromOnnx( const char* fileName, NeoML::CDnn& dnn, CArray<const char*>& inputs,
+    CArray<const char*>& outputs, CMap<CString, CString>& metadata );
+NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn& dnn, CArray<const char*>& inputs,
+    CArray<const char*>& outputs, CMap<CString, CString>& metadata );
 ```
 
 Загружает сеть из файла или из буфера.
@@ -28,6 +30,8 @@ NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn&
 
 Для каждого выхода сети в `dnn` будет создан `CSinkLayer` с таким же именем. Также имена выходов будут добавлены в массив `outputs`.
 
+`metadata_props` модели будут записаны в `metadata`.
+
 Указатели в массивах `inputs` и `outputs` привязаны к именам соответствующих слоёв и станут невалидными после переименования или разрушения соответствующих слоёв.
 
 ## Сборка

diff --git a/NeoOnnx/Onnx2NeoML/Onnx2NeoML.cpp b/NeoOnnx/Onnx2NeoML/Onnx2NeoML.cpp
@@ -29,8 +29,9 @@ int ConvertOnnx2NeoML( const char* inputOnnxFilename, const char* outputDnnArchi
     IMathEngine& mathEng = GetDefaultCpuMathEngine();
     CRandom random( 0x123 );
     CDnn net( random, mathEng );
+    CMap<CString, CString> metaData;
     try {
-        NeoOnnx::LoadFromOnnx( inputOnnxFilename, net, inputs, outputs );
+        NeoOnnx::LoadFromOnnx( inputOnnxFilename, net, inputs, outputs, metaData );
         {
             CArchiveFile file( outputDnnArchiveFileName, CArchive::store );
             CArchive archive( &file, CArchive::SD_Storing );

diff --git a/NeoOnnx/include/NeoOnnx/NeoOnnxImport.h b/NeoOnnx/include/NeoOnnx/NeoOnnxImport.h
@@ -28,6 +28,7 @@ namespace NeoOnnx {
 //
 // For every onnx graph output there will be CSinkLayer with the same name
 // Graph inputs' and outputs' names will be added to the corresponding CArray's
+// ONNX model's metadata_props will be written to the metadata CMap
 // Names' pointers are attached to the corresponding layers' names
 //
 // Input and output blobs have the following relations with the ONNX N-dimensional tensors:
@@ -41,10 +42,12 @@ namespace NeoOnnx {
 
 
 // Loads network "dnn" from onnx file "fileName"
-NEOONNX_API void LoadFromOnnx( const char* fileName, NeoML::CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs );
+NEOONNX_API void LoadFromOnnx( const char* fileName, NeoML::CDnn& dnn, CArray<const char*>& inputs,
+	CArray<const char*>& outputs, CMap<CString, CString>& metadata );
 
 // Loads network "dnn" from buffer with onnx data
-NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs );
+NEOONNX_API void LoadFromOnnx( const void* buffer, int bufferSize, NeoML::CDnn& dnn, CArray<const char*>& inputs,
+	CArray<const char*>& outputs, CMap<CString, CString>& metadata );
 
 } // namespace NeoOnnx
 
diff --git a/NeoOnnx/src/CMakeLists.txt b/NeoOnnx/src/CMakeLists.txt
@@ -212,8 +212,8 @@ endif()
 
 if(NeoOnnx_BUILD_SHARED AND NOT WIN32)
     set(EXPORT_SYMBOLS
-        _ZN7NeoOnnx12LoadFromOnnxEPKviRN5NeoML4CDnnERN4FObj6CArrayIPKcNS5_20CurrentMemoryManagerEEESB_
-        _ZN7NeoOnnx12LoadFromOnnxEPKcRN5NeoML4CDnnERN4FObj6CArrayIS1_NS5_20CurrentMemoryManagerEEES9_)
+        _ZN7NeoOnnx12LoadFromOnnxEPKviRN5NeoML4CDnnERN4FObj6CArrayIPKcNS5_20CurrentMemoryManagerEEESB_RNS5_4CMapINS5_7CStringESD_NS5_12CDefaultHashISD_EES9_EE
+        _ZN7NeoOnnx12LoadFromOnnxEPKcRN5NeoML4CDnnERN4FObj6CArrayIS1_NS5_20CurrentMemoryManagerEEES9_RNS5_4CMapINS5_7CStringESB_NS5_12CDefaultHashISB_EES7_EE)
     if(LINUX OR ANDROID)
         set(EXPORT_SYMBOLS_STR "${EXPORT_SYMBOLS};")
         configure_file(../NeoOnnx.version.in NeoOnnx.version)

diff --git a/NeoOnnx/src/NeoOnnxImport.cpp b/NeoOnnx/src/NeoOnnxImport.cpp
@@ -152,7 +152,17 @@ static void buildDnnFromGraphProto( const onnx::GraphProto& onnxGraph, int opset
 	}
 }
 
-void LoadFromOnnx( const char* fileName, CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs )
+// Extracts metadata from the model
+static void extractMetadata( const onnx::ModelProto& model, CMap<CString, CString>& metadata )
+{
+	metadata.Empty();
+	for( const onnx::StringStringEntryProto& entry : model.metadata_props() ) {
+		metadata.Add( CString( entry.key().data() ), CString( entry.value().data() ) );
+	}
+}
+
+void LoadFromOnnx( const char* fileName, CDnn& dnn, CArray<const char*>& inputs,
+	CArray<const char*>& outputs, CMap<CString, CString>& metadata )
 {
 	GOOGLE_PROTOBUF_VERIFY_VERSION;
 
@@ -169,6 +179,7 @@ void LoadFromOnnx( const char* fileName, CDnn& dnn, CArray<const char*>& inputs,
 		}
 
 		buildDnnFromGraphProto( model.graph(), getOpsetVersion( model ), dnn, inputs, outputs );
+		extractMetadata( model, metadata );
 	} catch( ... ) {
 		input.close();
 		google::protobuf::ShutdownProtobufLibrary();
@@ -179,7 +190,8 @@ void LoadFromOnnx( const char* fileName, CDnn& dnn, CArray<const char*>& inputs,
 	google::protobuf::ShutdownProtobufLibrary();
 }
 
-void LoadFromOnnx( const void* buffer, int bufferSize, CDnn& dnn, CArray<const char*>& inputs, CArray<const char*>& outputs )
+void LoadFromOnnx( const void* buffer, int bufferSize, CDnn& dnn, CArray<const char*>& inputs,
+	CArray<const char*>& outputs, CMap<CString, CString>& metadata )
 {
 	GOOGLE_PROTOBUF_VERIFY_VERSION;
 
@@ -193,6 +205,7 @@ void LoadFromOnnx( const void* buffer, int bufferSize, CDnn& dnn, CArray<const c
 		}
 
 		buildDnnFromGraphProto( model.graph(), getOpsetVersion( model ), dnn, inputs, outputs );
+		extractMetadata( model, metadata );
 	} catch( ... ) {
 		google::protobuf::ShutdownProtobufLibrary();
 		throw;

diff --git a/NeoProxy/src/NeoProxy.cpp b/NeoProxy/src/NeoProxy.cpp
@@ -584,8 +584,9 @@ const struct CDnnDesc* CreateDnnFromOnnxFile( const struct CDnnMathEngineDesc* m
 	try {
 		CArray<const char*> inputs;
 		CArray<const char*> outputs;
+		CMap<CString, CString> metadata;
 		dnnDesc = FINE_DEBUG_NEW CDnnDescImpl( mathEngine );
-		NeoOnnx::LoadFromOnnx( fileName, dnnDesc->Dnn(), inputs, outputs );
+		NeoOnnx::LoadFromOnnx( fileName, dnnDesc->Dnn(), inputs, outputs, metadata );
 		dnnDesc->SetInputNames( inputs );
 		dnnDesc->SetOutputNames( outputs );
 		return dnnDesc;
@@ -626,8 +627,9 @@ const struct CDnnDesc* CreateDnnFromOnnxBuffer( const struct CDnnMathEngineDesc*
 	try {
 		CArray<const char*> inputs;
 		CArray<const char*> outputs;
+		CMap<CString, CString> metadata;
 		dnnDesc = FINE_DEBUG_NEW CDnnDescImpl( mathEngine );
-		NeoOnnx::LoadFromOnnx( buffer, bufferSize, dnnDesc->Dnn(), inputs, outputs );
+		NeoOnnx::LoadFromOnnx( buffer, bufferSize, dnnDesc->Dnn(), inputs, outputs, metadata );
 		dnnDesc->SetInputNames( inputs );
 		dnnDesc->SetOutputNames( outputs );
 		return dnnDesc;