From ca4a5da013524d780cb362ec438cab13db06a727 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Wed, 11 May 2022 17:54:48 +0800 Subject: [PATCH] feat(serialization): add new serialization format define GitOrigin-RevId: 177bfdd56bad3feeb2c727ded05e627c271fe2ab --- src/CMakeLists.txt | 3 +- src/serialization/impl/schema_v2.fbs | 222 +++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 src/serialization/impl/schema_v2.fbs diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 41e97c70d..01efb2ef1 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -185,7 +185,8 @@ if(MGB_WITH_FLATBUFFERS) ${CMAKE_CURRENT_BINARY_DIR}/serialization/impl/opr_param_defs.fbs ${CMAKE_CURRENT_BINARY_DIR}/serialization/impl/mgb_opr_param_defs.fbs ${CMAKE_CURRENT_SOURCE_DIR}/opr/impl/mgb_cpp_opr.fbs - ${CMAKE_CURRENT_SOURCE_DIR}/serialization/impl/schema.fbs) + ${CMAKE_CURRENT_SOURCE_DIR}/serialization/impl/schema.fbs + ${CMAKE_CURRENT_SOURCE_DIR}/serialization/impl/schema_v2.fbs) list( APPEND FLATBUFFERS_SCHEMA_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/serialization/impl diff --git a/src/serialization/impl/schema_v2.fbs b/src/serialization/impl/schema_v2.fbs new file mode 100644 index 000000000..2e1f58b90 --- /dev/null +++ b/src/serialization/impl/schema_v2.fbs @@ -0,0 +1,222 @@ +include "dtype.fbs"; +include "opr_param_defs.fbs"; +include "mgb_opr_param_defs.fbs"; +include "mgb_cpp_opr.fbs"; + +namespace mgb.serialization.fbs.v2; + +file_identifier "mge2"; + +table CompNode { + logical_locator:string; +} + +table DefaultTensorFormat{} + +table Image2DPackedTensorFormat{ + align_axis: ubyte; +} + +table LowbitsAlignedTensorFormat{ + size_nbits: ubyte; + align_size_in_bits: ubyte; +} + +/// The Tensor Format +union TensorFormat { + DefaultTensorFormat = 1, + Image2DPackedTensorFormat = 2, + LowbitsAlignedTensorFormat = 3, +} + +/// Opaque byte buffer defined by operator implementation +table Blob { + data:[ubyte]; +} + +table Tensor { + name:string; + shape:[uint]; + comp_node:CompNode; + dtype:DType; + format:TensorFormat; + /// The tensor raw data + data:[ubyte]; +} + +table Reserved0 {} +table DeprecatedParam {} + +union OperatorParam { + param.Empty = 1, + param.Axis = 2, + param.Convolution = 3, + param.MaskPropagate = 4, + param.ConvPooling = 5, + param.ConvBias = 6, + param.SeparableConv = 7, + param.Images2Neibs = 8, + param.Pooling = 9, + param.LRN = 10, + param.BN = 11, + param.ROIPooling = 12, + param.WarpPerspective = 13, + param.SpatialTfGridGenerator = 14, + param.SpatialTfSampler = 15, + param.MGBAddUpdate = 16, + param.Elemwise = 17, + param.ElemwiseMultiType = 18, + param.PowC = 19, + param.MatrixMul = 20, + //Reserved for param.Winograd = 21, + DeprecatedParam = 21, + param.SVD = 22, + param.Reduce = 23, + param.Cumsum = 24, + param.CondTake = 25, + param.Argsort = 26, + param.IndexingRemap = 27, + param.MGBSleep = 28, + param.Linspace = 29, + param.LinspaceFull = 30, + param.Eye = 31, + param.UniformRNG = 32, + param.GaussianRNG = 33, + param.Flip = 34, + param.Rotate = 35, + param.ROICopy = 36, + param.CvtColor = 37, + param.WarpAffine = 38, + param.GaussianBlur = 39, + param.Resize = 40, + param.Convolution3D = 41, + param.Conv3DBias = 42, + param.SeparableConv3D = 43, + param.TopK = 44, + param.RelayoutFormat = 45, + param.SeparableFilter = 46, + param.LocalShare = 47, + param.ROIAlign = 48, + param.DeformablePSROIPooling = 49, + param.BatchConvBias = 50, + param.DType = 51, + param.PersistentOutputStorage = 52, + param.OptionalAxis = 53, + param.OptionalAxisV1 = 54, + param.ExecutionPolicy = 55, + param.AssertEqual = 56, + param.FpgaConv = 57, + param.CollectiveComm = 58, + param.CondExecPred = 59, + param.CondExecPredLogical = 60, + param.CondExecMark = 61, + param.CondExecMerge = 62, + param.Host2DeviceCopy = 63, + param.Dimshuffle = 64, + param.AxisAddRemove = 65, + param.IndexDescMaskDump = 66, + DType = 67, + param.Remap = 68, + param.NMSKeep = 69, + param.AdaptivePooling = 70, + param.NvOf = 71, + param.DctChannelSelect = 72, + param.FakeQuant = 73, + param.TQT = 74, + param.Correlation = 75, + param.LSQ = 76, + param.GammaRNG = 77, + param.PoissonRNG = 78, + param.PermutationRNG = 79, + param.BetaRNG = 80, + param.SlidingWindowTranspose = 81, + param.Padding = 82, + param.ShuffleRNG = 83, + param.CheckNonFinite = 84, + param.LayerNorm = 85, + param.Dropout = 86, + param.RNNCell = 87, + param.RNN = 88, + param.LSTM = 89, + param.Softmax = 90, + param.Diag = 91, +} + +table Operator { + /// the Operator type id + type:string; + /// sometime type maybe not exist, so add type_id + type_id:ulong; + name:string; + + /// Operator parameter + param:OperatorParam; + /// Operator may want to save more than one OperatorParam + additional_params:[OperatorParam]; + + /// ID of the input tensor in the middle_tensors of a model + inputs:[uint]; + + /// ID of the output tensor in the middle_tensors of a model + outputs:[uint]; + + comp_node:[CompNode]; + output_dtype:DType; + + /// the const value in tensor format of the Operator + tensors:[Tensor]; + + /// opr version, with develop of MegEngine, some opr may have multi version + opr_version:uint; + + /// the order of the Operator in the graph + priority:int = 0; + + /// custom may want to save big, opaque byte buffers. + custom_data:[Blob]; +} + +table Metadata { + is_valid:bool; + graph_modified:bool; + optimize_options:ulong; + user_info:string; +} + +table MiddleTensor { + name:string; + shape:[uint]; + comp_node:CompNode; + dtype:DType; + format:TensorFormat; +} + +table OutputVar { + /// the id of the middle tensor in graph, the same as the inputs in Operator + compact_id:uint; + original_id:uint; +} + +table Model { + /// the megengine version when serialize the model + mge_version:uint; + + /// model version, now model support: + /// version v1: the original fbs serialization version + /// version v2: support backward and poor forward compatibility + model_version:uint; + + oprs:[Operator]; + + /// the tensors produce and consume by the Operators, not the input or + /// output tensor + middle_tensors:[MiddleTensor]; + + output_vars_idx:[OutputVar]; + + nr_shared_tensor:uint; + /// the Metadata to storage the custom data or some flags + metadata:Metadata; +} + +root_type Model; -- GitLab