bool ConvertShape(Layer layer, ModelBuilder net)
{
var shape = IRShapeInferenceHelper.ShapeInference.OnnxLayoutToTensorShape(layer.pool);
var permutations = shape.Get8DPermutationsForNCHWPermutationsAndShape(k_FromNCHWtoNHWC);
// Preserve symbolic shape by operating on int array instead of TensorShape, which would resolve unknown dimensions
layer.pool = TensorExtensions.Permute(IRShapeInferenceHelper.ShapeInference.OnnxLayoutToTensorShapeLayout(layer.pool), permutations);
return(true);
}
static public int[] Permutation4DTo8D(int[] permutations)
{
if (permutations.Length == TensorShape.MaxRank)
{
return(permutations);
}
int batchOldAxis = TensorExtensions.Convert4DTo8DAxis(permutations[0]);
int heighOldAxis = TensorExtensions.Convert4DTo8DAxis(permutations[1]);
int widthOldIndex = TensorExtensions.Convert4DTo8DAxis(permutations[2]);
int channeOldIndex = TensorExtensions.Convert4DTo8DAxis(permutations[3]);
return(new int[] { 0, 1, batchOldAxis, 3, 4, heighOldAxis, widthOldIndex, channeOldIndex });
}
bool ConvertNormal(Layer layer, ModelBuilder net)
{
if (layer.inputs.Length == 1)
{
return(true);
}
var shape = new TensorShape(layer.pool);
var permutations = shape.Get8DPermutationsForNCHWPermutationsAndShape(k_FromNCHWtoNHWC);
// Preserve symbolic shape by operating on int array instead of TensorShape, which would resolve unknown dimensions
layer.pool = TensorExtensions.Permute(layer.pool, permutations);
return(true);
}
int[] MergeTranspose(int[] transpose0, int[] tranpose1)
{
int[] permutations = new int[] { 0, 1, 2, 3, 4, 5, 6, 7 };
if (transpose0.Length == 4)
{
permutations[2] = TensorExtensions.Convert4DTo8DAxis(transpose0[0]);
permutations[5] = TensorExtensions.Convert4DTo8DAxis(transpose0[1]);
permutations[6] = TensorExtensions.Convert4DTo8DAxis(transpose0[2]);
permutations[7] = TensorExtensions.Convert4DTo8DAxis(transpose0[3]);
}
else
{
permutations[0] = transpose0[0];
permutations[1] = transpose0[1];
permutations[2] = transpose0[2];
permutations[3] = transpose0[3];
permutations[4] = transpose0[4];
permutations[5] = transpose0[5];
permutations[6] = transpose0[6];
permutations[7] = transpose0[7];
}
int[] combinePermutations = new int[] { 0, 1, 2, 3, 4, 5, 6, 7 };
if (tranpose1.Length == 4)
{
combinePermutations[2] = TensorExtensions.Convert4DTo8DAxis(tranpose1[0]);
combinePermutations[5] = TensorExtensions.Convert4DTo8DAxis(tranpose1[1]);
combinePermutations[6] = TensorExtensions.Convert4DTo8DAxis(tranpose1[2]);
combinePermutations[7] = TensorExtensions.Convert4DTo8DAxis(tranpose1[3]);
}
else
{
combinePermutations[0] = tranpose1[0];
combinePermutations[1] = tranpose1[1];
combinePermutations[2] = tranpose1[2];
combinePermutations[3] = tranpose1[3];
combinePermutations[4] = tranpose1[4];
combinePermutations[5] = tranpose1[5];
combinePermutations[6] = tranpose1[6];
combinePermutations[7] = tranpose1[7];
}
permutations = TensorExtensions.Permute(permutations, combinePermutations);
return(permutations);
}
static public TensorShape?InferOutputShapeNCHW(Layer layer, int[] inputRanks, TensorShape[] inputShapes)
{
switch (layer.type)
{
case Layer.Type.Conv3D:
{
TensorShape X = inputShapes[0];
// N C D H W, constructor is N D H W C
// => N = N C = D, D = H, H = W, W = C
// TODO helper function for that
X = new TensorShape(X.batch, X.height, X.width, X.channels, X.depth);
var K = layer.datasets[0].shape;
Assert.IsNotNull(layer.stride);
Assert.IsNotNull(layer.pad);
var pad = X.AdjustPadToKernel(K, layer.stride, layer.pad);
var O = X.ApplyKernel(K, layer.stride, pad);
return(new TensorShape(O.batch, O.channels, O.depth, O.height, O.width));
}
case Layer.Type.Conv2D:
case Layer.Type.DepthwiseConv2D:
{
TensorShape X = inputShapes[0];
// N C H W, constructor is N H W C
// => N = N C = H, H = W, H = C
// TODO helper function for that
X = new TensorShape(X.batch, X.width, X.channels, X.height);
var K = layer.datasets[0].shape;
Assert.IsNotNull(layer.stride);
Assert.IsNotNull(layer.pad);
var pad = X.AdjustPadToKernel(K, layer.stride, layer.pad);
var O = X.ApplyKernel(K, layer.stride, pad);
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.Conv2DTrans:
{
TensorShape X = inputShapes[0];
// N C H W, constructor is N H W C
// => N = N C = H, H = W, H = C
// TODO helper function for that
X = new TensorShape(X.batch, X.width, X.channels, X.height);
var K = layer.datasets[0].shape;
Assert.IsNotNull(layer.stride);
Assert.IsNotNull(layer.pad);
// pool size is treated as output_adjustment aka output_padding here
var outputAdjustment = layer.pool;
var pad = X.AdjustPadToKernel(K, layer.stride, layer.pad);
var O = X.ApplyKernelInverse(K, layer.stride, pad, outputAdjustment);
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.GlobalMaxPool2D:
case Layer.Type.GlobalAvgPool2D:
{
TensorShape X = inputShapes[0];
int rankX = inputRanks[0];
List <int> xShape = ShapeToOnnxLayout(X, rankX);
for (int i = 2; i < xShape.Count; i++)
{
xShape[i] = 1;
}
return(OnnxLayoutToTensorShape(xShape.ToArray()));
}
case Layer.Type.Dense:
{
TensorShape X = inputShapes[0];
X = new TensorShape(X.batch, X.width, X.channels, X.height);
Assert.IsNotNull(layer.datasets);
var W = layer.datasets[0].shape;
var O = new TensorShape(X.flatHeight, W.flatWidth);
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.MatMul:
{
TensorShape X = inputShapes[0];
int rankX = inputRanks[0];
List <int> xShape = ShapeToOnnxLayout(X, rankX);
TensorShape Y = inputShapes[1];
int rankY = inputRanks[1];
List <int> yShape = ShapeToOnnxLayout(Y, rankY);
int rankO = Mathf.Max(rankX, rankY);
for (int i = 0; i < rankO - rankX; i++)
{
xShape.Insert(0, 1);
}
for (int i = 0; i < rankO - rankY; i++)
{
yShape.Insert(0, 1);
}
List <int> oShape = new List <int>();
for (int i = 0; i < rankO - 2; i++)
{
oShape.Add(Mathf.Max(xShape[i], yShape[i]));
}
oShape.Add(xShape[rankO - 2]);
oShape.Add(yShape[rankO - 1]);
return(OnnxLayoutToTensorShape(oShape.ToArray()));
}
case Layer.Type.Border3D:
{
TensorShape X = inputShapes[0];
X = new TensorShape(X.batch, X.height, X.width, X.channels, X.depth);
Assert.IsNotNull(layer.pad);
var O = X.ApplyBorder(layer.pad);
return(new TensorShape(O.batch, O.channels, O.depth, O.height, O.width));
}
case Layer.Type.Border2D:
case Layer.Type.Pad2DReflect:
case Layer.Type.Pad2DSymmetric:
case Layer.Type.Pad2DEdge:
{
if (inputShapes.Length > 1)
{
return(null);
}
TensorShape X = inputShapes[0];
X = new TensorShape(X.batch, X.width, X.channels, X.height);
Assert.IsNotNull(layer.pad);
var O = X.ApplyBorder(layer.pad);
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.Upsample2D:
{
if (inputShapes.Length > 1)
{
return(null);
}
TensorShape X = inputShapes[0];
// pool size is treated as upsample coefficient here
Assert.IsNotNull(layer.pool);
Assert.AreEqual(layer.pool.Length, 4);
return(new TensorShape(X.batch * layer.pool[0], X.height * layer.pool[1], X.width * layer.pool[2], X.channels * layer.pool[3]));
}
case Layer.Type.Upsample3D:
{
if (inputShapes.Length > 1)
{
return(null);
}
TensorShape X = inputShapes[0];
// pool size is treated as upsample coefficient here
Assert.IsNotNull(layer.pool);
Assert.AreEqual(layer.pool.Length, 5);
return(new TensorShape(X.batch * layer.pool[0], X.depth * layer.pool[1], X.height * layer.pool[2], X.width * layer.pool[3], X.channels * layer.pool[4]));
}
case Layer.Type.Resample2D:
{
TensorShape X = inputShapes[0];
if (inputShapes.Length > 1)
{
return(null);
}
else
{
// pool is treated as resample size here
var size = layer.pool;
Assert.IsNotNull(size);
Assert.AreEqual(size.Length, 4);
return(new TensorShape(size[0], size[1], size[2], size[3]));
}
}
case Layer.Type.TopKIndices:
case Layer.Type.TopKValues:
{
// Calculated at runtime: same shape as input 0 with k elements in the dimension specified by axis
return(null);
}
case Layer.Type.NonMaxSuppression:
{
int maxOutputBoxesPerClass = 0;
if (layer.pool.Length > 0)
{
maxOutputBoxesPerClass = layer.pool[0];
}
if (maxOutputBoxesPerClass <= 0)
{
return(null);
}
return(new TensorShape(maxOutputBoxesPerClass, 3));
}
case Layer.Type.NonZero:
{
// Calculated at runtime
return(null);
}
case Layer.Type.Add:
case Layer.Type.Sub:
case Layer.Type.Mul:
case Layer.Type.Div:
case Layer.Type.Pow:
case Layer.Type.Min:
case Layer.Type.Max:
case Layer.Type.Mean:
case Layer.Type.Greater:
case Layer.Type.GreaterEqual:
case Layer.Type.Less:
case Layer.Type.LessEqual:
case Layer.Type.Equal:
case Layer.Type.LogicalOr:
case Layer.Type.LogicalAnd:
case Layer.Type.LogicalXor:
{
int rankO = 0;
for (int i = 0; i < inputRanks.Length; i++)
{
rankO = Mathf.Max(inputRanks[i], rankO);
}
var O = new List <int>();
for (int i = 0; i < rankO; i++)
{
O.Add(1);
}
for (int i = 0; i < inputShapes.Length; i++)
{
TensorShape X = inputShapes[i];
int rankX = inputRanks[i];
List <int> xShape = ShapeToOnnxLayout(X, rankX);
for (int k = 0; k < rankO - rankX; k++)
{
xShape.Insert(0, 1);
}
for (int k = 0; k < rankO; k++)
{
O[k] = Math.Max(O[k], xShape[k]);
}
}
return(OnnxLayoutToTensorShape(O.ToArray()));
}
case Layer.Type.Range:
{
return(null); // only const support
}
case Layer.Type.ReduceL1:
case Layer.Type.ReduceL2:
case Layer.Type.ReduceLogSum:
case Layer.Type.ReduceLogSumExp:
case Layer.Type.ReduceMax:
case Layer.Type.ReduceMean:
case Layer.Type.ReduceMin:
case Layer.Type.ReduceProd:
case Layer.Type.ReduceSum:
case Layer.Type.ReduceSumSquare:
case Layer.Type.ArgMax:
case Layer.Type.ArgMin:
{
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
var xShape = ShapeToOnnxLayout(X, rank);
var axis = layer.axis;
if (axis < 0)
{
axis = rank + axis;
}
xShape[axis] = 1;
if (layer.alpha != 1.0f) // keepdim == 0
{
xShape.RemoveAt(axis);
}
return(OnnxLayoutToTensorShape(xShape.ToArray()));
}
case Layer.Type.Transpose:
{
TensorShape X = inputShapes[0];
var permutations = layer.pool;
if (permutations == null)
{
return(new TensorShape(X.batch, X.width));
}
else
{
int rank = inputRanks[0];
List <int> xShape = ShapeToOnnxLayout(X, rank);
// Permutations may already be in padded form for op purposes, so strip down to match rank
permutations = permutations.Take(rank).ToArray();
var oShape = TensorExtensions.Permute(xShape.ToArray(), permutations);
return(OnnxLayoutToTensorShape(oShape));
}
}
case Layer.Type.MaxPool2D:
case Layer.Type.AvgPool2D:
{
TensorShape X = inputShapes[0];
X = new TensorShape(X.batch, X.width, X.channels, X.height);
Assert.IsNotNull(layer.pool);
Assert.IsNotNull(layer.stride);
Assert.IsNotNull(layer.pad);
var pad = X.AdjustPadToPool(layer.pool, layer.stride, layer.pad);
var O = X.ApplyPool(layer.pool, layer.stride, pad);
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.Load:
{
return(layer.datasets[0].shape);
}
case Layer.Type.DepthToSpace:
{
TensorShape X = inputShapes[0];
X = new TensorShape(X.batch, X.width, X.channels, X.height);
// pool size is treated as blocksize here
Assert.IsNotNull(layer.pool);
Assert.AreEqual(layer.pool.Length, 2);
Assert.AreEqual(X.channels % (layer.pool[0] * layer.pool[1]), 0);
var O = new TensorShape(X.batch, X.height * layer.pool[1], X.width * layer.pool[0], X.channels / (layer.pool[0] * layer.pool[1]));
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.SpaceToDepth:
{
TensorShape X = inputShapes[0];
X = new TensorShape(X.batch, X.width, X.channels, X.height);
// pool size is treated as blocksize here
Assert.IsNotNull(layer.pool);
Assert.AreEqual(layer.pool.Length, 2);
var O = new TensorShape(X.batch, X.height / layer.pool[1], X.width / layer.pool[0], X.channels * (layer.pool[0] * layer.pool[1]));
return(new TensorShape(O.batch, O.channels, O.height, O.width));
}
case Layer.Type.RandomNormal:
case Layer.Type.RandomUniform:
{
Assert.IsNotNull(layer.pool);
// pool size is treated as shape constant, if not empty
// otherwise shape of the previous tensor is used
if (layer.pool.Length > 0)
{
return(new TensorShape(layer.pool));
}
else
{
return(inputShapes[0]);
}
}
case Layer.Type.Multinomial:
{
TensorShape X = inputShapes[0];
Assert.IsNotNull(layer.pool);
Assert.AreEqual(layer.pool.Length, 1);
return(new TensorShape(X.batch, layer.pool[0]));
}
case Layer.Type.OneHot:
{
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
var nchwShape = ShapeToOnnxLayout(X, rank);
int depth = layer.pool[0];
nchwShape.Add(depth);
for (int i = 0; i < 4 - rank - 1; i++)
{
nchwShape.Add(1);
}
return(new TensorShape(nchwShape[0], nchwShape[1], nchwShape[2], nchwShape[3]));
}
case Layer.Type.LSTM:
{
TensorShape X = inputShapes[0];
var nchwShape = new List <int> {
X.batch, X.height, X.width, X.channels
};
int hiddenSize = layer.pool[0];
// The first output, Y, is rank 4; Other outputs are handled as identity layers
return(new TensorShape(nchwShape[0], 1, nchwShape[1], hiddenSize));
}
case Layer.Type.Flatten:
{
TensorShape X = inputShapes[0];
return(X.Flatten());
}
case Layer.Type.Tile:
{
if (inputShapes.Length > 1)
{
return(null);
}
var inputShape = ShapeToOnnxLayout(inputShapes[0], inputRanks[0]);
var scale = layer.pool.ToArray();
Assert.IsNotNull(scale);
Assert.AreEqual(scale.Length, inputShape.Count);
for (int i = 0; i < scale.Length; i++)
{
scale[i] *= inputShape[i];
}
return(OnnxLayoutToTensorShape(scale));
}
case Layer.Type.ConstantOfShape:
{
if (layer.axis == 1)
{
return(inputShapes[0]);
}
if (inputShapes.Length == 1)
{
return(null);
}
else
{
return(OnnxLayoutToTensorShape(layer.pool));
}
}
case Layer.Type.Reshape:
{
if (inputShapes.Length > 1)
{
return(null);
}
// TODO shape to onnx shape given rank
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
var nchwShape = ShapeToOnnxLayout(X, rank);
var unknownIndex = -1;
var multipleOf = 1;
var size = layer.pool.ToArray();
for (var i = 0; i < size.Length; ++i)
{
if (size[i] == 0)
{
size[i] = nchwShape[i];
}
if (size[i] < 0)
{
unknownIndex = i;
}
else
{
multipleOf *= size[i];
}
}
if (unknownIndex != -1)
{
size[unknownIndex] = X.length / multipleOf;
}
return(OnnxLayoutToTensorShape(size));
}
case Layer.Type.Expand:
{
if (inputShapes.Length > 1)
{
return(null);
}
var size = layer.pool.ToList();
var inputShape = ShapeToOnnxLayout(inputShapes[0], inputRanks[0]);
int rankO = Math.Max(size.Count, inputShape.Count);
for (int i = 0; i < rankO - size.Count; i++)
{
size.Insert(0, 1);
}
for (int i = 0; i < rankO - inputShape.Count; i++)
{
inputShape.Insert(0, 1);
}
var tiledShape = new int[rankO];
for (int i = 0; i < rankO; i++)
{
tiledShape[i] = Mathf.Max(size[i], inputShape[i]);
}
return(OnnxLayoutToTensorShape(tiledShape));
}
case Layer.Type.Concat:
{
var shape = ShapeToOnnxLayout(inputShapes[0], inputRanks[0]);
var axis = layer.axis;
if (axis < 0)
{
axis += inputRanks[0];
}
for (int i = 1; i < inputShapes.Length; i++)
{
var shapei = ShapeToOnnxLayout(inputShapes[i], inputRanks[i]);
shape[axis] += shapei[axis];
}
return(OnnxLayoutToTensorShape(shape.ToArray()));
}
case Layer.Type.Gather:
{
var input0Shape = inputShapes[0];
var input1Shape = inputShapes[1];
int rank0 = inputRanks[0];
int rank1 = inputRanks[1];
var shape = ShapeToOnnxLayout(input0Shape, rank0);
var indicies = ShapeToOnnxLayout(input1Shape, rank1);
var axis = layer.axis;
if (axis < 0)
{
axis += rank0;
}
shape.InsertRange(axis, indicies);
return(OnnxLayoutToTensorShape(shape.ToArray()));
}
// elementwise operations
case Layer.Type.Nop:
case Layer.Type.ScaleBias:
case Layer.Type.Normalization:
case Layer.Type.LRN:
case Layer.Type.Dropout:
case Layer.Type.LogicalNot:
case Layer.Type.Sign:
case Layer.Type.Where:
{
// works in place, keeps the same shape size
return(inputShapes[0]);
}
case Layer.Type.Activation:
{
TensorShape X = inputShapes[0];
// LSTMs have multiple outputs, so deal with those separately
if (layer.activation == Layer.Activation.None && layer.pad.Length > 0 &&
layer.name.IndexOf("lstm", StringComparison.OrdinalIgnoreCase) >= 0)
{
int rank = layer.pad[0];
switch (rank)
{
case 4:
// Y
return(X);
case 3:
// Y_h, Y_c: seq_length is stripped off
return(new TensorShape(X[1], X[2], X[3]));
}
}
// works in place, keeps the same shape size
return(X);
}
case Layer.Type.Shape:
{
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
return(new TensorShape(rank));
}
case Layer.Type.Squeeze:
{
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
var nchwShape = ShapeToOnnxLayout(X, rank);
var squeezedShape = new List <int>();
for (int i = 0; i < nchwShape.Count; i++)
{
if (!layer.pool.Contains(i))
{
squeezedShape.Add(nchwShape[i]);
}
}
return(OnnxLayoutToTensorShape(squeezedShape.ToArray()));
}
case Layer.Type.Unsqueeze:
{
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
if (rank < 0)
{
return(null);
}
var nchwShape = ShapeToOnnxLayout(X, rank);
if (rank == 0)
{
return(new TensorShape(new int[] { 1, 1, 1, 1 }));
}
for (int a = 0; a < layer.pool.Length; a++)
{
var axis = layer.pool[a];
if (axis < 0)
{
axis += rank;
}
nchwShape.Insert(axis, 1);
rank++;
}
return(OnnxLayoutToTensorShape(nchwShape.ToArray()));
}
case Layer.Type.StridedSlice:
{
if (inputShapes.Length > 1)
{
return(null);
}
TensorShape X = inputShapes[0];
int rank = inputRanks[0];
var nchwShape = ShapeToOnnxLayout(X, rank);
var starts = layer.pad.ToArray();
var ends = layer.pool.ToArray();
var steps = layer.stride.ToArray();
var axes = layer.axes.ToArray();
var onnxStarts = Enumerable.Repeat(0, rank).ToArray();
var onnxEnds = Enumerable.Repeat(int.MaxValue, rank).ToArray(); // by default copy the whole axis till the end
var onnxSteps = Enumerable.Repeat(1, rank).ToArray();
// NOTE: begin=0, end=0, stride=1 <= full range from existing axis
// begin=0, end=inf,stride=1 <= full range from existing axis
// begin=0, end=X, stride=1 <= full range from existing axis, if X==last element on this axis
// begin=0, end=0, stride=0 <= new axis OR shrink axis to single 1st element
// begin=N, end=N, stride=0 <= shrink axis to single Nth element
// These notes are copied from TensorExtensions.ApplyStridedSlice(...)
for (int i = 0; i < axes.Length; ++i)
{
var axis = axes[i];
if (axis < 0)
{
axis += rank;
}
axis = Math.Min(Math.Max(axis, 0), rank);
onnxStarts[axis] = starts[i];
onnxEnds[axis] = ends[i];
onnxSteps[axis] = steps[i];
}
var sliced = new int[rank];
for (int i = 0; i < rank; ++i)
{
// NOTE: begin=0, end=0, stride=1 <= full range from the existing axis
// begin=0, end=X, stride=1 <= full range from the existing axis, if X==last element on this axis
// begin=0, end=0, stride=0 <= new axis OR shrink axis to a single 1st element
// begin=N, end=N, stride=0 <= shrink axis to a single Nth element
int ei = TensorExtensions.WrapIndex(onnxEnds[i], nchwShape[i]);
int si = TensorExtensions.WrapIndex(onnxStarts[i], nchwShape[i]);
if (onnxSteps[i] > 0)
{
sliced[i] = (int)Mathf.Round((float)(Math.Min(ei, nchwShape[i]) - Math.Min(si, nchwShape[i] - 1)) / (float)(Mathf.Abs(onnxSteps[i])));
}
else
{
bool inclusive = onnxEnds[i] < -nchwShape[i]; // edge case when ends is negative and bigger than nchwShape
sliced[i] = (int)Mathf.Round((float)(Math.Min(si, nchwShape[i] - 1) - Math.Min(ei, nchwShape[i]) + (inclusive ? 1 : 0)) / (float)(Mathf.Abs(onnxSteps[i])));
}
}
return(OnnxLayoutToTensorShape(sliced.ToArray()));
}
default:
throw new NotImplementedException("InferOutputShapeNCHW: Unhandled layer: " + layer.ToString());
}
}
void Rewrite(ref Model model)
{
IRShapeInferenceHelper.RankInference.ListTemporaryTensorRanks(model, out m_RanksByName);
var inputShapes = new Dictionary <string, TensorShape>();
foreach (var i in model.inputs)
{
if (!ModelAnalyzer.IsInputShapeAcceptablyKnowForShapeInference(i))
{
continue;
}
inputShapes.Add(i.name, new TensorShape(i.shape));
}
IRShapeInferenceHelper.ShapeInference.ListTemporaryTensorShapesNCHW(model, inputShapes, ref m_RanksByName, out m_ShapesByName);
var nhwc = model.ShallowCopy();
nhwc.layers.Clear();
nhwc.layout = "NHWC";
// TF2ONNX transpose pattern -> part of the model are in NHWC and not NCHW
// * identify those
// * transpose inputs to NCHW
// * remove layout transposes
// * convert axis/constants accordingly
LayoutTransposeRemovalHelper transposeRemoval = new LayoutTransposeRemovalHelper();
m_isModelExportedFromNHWC = transposeRemoval.InferAllLayersChannelOrder(model, out m_layersChannelOrder);
if (m_isModelExportedFromNHWC && !transposeRemoval.IsImporterLikelyNHWCLayout(model.ProducerName))
{
nhwc.Warnings.Add(new Model.ImporterWarning("model", "model detected as NCHW, but not natively in this layout, behavior might be erroneous"));
}
// remove layout change transposes
if (m_isModelExportedFromNHWC)
{
transposeRemoval.RemoveAllChannelLayoutTransposes(ref model, m_layersChannelOrder);
}
var modelBuilder = new ModelBuilder(nhwc);
for (int i = 0; i < nhwc.inputs.Count; i++)
{
Model.Input input = nhwc.inputs[i];
int[] shape = input.shape;
var tensorShape = new TensorShape(shape);
int[] rankPermutations = GetChannelsLastPermutationsFromRank(input.rank);
int[] permutations = tensorShape.Get8DPermutationsForNCHWPermutationsAndShape(rankPermutations);
// Preserve symbolic shape by operating on int array instead of TensorShape, which would resolve unknown dimensions
if (m_isModelExportedFromNHWC) // transpose input shape if importer preserved NHWC layout
{
if (m_layersChannelOrder[input.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NCHW)
{
input.shape = TensorExtensions.Permute(shape, permutations);
}
else
{
var onnxShape = new List <int> {
shape[2], shape[5], shape[6], shape[7]
};
onnxShape.RemoveRange(input.rank, 4 - input.rank);
input.shape = IRShapeInferenceHelper.ShapeInference.BarracudaLayoutToTensorShapeLayout(onnxShape.ToArray());
}
}
else
{
input.shape = TensorExtensions.Permute(shape, permutations);
}
nhwc.inputs[i] = input;
}
// NCHW -> Barracuda NHWC rewriter (some layer need to insert aditional layers to be Barracuda compatible)
var rewriters = InstantiateRewriterNCHWToNHWC();
// NHWC -> Barracuda NHWC rewriter (axis and constant padding padding)
var rewritersNHWC = InstantiateRewriterNHWCToNHWC();
foreach (var l in model.layers)
{
// Some nodes output multiple layers (e.g. LSTM), so don't process or include those layers
if (nhwc.layers.Exists(alreadyOutputLayer => alreadyOutputLayer.name == l.name))
{
continue;
}
if (m_layersChannelOrder.TryGetValue(l.name, out LayoutTransposeRemovalHelper.ChannelsOrder layerChannelOrder))
{
if (m_isModelExportedFromNHWC && (layerChannelOrder == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC))
{
if (!rewritersNHWC.TryGetValue(l.type, out Func <Layer, ModelBuilder, bool> rwNCHW) || rwNCHW(l, modelBuilder))
{
nhwc.layers.Add(l);
}
continue;
}
}
if (!rewriters.TryGetValue(l.type, out Func <Layer, ModelBuilder, bool> rw) || rw(l, modelBuilder))
{
// Either no re-write was needed or the layer was not replaced
nhwc.layers.Add(l);
}
}
// We need to correct constants to have broadcast work correctly
// ONNX: 1,64,32 + c:32
// Barracuda: 1,_32,64 + c:_,_,32,64 and not c:32,_,_,_
// X:5,7 + c: 6,9,5,7 = 6,9,5,7
// X: 5,_,_,7 + c: 6,5,7,9 = ???
CorrectConstantsForBroadCast(ref nhwc);
CorrectDynamicInputsForBroadCast(ref nhwc);
// for NHWC importers, perform slightly more aggressive output shape check
// => add transposes to match onnx layout
if (transposeRemoval.IsImporterLikelyNHWCLayout(model.ProducerName))
{
CorrectOutputLayoutToMatchNHWCLayout(ref nhwc);
}
model = nhwc;
}
void ConcatenateTransposes(ref Model model)
{
var remap = new Dictionary<string, string>();
var layerReferences = new Dictionary<string, int>();
for (int l = 0; l < model.layers.Count; ++l)
{
Layer layer = model.layers[l];
string[] layerInputs = layer.inputs;
for (int i = 0; i < layerInputs.Length; i++)
{
if (layerReferences.TryGetValue(layerInputs[i], out int count))
count++;
else
count = 0;
layerReferences[layerInputs[i]] = count;
}
if (model.outputs.Contains(layer.name))
{
if (layerReferences.TryGetValue(layer.name, out int count))
count++;
else
count = 0;
layerReferences[layer.name] = count;
}
}
for (int l = 0; l < model.layers.Count - 1; ++l)
{
var layer = model.layers[l];
var nextLayer = model.layers[l + 1];
if (remap.ContainsKey(layer.name)) // This layer will get removed
continue;
string[] layerInputs = layer.inputs;
for (int i = 0; i < layerInputs.Length; i++)
{
if (remap.TryGetValue(layerInputs[i], out string replacement))
layerInputs[i] = replacement;
}
if (layer.flags.HasFlag(Layer.Flags.Preserve))
continue;
bool reverseMerge = nextLayer.flags.HasFlag(Layer.Flags.Preserve);
// Only concatenate serial transpose layers
if (layer.type == Layer.Type.Transpose
&& nextLayer.type == Layer.Type.Transpose
&& nextLayer.inputs.Contains(layer.name)
&& layerReferences.TryGetValue(layer.name, out int references)
&& references <= 1
&& layerReferences.TryGetValue(nextLayer.name, out references)
&& references <= 1)
{
int[] permutations = new int[] { 0, 1, 2, 3, 4, 5, 6, 7 };
if (layer.pool.Length == 4)
{
permutations[2] = TensorExtensions.Convert4DTo8DAxis(layer.pool[0]);
permutations[5] = TensorExtensions.Convert4DTo8DAxis(layer.pool[1]);
permutations[6] = TensorExtensions.Convert4DTo8DAxis(layer.pool[2]);
permutations[7] = TensorExtensions.Convert4DTo8DAxis(layer.pool[3]);
}
else
{
permutations[0] = layer.pool[0];
permutations[1] = layer.pool[1];
permutations[2] = layer.pool[2];
permutations[3] = layer.pool[3];
permutations[4] = layer.pool[4];
permutations[5] = layer.pool[5];
permutations[6] = layer.pool[6];
permutations[7] = layer.pool[7];
}
int[] combinePermutations = new int[] { 0, 1, 2, 3, 4, 5, 6, 7 };
if (nextLayer.pool.Length == 4)
{
combinePermutations[2] = TensorExtensions.Convert4DTo8DAxis(nextLayer.pool[0]);
combinePermutations[5] = TensorExtensions.Convert4DTo8DAxis(nextLayer.pool[1]);
combinePermutations[6] = TensorExtensions.Convert4DTo8DAxis(nextLayer.pool[2]);
combinePermutations[7] = TensorExtensions.Convert4DTo8DAxis(nextLayer.pool[3]);
}
else
{
combinePermutations[0] = nextLayer.pool[0];
combinePermutations[1] = nextLayer.pool[1];
combinePermutations[2] = nextLayer.pool[2];
combinePermutations[3] = nextLayer.pool[3];
combinePermutations[4] = nextLayer.pool[4];
combinePermutations[5] = nextLayer.pool[5];
combinePermutations[6] = nextLayer.pool[6];
combinePermutations[7] = nextLayer.pool[7];
}
permutations = TensorExtensions.Permute(permutations, combinePermutations);
if (reverseMerge)
{
remap[layer.name] = nextLayer.name;
nextLayer.pool = permutations;
nextLayer.inputs = layer.inputs.ToArray();
}
else
{
layer.pool = permutations;
remap[nextLayer.name] = layer.name;
}
}
}
var removeLayers = remap.Keys;
model.layers.RemoveAll(l => removeLayers.Contains(l.name));
}
void ConcatenateTransposes(ref Model model)
{
var remap = new Dictionary <string, string>();
var layerReferences = new Dictionary <string, int>();
for (int l = 0; l < model.layers.Count - 1; ++l)
{
Layer layer = model.layers[l];
string[] layerInputs = layer.inputs;
for (int i = 0; i < layerInputs.Length; i++)
{
if (layerReferences.TryGetValue(layerInputs[i], out int count))
{
count++;
}
else
{
count = 0;
}
layerReferences[layerInputs[i]] = count;
}
}
for (int l = 0; l < model.layers.Count - 1; ++l)
{
var layer = model.layers[l];
var nextLayer = model.layers[l + 1];
if (layer.flags.HasFlag(Layer.Flags.Preserve))
{
continue;
}
if (remap.ContainsKey(layer.name)) // This layer will get removed
{
continue;
}
string[] layerInputs = layer.inputs;
for (int i = 0; i < layerInputs.Length; i++)
{
if (remap.TryGetValue(layerInputs[i], out string replacement))
{
layerInputs[i] = replacement;
}
}
// Only concatenate serial transpose layers
if (layer.type == Layer.Type.Transpose &&
nextLayer.type == Layer.Type.Transpose &&
nextLayer.inputs.Contains(layer.name) &&
layerReferences.TryGetValue(layer.name, out int references) &&
references <= 1 &&
layerReferences.TryGetValue(nextLayer.name, out references) &&
references <= 1)
{
int[] permutations = layer.pool;
int[] combinePermutations = nextLayer.pool;
permutations = TensorExtensions.Permute(permutations, combinePermutations);
layer.pool = permutations;
remap[nextLayer.name] = layer.name;
}
}
var removeLayers = remap.Keys;
model.layers.RemoveAll(l => removeLayers.Contains(l.name));
}
Dictionary <Layer.Type, Func <Layer, ModelBuilder, bool> > InstantiateRewriterNCHWToNHWC()
{
var rewriters = new Dictionary <Layer.Type, Func <Layer, ModelBuilder, bool> >();
// return true if layer should be included in rewritten model, false if it was replaced
rewriters.Add(Layer.Type.Load, ConvertDatasets);
rewriters.Add(Layer.Type.Reshape, (layer, net) =>
{
// TODO reshape with pool as constant
string input0 = layer.inputs[0];
if (!m_RanksByName.TryGetValue(input0, out int?input0Rank) || !input0Rank.HasValue)
{
throw new Exception($"Must have input rank for {input0} in order to convert Reshape to NHWC");
}
int outputRank = 4;
Layer nchwTranspose;
// TODO cleanup?
if (input0Rank.Value == 1)
{
nchwTranspose = net.Identity($"Transpose_{input0}_For_{layer.name}", input0);
}
else if (input0Rank.Value == 2)
{
nchwTranspose = net.Transpose($"Transpose_{input0}_For_{layer.name}", input0, k_FromNHWCtoNCHW);
}
else if (input0Rank.Value == 3)
{
nchwTranspose = net.Transpose($"Transpose_{input0}_For_{layer.name}", input0, k_FromN1WCtoNCH);
}
else if (input0Rank.Value == 4)
{
nchwTranspose = net.Transpose($"Transpose_{input0}_For_{layer.name}", input0, k_FromNHWCtoNCHW);
}
else if (input0Rank.Value == 5)
{
nchwTranspose = net.Transpose($"Transpose_{input0}_For_{layer.name}", input0, new[] { 0, 1, 2, 3, 7, 4, 5, 6 });
}
else
{
// TODO 8D?
nchwTranspose = net.Transpose($"Transpose_{input0}_For_{layer.name}", input0, new[] { 0, 1, 2, 7, 3, 4, 5, 6 });
}
Layer reshape = null;
if (layer.inputs.Length > 1)
{
string input1 = layer.inputs[1];
if (!m_RanksByName.TryGetValue(input1, out int?input1Rank) || !input1Rank.HasValue)
{
throw new Exception($"Must have input rank for {input1} in order to convert Reshape to NHWC");
}
if (input1Rank.Value == 1) // shape is in the tensor
{
if (!m_ShapesByName.TryGetValue(input1, out TensorShape? input1Shape) || !input1Shape.HasValue)
{
throw new Exception($"Must have input shape for {input1} in order to convert Reshape to NHWC");
}
outputRank = input1Shape.Value[TensorShape.DataBatch];
}
reshape = net.Reshape($"{layer.name}_NCHW", nchwTranspose, input1);
}
else if (layer.pool.Length > 0)
{
outputRank = layer.pool.Length;
var shape = IRShapeInferenceHelper.ShapeInference.OnnxLayoutToTensorShapeLayout(layer.pool);
reshape = net.Reshape($"{layer.name}_NCHW", nchwTranspose, shape);
}
// TODO cleanup?
if (outputRank == 1)
{
nchwTranspose = net.Identity(layer.name, reshape);
}
else if (outputRank == 2)
{
nchwTranspose = net.Transpose(layer.name, reshape, k_FromNCHWtoNHWC);
}
else if (outputRank == 3)
{
net.Transpose(layer.name, reshape, k_FromNCHtoN1WC);
}
else if (outputRank == 4)
{
net.Transpose(layer.name, reshape, k_FromNCHWtoNHWC);
}
else if (outputRank == 5)
{
net.Transpose(layer.name, reshape, new[] { 0, 1, 2, 3, 5, 6, 7, 4 });
}
else
{
// TODO 8D?
net.Transpose(layer.name, reshape, new[] { 0, 1, 2, 4, 5, 6, 7, 3 });
}
return(false);
});
rewriters.Add(Layer.Type.Expand, ConvertShape);
rewriters.Add(Layer.Type.Shape, (layer, net) =>
{
if (layer.axis >= 0)
{
ConvertAxis(layer, net);
}
return(true);
});
rewriters.Add(Layer.Type.Transpose, (layer, net) =>
{
int[] permutations = layer.pool;
int rank = layer.pool.Length;
int[] onnxTranspose = layer.pool;
// TODO cleanup?
switch (rank)
{
case 2:
{
// onnx : 5,7 => 5,7 / 7,5
// barracuda : 5,_,_,7 => 5,_,_,7 / 7,_,_,5
layer.pool = new[] { 0, 1, 2, 3 };
layer.pool[0] = onnxTranspose[0] == 1 ? 3 : onnxTranspose[0];
layer.pool[3] = onnxTranspose[1] == 1 ? 3 : onnxTranspose[1];
return(true);
}
case 3:
{
// onnx : 5,7,3 => 5,7,3 / 7,5,3 / 7,3,5 ...
// barracuda : 5,_,7,3 => 7,_,3,5 / 7,_,5,3 ...
layer.pool = new[] { 0, 1, 2, 3 };
layer.pool[0] = onnxTranspose[0] == 1 ? 3 : onnxTranspose[0] == 2 ? 2 : onnxTranspose[0];
layer.pool[3] = onnxTranspose[1] == 1 ? 3 : onnxTranspose[1] == 2 ? 2 : onnxTranspose[1];
layer.pool[2] = onnxTranspose[2] == 1 ? 3 : onnxTranspose[2] == 2 ? 2 : onnxTranspose[2];
return(true);
}
case 4:
{
layer.pool = new[] { 0, 1, 2, 3 };
layer.pool[0] = onnxTranspose[0] == 1 ? 3 : onnxTranspose[0] == 2 ? 1 : onnxTranspose[0] == 3 ? 2 : onnxTranspose[0];
layer.pool[3] = onnxTranspose[1] == 1 ? 3 : onnxTranspose[1] == 2 ? 1 : onnxTranspose[1] == 3 ? 2 : onnxTranspose[1];
layer.pool[1] = onnxTranspose[2] == 1 ? 3 : onnxTranspose[2] == 2 ? 1 : onnxTranspose[2] == 3 ? 2 : onnxTranspose[2];
layer.pool[2] = onnxTranspose[3] == 1 ? 3 : onnxTranspose[3] == 2 ? 1 : onnxTranspose[3] == 3 ? 2 : onnxTranspose[3];
return(true);
}
case 5:
{
// onnx : 5,7,3,4,9 => 5,9,4,7,3 / 3,9,4,7,5 ...
layer.pool = new[] { 0, 1, 2, 3, 4, 5, 6, 7 };
// [1,1,N,1,D,H,W,C]
layer.pool[2] = onnxTranspose[0] == 0 ? 2 : onnxTranspose[0] == 1 ? 7 : onnxTranspose[0] + 2;
layer.pool[7] = onnxTranspose[1] == 0 ? 2 : onnxTranspose[1] == 1 ? 7 : onnxTranspose[1] + 2;
layer.pool[4] = onnxTranspose[2] == 0 ? 2 : onnxTranspose[2] == 1 ? 7 : onnxTranspose[2] + 2;
layer.pool[5] = onnxTranspose[3] == 0 ? 2 : onnxTranspose[3] == 1 ? 7 : onnxTranspose[3] + 2;
layer.pool[6] = onnxTranspose[4] == 0 ? 2 : onnxTranspose[4] == 1 ? 7 : onnxTranspose[4] + 2;
return(true);
}
default:
{
// TODO 8D?
layer.pool = new[] { 0, 1, 2, 3, 4, 5, 6, 7 };
// NCTDHW
layer.pool[2] = onnxTranspose[0] == 0 ? 2 : onnxTranspose[0] == 1 ? 7 : onnxTranspose[0] + 1;
layer.pool[7] = onnxTranspose[1] == 0 ? 2 : onnxTranspose[1] == 1 ? 7 : onnxTranspose[1] + 1;
layer.pool[3] = onnxTranspose[2] == 0 ? 2 : onnxTranspose[2] == 1 ? 7 : onnxTranspose[2] + 1;
layer.pool[4] = onnxTranspose[3] == 0 ? 2 : onnxTranspose[3] == 1 ? 7 : onnxTranspose[3] + 1;
layer.pool[5] = onnxTranspose[4] == 0 ? 2 : onnxTranspose[4] == 1 ? 7 : onnxTranspose[4] + 1;
layer.pool[6] = onnxTranspose[5] == 0 ? 2 : onnxTranspose[5] == 1 ? 7 : onnxTranspose[5] + 1;
return(true);
}
}
});
rewriters.Add(Layer.Type.Unsqueeze, (layer, net) =>
{
// Replace w/ a Transpose since Barracuda tensors are full rank (i.e. grab an unused dimension)
string input0 = layer.inputs[0];
if (!m_RanksByName.TryGetValue(input0, out int?input0Rank) || !input0Rank.HasValue)
{
throw new Exception($"Must have input rank for {input0} in order to convert axis for Unsqueeze");
}
var axis = layer.pool[0];
if (axis < 0)
{
axis = input0Rank.Value + 1 - axis;
}
var transpose = UnSqueezeAxisPermutationForMappingNCHWLayoutToBarracuda(input0Rank.Value, axis);
net.Transpose(layer.name, input0, transpose);
return(false);
});
rewriters.Add(Layer.Type.Squeeze, (layer, net) =>
{
// Replace w/ a Transpose since Barracuda tensors are full rank
string input0 = layer.inputs[0];
if (!m_RanksByName.TryGetValue(input0, out int?input0Rank) || !input0Rank.HasValue)
{
throw new Exception($"Must have input rank for {input0} in order to convert axis for Squeeze");
}
var axis = layer.pool[0];
if (axis < 0)
{
axis = input0Rank.Value + 1 - axis;
}
var transpose = SqueezeAxisPermutationForMappingNCHWLayoutToBarracuda(input0Rank.Value, axis);
net.Transpose(layer.name, input0, transpose);
return(false);
});
rewriters.Add(Layer.Type.Flatten, (layer, net) =>
{
string input0 = layer.inputs[0];
if (!m_RanksByName.TryGetValue(input0, out int?input0Rank) || !input0Rank.HasValue)
{
throw new Exception($"Must have input rank for {input0} in order to convert Flatten to NHWC");
}
Layer nchwTranspose = net.Transpose($"Transpose_{input0}_For_{layer.name}", input0, input0Rank.Value == 3 ? k_FromN1WCtoNCH : k_FromNHWCtoNCHW);
net.Flatten(layer.name, nchwTranspose);
// No need to transpose back b/c final shape is always NC (rank 2)
return(false);
});
rewriters.Add(Layer.Type.Concat, ConvertAxis);
rewriters.Add(Layer.Type.StridedSlice, (layer, net) =>
{
int rank = 4;
if (m_RanksByName.ContainsKey(layer.name) && m_RanksByName[layer.name] != null)
{
rank = m_RanksByName[layer.name].Value;
}
var name = layer.name;
var starts = layer.pad;
var ends = layer.pool;
var steps = layer.stride;
var axes = layer.axes;
var onnxStarts = Enumerable.Repeat(0, rank).ToArray();
var onnxEnds = Enumerable.Repeat(int.MaxValue, rank).ToArray(); // by default copy the whole axis till the end
var onnxSteps = Enumerable.Repeat(1, rank).ToArray();
// NOTE: begin=0, end=0, stride=1 <= full range from existing axis
// begin=0, end=inf,stride=1 <= full range from existing axis
// begin=0, end=X, stride=1 <= full range from existing axis, if X==last element on this axis
// begin=0, end=0, stride=0 <= new axis OR shrink axis to single 1st element
// begin=N, end=N, stride=0 <= shrink axis to single Nth element
// These notes are copied from TensorExtensions.ApplyStridedSlice(...)
for (int i = 0; i < axes.Length; ++i)
{
var axis = axes[i];
if (axis < 0)
{
axis += rank;
}
axis = Math.Min(Math.Max(axis, 0), rank);
onnxStarts[axis] = starts[i];
onnxEnds[axis] = ends[i];
onnxSteps[axis] = steps[i];
}
layer.pad = PermuteToBarracuda(onnxStarts, rank, 0);
layer.pool = PermuteToBarracuda(onnxEnds, rank, int.MaxValue);
layer.stride = PermuteToBarracuda(onnxSteps, rank, 1);
return(true);
});
rewriters.Add(Layer.Type.Tile, (layer, net) =>
{
if (layer.inputs.Length == 1)
{
layer.pool = TensorExtensions.Permute(layer.pool, k_FromNCHWtoNHWC);
}
return(true);
});
rewriters.Add(Layer.Type.Activation, ConvertActivation);
rewriters.Add(Layer.Type.Gather, ConvertAxis);
rewriters.Add(Layer.Type.TopKIndices, ConvertAxis);
rewriters.Add(Layer.Type.TopKValues, ConvertAxis);
rewriters.Add(Layer.Type.LSTM, (layer, net) => ExpandOpsPass.ConvertLSTM(layer, net, m_Ops));
rewriters.Add(Layer.Type.RandomNormal, ConvertNormal);
rewriters.Add(Layer.Type.RandomUniform, ConvertNormal);
rewriters.Add(Layer.Type.ReduceMax, Reduce);
rewriters.Add(Layer.Type.ReduceMean, Reduce);
rewriters.Add(Layer.Type.ReduceMin, Reduce);
rewriters.Add(Layer.Type.ReduceProd, Reduce);
rewriters.Add(Layer.Type.ReduceSum, Reduce);
rewriters.Add(Layer.Type.ArgMax, Reduce);
rewriters.Add(Layer.Type.ArgMin, Reduce);
rewriters.Add(Layer.Type.Upsample2D, Upsample);
rewriters.Add(Layer.Type.Resample2D, Upsample);
rewriters.Add(Layer.Type.Upsample3D, Upsample);
rewriters.Add(Layer.Type.MatMul, (layer, net) =>
{
string input0 = layer.inputs[0];
if (!m_RanksByName.TryGetValue(input0, out int?input0Rank) || !input0Rank.HasValue)
{
throw new Exception($"Must have input rank for {input0} in order to convert axis for NHWC op");
}
string input1 = layer.inputs[1];
if (!m_RanksByName.TryGetValue(input1, out int?input1Rank) || !input1Rank.HasValue)
{
throw new Exception($"Must have input rank for {input1} in order to convert axis for NHWC op");
}
layer.pool = new[] { input0Rank.Value, input1Rank.Value };
return(true);
});
return(rewriters);
}
Dictionary <Layer.Type, Func <Layer, ModelBuilder, bool> > InstantiateRewriterNHWCToNHWC()
{
var rewritersNHWC = new Dictionary <Layer.Type, Func <Layer, ModelBuilder, bool> >();
// TODO, upsample is sometimes in NHWC mode
rewritersNHWC.Add(Layer.Type.Reshape, (layer, net) =>
{
if (layer.inputs.Length == 1)
{
var size = layer.pool;
// Don't use Tensorshape as this can remove a wild card
const int _ = 1;
if (size.Length == 1)
{
layer.pool = new[] { _, _, size[0], _, _, 1, 1, 1 }
}
; // [1,1,N,1,1,1,1,1]
else if (size.Length == 2)
{
layer.pool = new[] { _, _, size[0], _, _, 1, 1, size[1] }
}
; // [1, 1, N, 1, 1, 1, 1, C]
else if (size.Length == 3)
{
layer.pool = new[] { _, _, size[0], _, _, _, size[1], size[2] }
}
; // [1,1,N,1,1,1,W,C]
else if (size.Length == 4)
{
layer.pool = new[] { _, _, size[0], _, _, size[1], size[2], size[3] }
}
; // [1,1,N,1,1,H,W,C]
else if (size.Length == 5)
{
layer.pool = new[] { _, _, size[0], _, size[1], size[2], size[3], size[4] }
}
; // [1,1,N,1,D,H,W,C]
else if (size.Length == 6)
{
layer.pool = new[] { _, _, size[0], size[1], size[2], size[3], size[4], size[5] }
}
; // [1,1,N,T,D,H,W,C]
else
{
layer.pool = new[] { size[0], size[1], size[2], size[3], size[4], size[5], size[6], size[7] }
}; // [S,R,N,T,D,H,W,C]
}
return(true);
});
rewritersNHWC.Add(Layer.Type.Transpose, (layer, net) =>
{
var size = layer.pool;
if (size.Length == 1)
{
layer.pool = new[] { 0, 1, 2, 3 }; // [N,_,_,_]
layer.pool[0] = size[0];
}
else if (size.Length == 2)
{
layer.pool = new[] { 0, 1, 2, 3 }; // [N, _, _, C]
layer.pool[0] = size[0] == 0 ? 0 : size[0] + 2;
layer.pool[3] = size[1] == 0 ? 0 : size[1] + 2;
}
else if (size.Length == 3)
{
layer.pool = new[] { 0, 1, 2, 3 }; // [N, _, W, C]
layer.pool[0] = size[0] == 0 ? 0 : size[0] + 1;
layer.pool[2] = size[1] == 0 ? 0 : size[1] + 1;
layer.pool[3] = size[2] == 0 ? 0 : size[2] + 1;
}
else if (size.Length == 4)
{
layer.pool = size; // [N,H,W,C]
}
else if (size.Length == 5)
{
layer.pool = new[] { 0, 1, 2, 3, 4, 5, 6, 7 }; // [_,_,N,_,D,H,W,C]
layer.pool[2] = size[0] == 0 ? 2 : size[0] + 3;
layer.pool[4] = size[1] == 0 ? 2 : size[1] + 3;
layer.pool[5] = size[2] == 0 ? 2 : size[2] + 3;
layer.pool[6] = size[3] == 0 ? 2 : size[3] + 3;
layer.pool[7] = size[4] == 0 ? 2 : size[4] + 3;
}
else if (size.Length == 6)
{
layer.pool = new[] { 0, 1, 2, 3, 4, 5, 6, 7 }; // [1,1,N,T,D,H,W,C]
layer.pool[2] = size[0] + 2;
layer.pool[3] = size[1] + 2;
layer.pool[4] = size[2] + 2;
layer.pool[5] = size[3] + 2;
layer.pool[6] = size[4] + 2;
layer.pool[7] = size[5] + 2;
}
else
{
layer.pool = new[] { size[0], size[1], size[2], size[3], size[4], size[5], size[6], size[7] }
}; // [S,R,N,T,D,H,W,C]
return(true);
});
rewritersNHWC.Add(Layer.Type.Gather, ConvertGatherNHWC);
rewritersNHWC.Add(Layer.Type.Concat, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ReduceMax, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ReduceMean, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ReduceMin, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ReduceProd, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ReduceSum, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ArgMax, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.ArgMin, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.Activation, ConvertAxisNHWC);
rewritersNHWC.Add(Layer.Type.StridedSlice, (layer, net) =>
{
int rank = 4;
if (m_RanksByName.ContainsKey(layer.name) && m_RanksByName[layer.name] != null)
{
rank = m_RanksByName[layer.name].Value;
}
var name = layer.name;
var starts = layer.pad;
var ends = layer.pool;
var steps = layer.stride;
var axes = layer.axes;
var onnxStarts = Enumerable.Repeat(0, rank).ToArray();
var onnxEnds = Enumerable.Repeat(int.MaxValue, rank).ToArray(); // by default copy the whole axis till the end
var onnxSteps = Enumerable.Repeat(1, rank).ToArray();
// NOTE: begin=0, end=0, stride=1 <= full range from existing axis
// begin=0, end=inf,stride=1 <= full range from existing axis
// begin=0, end=X, stride=1 <= full range from existing axis, if X==last element on this axis
// begin=0, end=0, stride=0 <= new axis OR shrink axis to single 1st element
// begin=N, end=N, stride=0 <= shrink axis to single Nth element
// These notes are copied from TensorExtensions.ApplyStridedSlice(...)
for (int i = 0; i < axes.Length; ++i)
{
var axis = axes[i];
if (axis < 0)
{
axis += rank;
}
axis = Math.Min(Math.Max(axis, 0), rank);
onnxStarts[axis] = starts[i];
onnxEnds[axis] = ends[i];
onnxSteps[axis] = steps[i];
}
switch (rank)
{
case 1:
layer.pad = new[] { 0, 0, onnxStarts[0], 0, 0, 0, 0, 0 };
layer.pool = new[] { int.MaxValue, int.MaxValue, onnxEnds[0], int.MaxValue, int.MaxValue, int.MaxValue, int.MaxValue, int.MaxValue };
layer.stride = new[] { 1, 1, onnxSteps[0], 1, 1, 1, 1, 1 };
break;
case 2:
layer.pad = new[] { 0, 0, onnxStarts[0], 0, 0, 0, 0, onnxStarts[1] };
layer.pool = new[] { int.MaxValue, int.MaxValue, onnxEnds[0], int.MaxValue, int.MaxValue, int.MaxValue, int.MaxValue, onnxEnds[1] };
layer.stride = new[] { 1, 1, onnxSteps[0], 1, 1, 1, 1, onnxSteps[1] };
break;
case 3:
layer.pad = new[] { 0, 0, onnxStarts[0], 0, 0, 0, onnxStarts[1], onnxStarts[2] };
layer.pool = new[] { int.MaxValue, int.MaxValue, onnxEnds[0], int.MaxValue, int.MaxValue, int.MaxValue, onnxEnds[1], onnxEnds[2] };
layer.stride = new[] { 1, 1, onnxSteps[0], 1, 1, 1, onnxSteps[1], onnxSteps[2] };
break;
case 4:
layer.pad = new[] { 0, 0, onnxStarts[0], 0, 0, onnxStarts[1], onnxStarts[2], onnxStarts[3] };
layer.pool = new[] { int.MaxValue, int.MaxValue, onnxEnds[0], int.MaxValue, int.MaxValue, onnxEnds[1], onnxEnds[2], onnxEnds[3] };
layer.stride = new[] { 1, 1, onnxSteps[0], 1, 1, onnxSteps[1], onnxSteps[2], onnxSteps[3] };
break;
case 5:
layer.pad = new[] { 0, 0, onnxStarts[0], 0, onnxStarts[1], onnxStarts[2], onnxStarts[3], onnxStarts[4] };
layer.pool = new[] { int.MaxValue, int.MaxValue, onnxEnds[0], int.MaxValue, onnxEnds[1], onnxEnds[2], onnxEnds[3], onnxEnds[4] };
layer.stride = new[] { 1, 1, onnxSteps[0], 1, onnxSteps[1], onnxSteps[2], onnxSteps[3], onnxSteps[4] };
break;
default:
throw new ArgumentException($"Unsupported tensor rank {rank} for StridedSlice");
}
return(true);
});
rewritersNHWC.Add(Layer.Type.Flatten, (layer, net) =>
{
layer.type = Layer.Type.Nop;
return(true);
});
rewritersNHWC.Add(Layer.Type.Squeeze, (layer, net) =>
{
int input0Rank = 4;
if (m_RanksByName.ContainsKey(layer.inputs[0]) && m_RanksByName[layer.inputs[0]] != null)
{
input0Rank = m_RanksByName[layer.inputs[0]].Value;
}
int rank = input0Rank;
var combinePermutations = new[] { 0, 1, 2, 3 };
for (int i = 0; i < layer.pool.Length; i++)
{
int axis = layer.pool[i];
if (axis < 0)
{
axis = rank + 1 - axis;
}
var transpose = SqueezeAxisPermutationForMappingNHWCLayoutToBarracuda(rank, axis);
// there could be a 4 / 8D shape mismatch
if (transpose.Length == 8 && combinePermutations.Length == 4)
{
combinePermutations = Permutation4DTo8D(combinePermutations);
}
combinePermutations = TensorExtensions.Permute(transpose, combinePermutations);
rank--;
}
layer.type = Layer.Type.Transpose;
layer.pool = combinePermutations;
return(true);
});
rewritersNHWC.Add(Layer.Type.Unsqueeze, (layer, net) =>
{
int input0Rank = 4;
if (m_RanksByName.ContainsKey(layer.inputs[0]) && m_RanksByName[layer.inputs[0]] != null)
{
input0Rank = m_RanksByName[layer.inputs[0]].Value;
}
int rank = input0Rank;
var combinePermutations = new[] { 0, 1, 2, 3 };
for (int i = 0; i < layer.pool.Length; i++)
{
int axis = layer.pool[i];
if (axis < 0)
{
axis = rank + 1 - axis;
}
var transpose = UnSqueezeAxisPermutationForMappingNHWCLayoutToBarracuda(rank, axis);
// there could be a 4 / 8D shape mismatch
if (transpose.Length == 8 && combinePermutations.Length == 4)
{
combinePermutations = Permutation4DTo8D(combinePermutations);
}
combinePermutations = TensorExtensions.Permute(transpose, combinePermutations);
rank++;
}
layer.type = Layer.Type.Transpose;
layer.pool = combinePermutations;
return(true);
});
rewritersNHWC.Add(Layer.Type.Load, (layer, net) =>
{
int rank = layer.axis;
if (rank != 2 && rank != 3)
{
return(true);
}
var constX = layer.DataSetToTensor(0);
var shape = constX.shape;
switch (rank)
{
case 2:
// _,_,N,_,_,C,_,_ => _,_,N,_,_,_,_,C
shape = new TensorShape(shape.batch, shape.height);
break;
case 3:
// _,_,N,_,_,W,C,_ => _,_,N,_,_,_,W,C
shape = new TensorShape(shape.batch, shape.height, shape.width);
break;
}
var reshapedX = m_Ops.Reshape(constX, shape);
layer.ApplyTensorToDataSet(reshapedX, 0);
reshapedX.Dispose();
constX.Dispose();
return(true);
});
rewritersNHWC.Add(Layer.Type.MatMul, (layer, net) =>
{
string input0 = layer.inputs[0];
if (!m_RanksByName.TryGetValue(input0, out int?input0Rank) || !input0Rank.HasValue)
{
throw new Exception($"Must have input rank for {input0} in order to convert axis for NHWC op");
}
string input1 = layer.inputs[1];
if (!m_RanksByName.TryGetValue(input1, out int?input1Rank) || !input1Rank.HasValue)
{
throw new Exception($"Must have input rank for {input1} in order to convert axis for NHWC op");
}
layer.pool = new[] { input0Rank.Value, input1Rank.Value };
int outputRank = Math.Max(input0Rank.Value, input1Rank.Value);
if (outputRank <= 2)
{
return(true);
}
Layer input0Transposed = net.Transpose($"Transpose_For_{input0}", input0, input0Rank.Value == 3 ? k_FromNCHtoN1WC : k_ToNHWC);
Layer input1Transposed = net.Transpose($"Transpose_For_{input1}", input1, input1Rank.Value == 3 ? k_FromNCHtoN1WC : k_ToNHWC);
string originalLayerName = layer.name;
layer.name = $"{layer.name}_NHWC";
layer.inputs[0] = input0Transposed.name;
layer.inputs[1] = input1Transposed.name;
net.model.layers.Add(layer);
net.Transpose(originalLayerName, layer.name, outputRank == 3 ? k_FromN1WCtoNCH : k_ToNCHW);
return(false);
});
rewritersNHWC.Add(Layer.Type.Pad, PadNHWC);
return(rewritersNHWC);
}
