Open music-dino opened 1 month ago
Wraps an mgx::parameter
IActivationLayer * addActivation (ITensor &input, ActivationType type)
Activations which require non-default alpha and/or beta params, must set them after the layer is created via the appropriate method Type can be changed after the layer is created. ActivationTypes that map to a single MGX operator:
ActivationTypes that map to a combination of MGX operators:
Unsure:
Not part of ONNX:
ILRNLayer * addLRN (ITensor &input, int64_t window, float alpha, float beta, float k)
Maps to a single MGX operator - "lrn" WindowSize, Alpha, Beta, and K can be changed after the layer is created
IScaleLayer * addScale (ITensor &input, ScaleMode mode, Weights shift, Weights scale, Weights power)
Doesn't seem to have ONNX equivalent
ISoftMaxLayer * addSoftMax (ITensor &input)
Axis must be set after the layer is created Maps to a single MGX operator - "softmax"
IConcatenationLayer * addConcatenation (ITensor *const *inputs, int32_t nbInputs)
Axis must be set after the layer is created Maps to a single MGX operator - "concat"
IElementWiseLayer * addElementWise (ITensor &input1, ITensor &input2, ElementWiseOperation op)
Operation can be changed after the layer is created ElementWiseOperations mapping to mgx operators:
ElementWiseOperations without mapping to mgx operators:
IUnaryLayer * addUnary (ITensor &input, UnaryOperation operation)
Operation can be changed after the layer is created UnaryOperations mapping to mgx operators:
IShuffleLayer * addShuffle (ITensor &input)
This layer shuffles data by applying in sequence: a transpose operation, a reshape operation and a second transpose operation. The dimension types of the output are those of the reshape dimension.
The layer has an optional second input. If present, it must be a 1D Int32 shape tensor, and the reshape dimensions are taken from it.
Maps to a combination of mgx operators (transpose and reshape).
IOneHotLayer * addOneHot (ITensor &indices, ITensor &values, ITensor &depth, int32_t axis)
Axis can be changed after the layer is created Maps to a combination of mgx operators
IReduceLayer * addReduce (ITensor &input, ReduceOperation operation, uint32_t reduceAxes, bool keepDimensions)
Operation, Axes, and KeepDimensions can be changed after the layer is created ReduceOperations mapping to mgx operators:
ITopKLayer * addTopK (ITensor &input, TopKOperation op, int32_t k, uint32_t reduceAxes)
Operation, K, and reduceAxes can be changed after the layer is created TopKOperations kMAX and kMIN correspond to the two possible values of the largest attribute? setInput(1) is used when K is to be provided as an input. MGX does not support this
IGatherLayer * addGather (ITensor &data, ITensor &indices, int32_t axis)
GatherMode must be set after the layer is created, unless kDEFAULT is desired GatherAxis and NbElementWiseDims can be changed after the layer is created Depending on GatherMode, the layer corresponds to different ONNX ops:
For mode kND the mgx "gathernd" operator can be used directly For mode kELEMENT the "gather" operator can be used in combination with other operators(NOTE: There seems to be a bug in parse_gather_elements.cpp:95, 0 is always passed for axis)
IGatherLayer * addGatherV2 (ITensor &data, ITensor &indices, GatherMode mode)
See above
IRaggedSoftMaxLayer * addRaggedSoftMax (ITensor &input, ITensor &bounds)
No equivalent ONNX or MGX operator
IMatrixMultiplyLayer * addMatrixMultiply (ITensor &input0, MatrixOperation op0, ITensor &input1, MatrixOperation op1)
Op0 and Op1 can be changed after the layer is created MatrixOperations kNONE and kTRANSPOSE can be supported, kVECTOR needs more investigation Can be implemented as GEMM with alpha and beta equal to 1
INonZeroLayer * addNonZero (ITensor &input)
Implemented via "nonzero" mgx op if input is not a constant, a literal is inserted otherwise
IConstantLayer * addConstant (Dims const &dimensions, Weights weights)
Wraps an mgx literal Dimensions and weights can be changed after layer is created
IIdentityLayer * addIdentity (ITensor &input)
Directly maps to mgx id operator
ICastLayer * addCast (ITensor &input, DataType toType)
toType can be changed after the layer is created Maps to mgx "convert" operator
IParametricReLULayer * addParametricReLU (ITensor &input, ITensor &slope)
Maps to mgx "prelu" operator
IConvolutionLayer * addConvolutionNd (ITensor &input, int64_t nbOutputMaps, Dims const &kernelSize, Weights kernelWeights, Weights biasWeights)
Input, nbOutputMaps, kernelSize, kernelWeights, and biasWeights can be changed after the layer is created Dilations, nbGroups, PaddingMode, PaddingNd, PostPadding, PrePadding, StrideNd all must be set after the layer is created if non-default values are to be used If setPaddingMode and setPrePadding/setPostPadding are both used, PaddingMode takes precedence The supported TRT padding modes are:
PrePadding and PostPadding amount to the same things as the pads attribute in ONNX. The difference being that ONNX prohibits pads being used when auto_pad is used, while TRT gives precedence to auto_pad/PaddingMode if both are used.
IPoolingLayer * addPoolingNd (ITensor &input, PoolingType type, Dims const &windowSize)
Corresponds to ONNX MaxPool, AveragePool, GlobalMaxPool and GlobalAveragePool For global pooling a window size equal to the input size should be used PoolingType kMAX_AVERAGE_BLEND does not correspond to an ONNX operator LpPooling is not supported PoolingType, and window size can be changed after the layer is created AverageCountExcludesPadding, PrePadding, PostPadding, PaddingMode, and Stride must be set after the layer is created if non-default values are to be used For padding information see addConvolutionNd
IDeconvolutionLayer * addDeconvolutionNd (ITensor &input, int64_t nbOutputMaps, Dims kernelSize, Weights kernelWeights, Weights biasWeights)
Corresponds to ONNX ConvTranspose operator Setting the output dimensions is not supported in TRT For other detail see addConvolutionNd
IScaleLayer * addScaleNd (ITensor &input, ScaleMode mode, Weights shift, Weights scale, Weights power, int32_t channelAxis)
TODO
IResizeLayer * addResize (ITensor &input)
Implemented as a combination of mgx operators OutputDimensions, Scales, ResizeMode, CoordinateTransformation, SelectorForSinglePixel, NearestRounding, CubicCoeff, ExcludeOutside must be set after the layer is created if non-default values are to be used InterpolationMode corresponds to ONNX mode, all are supported ResizeCoordinateTransformation corresponds to ONXX coordinate_transformation_mode, only half_pixel, asymmetric and align_corner are supported ResizeRoundMode corresponds to ONNX nearest_mode, all are supported Antialising is not supported
[ ] ILoop * addLoop ()
TODO
IIfConditional * addIfConditional ()
TODO
ISelectLayer * addSelect (ITensor &condition, ITensor &thenInput, ITensor &elseInput)
Equivalent to the "where" mgx/ONNX op
IAssertionLayer * addAssertion (ITensor &condition, char const *message)
No equivalent ONNX or MGX operator
IFillLayer * addFill (Dims const &dimensions, FillOperation op, DataType outputType)
Covers the RandomNormal and RandomNormalLike ONNX operators for KillOperation::kRANDOM_NORMAL, and RandomUniform and RandomUniformLike for FillOperation::kRANDOM_UNIFORM. Value of the seed attribute, if present, is ignored In MGX a literal with values from the appropriate distribution is created Operation and outputType can be changed after the layer is created Alpha and beta values must be changed after the layer is created if non-default values are to be used Alpha and beta have different meaning depending on FillOperation:
IPaddingLayer * addPaddingNd (ITensor &input, Dims const &prePadding, Dims const &postPadding)
Pre and post padding can be changed after the layer is created Only does zero padding
IDequantizeLayer * addDequantize (ITensor &input, ITensor &scale, DataType outputType)
TODO
IScatterLayer * addScatter (ITensor &data, ITensor &indices, ITensor &updates, ScatterMode mode)
Covers the ONNX ScatterElements and ScatterND operators for ScatterMode::kElement and ScatterMode::kND respectively Reduction is not supported Mode can be changed after the layer is created Axis must be set after the layer is created if non default value is to be used
...
The initial approach will be to wrap both layers and tensors around instructions.
Layers that are implemented by more than one mgx operator will hold references to the starting and ending instructions of the block of instructions that it's composed of, while the output tensor(s) will hold references to only the final instruction.
When a layer is modified after creation by use of a setter method, the existing mgx IR will be modified by use of replace_instruction
.
This approach allow us to use existing infrastructure for both instruction modification, which needs to be reflected in downstream instructions if the output type changes, and for creating composite operators, by reusing existing parsing code.
A difficulty that presents itself for any approach and this one especially are ONNX operators that hold subgraphs. An attempt will be made to validate if this approach can work for them by implementing the ILoopLayer. A good idea for how this layer works can be found in the trt onnx parser source code for parsing an ONNX Loop: https://github.com/onnx/onnx-tensorrt/blob/7ecb49a435bd881b9ac4011450315192885e5cc3/onnxOpImporters.cpp#L2810
The way INetworkDefinition tracks input and output tensors needs to be modified. MGX does not guarantee that output parameters will be inserted after input parameters, nor seemingly the relative ordering of the parameters.
Create an MGX API to mimic the TRT network definition APIs. Once the APIs are created, attempt to combine it with the TRT ONNX parser so that the network definition and layer manipulation APIs can be used on a network created by parsing an ONNX file.