public static TensorShape?[] ListTemporaryTensorShapes(Model model, IDictionary <string, TensorShape> inputShapes,
out IDictionary <string, TensorShape?> shapesByName)
{
Profiler.BeginSample("Barracuda.ListTemporaryTensorShapes");
var shapes = new List <TensorShape?>();
shapesByName = new Dictionary <string, TensorShape?>();
foreach (var entry in inputShapes)
{
shapesByName.Add(entry.Key, entry.Value);
}
TensorShape?Xn;
shapesByName.TryGetValue(GetDefaultInputName(model), out Xn); // default input
TensorShape?O = Xn;
foreach (var l in model.layers)
{
if (l.inputs.Length > 0 && shapesByName.TryGetValue(l.inputs[0], out TensorShape? xShape))
{
Xn = xShape;
}
else
{
Xn = O; // previous output is used, if-and-only-if layer has no explicit inputs
}
if (Xn == null)
{
shapes.Add(Xn);
shapesByName.Add(l.name, Xn);
continue;
}
TensorShape X = Xn.Value;
if (l.type == Layer.Type.Dense)
{
Assert.IsNotNull(l.datasets);
var W = l.datasets[0].shape;
O = new TensorShape(X.flatHeight, W.flatWidth);
}
else if (l.type == Layer.Type.MatMul)
{
if (!shapesByName.ContainsKey(l.inputs[1]) || shapesByName[l.inputs[1]] == null)
{
O = null;
break;
}
var Y = shapesByName[l.inputs[1]].Value;
// ONNX rank 2 : N,C => N,1,1,C
// rank 3 : one must be N C W, (batches = N) => N, 1, W, C
// rank 4 : one must be N C H W, (batches = N * C) => N H W C
// X and Y can be different ranks
var onnxXshape = new List <int> {
X.batch, X.channels, X.height, X.width
};
if (X.height == 1)
{
onnxXshape = new List <int> {
X.batch, X.channels, X.width, 1
}
}
;
var onnxYshape = new List <int> {
Y.batch, Y.channels, Y.height, Y.width
};
if (Y.height == 1)
{
onnxYshape = new List <int> {
Y.batch, Y.channels, Y.width, 1
}
}
;
int rankX = 0;
for (int i = 3; i >= 0; i--)
{
if (onnxXshape[i] != 1)
{
rankX = i + 1;
break;
}
}
int rankY = 0;
for (int i = 3; i >= 0; i--)
{
if (onnxYshape[i] != 1)
{
rankY = i + 1;
break;
}
}
int rankO = Math.Max(rankX, rankY);
// pad 1 on front of shape to both be rankO shape
for (int i = 0; i < (rankX - rankY); i++)
{
onnxYshape.Insert(0, 1);
}
onnxYshape.RemoveRange(4, onnxYshape.Count - 4);
for (int i = 0; i < (rankY - rankX); i++)
{
onnxXshape.Insert(0, 1);
}
onnxXshape.RemoveRange(4, onnxXshape.Count - 4);
if (rankO == 2)
{
O = new TensorShape(onnxXshape[0], 1, 1, onnxYshape[1]);
}
else if (rankO == 3)
{
O = new TensorShape(Math.Max(onnxXshape[0], onnxYshape[0]), 1, onnxYshape[2], onnxXshape[1]);
}
else
{
O = new TensorShape(Math.Max(onnxXshape[0], onnxYshape[0]), onnxXshape[2], onnxYshape[3], Math.Max(onnxXshape[1], onnxYshape[1]));
}
}
else if (
l.type == Layer.Type.Conv2D ||
l.type == Layer.Type.DepthwiseConv2D)
{
var K = l.datasets[0].shape;
Assert.IsNotNull(l.stride);
Assert.IsNotNull(l.pad);
var pad = X.AdjustPadToKernel(K, l.stride, l.pad);
O = X.ApplyKernel(K, l.stride, pad);
}
else if (
l.type == Layer.Type.Conv2DTrans)
{
var K = l.datasets[0].shape;
Assert.IsNotNull(l.stride);
Assert.IsNotNull(l.pad);
// pool size is treated as output_adjustment aka output_padding here
var outputAdjustment = l.pool;
var pad = X.AdjustPadToKernel(K, l.stride, l.pad);
O = X.ApplyKernelInverse(K, l.stride, pad, outputAdjustment);
}
else if (
l.type == Layer.Type.Upsample2D)
{
if (inputShapes.Count > 1)
{
O = null;
}
else
{
// pool size is treated as upsample coefficient here
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 2);
O = new TensorShape(X.batch, X.height * l.pool[1], X.width * l.pool[0], X.channels);
}
}
else if (
l.type == Layer.Type.Resample2D)
{
if (inputShapes.Count > 1)
{
O = null;
}
else
{
// pool is treated as resample size here
var size = l.pool;
Assert.IsNotNull(size);
Assert.AreEqual(size.Length, 2);
O = new TensorShape(X.batch, size[1], size[0], X.channels);
}
}
else if (
l.type == Layer.Type.DepthToSpace)
{
// pool size is treated as blocksize here
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 2);
Assert.AreEqual(X.channels % (l.pool[0] * l.pool[1]), 0);
O = new TensorShape(X.batch, X.height * l.pool[1], X.width * l.pool[0], X.channels / (l.pool[0] * l.pool[1]));
}
else if (
l.type == Layer.Type.SpaceToDepth)
{
// pool size is treated as blocksize here
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 2);
O = new TensorShape(X.batch, X.height / l.pool[1], X.width / l.pool[0], X.channels * (l.pool[0] * l.pool[1]));
}
else if (
l.type == Layer.Type.MaxPool2D ||
l.type == Layer.Type.AvgPool2D)
{
Assert.IsNotNull(l.pool);
Assert.IsNotNull(l.stride);
Assert.IsNotNull(l.pad);
var pad = X.AdjustPadToPool(l.pool, l.stride, l.pad);
O = X.ApplyPool(l.pool, l.stride, pad);
}
else if (
l.type == Layer.Type.GlobalMaxPool2D ||
l.type == Layer.Type.GlobalAvgPool2D)
{
O = new TensorShape(X.batch, 1, 1, X.channels);
}
else if (
l.type == Layer.Type.Border2D ||
l.type == Layer.Type.Pad2DReflect ||
l.type == Layer.Type.Pad2DSymmetric ||
l.type == Layer.Type.Pad2DEdge)
{
Assert.IsNotNull(l.pad);
O = X.ApplyBorder(l.pad);
}
else if (
l.type == Layer.Type.Conv3D ||
l.type == Layer.Type.Conv3DTrans ||
l.type == Layer.Type.Upsample3D ||
l.type == Layer.Type.MaxPool3D ||
l.type == Layer.Type.AvgPool3D ||
l.type == Layer.Type.GlobalMaxPool3D ||
l.type == Layer.Type.GlobalAvgPool3D ||
l.type == Layer.Type.Border3D)
{
throw new NotImplementedException();
}
else if (
l.type == Layer.Type.RandomNormal ||
l.type == Layer.Type.RandomUniform)
{
Assert.IsNotNull(l.pool);
// pool size is treated as shape constant, if not empty
// otherwise shape of the previous tensor is used
if (l.pool.Length > 0)
{
O = new TensorShape(l.pool);
}
else
{
O = X;
}
}
else if (
l.type == Layer.Type.Multinomial)
{
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 1);
O = new TensorShape(X.batch, l.pool[0]);
}
else if (
l.type == Layer.Type.OneHot)
{
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 1);
int features = X.flatWidth;
int depth = l.pool[0];
if (X.flatWidth == 1) // 1D input
{
O = new TensorShape(X.batch, depth);
}
else
{
O = new TensorShape(X.batch, 1, depth, features);
}
}
else if (
l.type == Layer.Type.Add ||
l.type == Layer.Type.Sub ||
l.type == Layer.Type.Mul ||
l.type == Layer.Type.Div ||
l.type == Layer.Type.Pow ||
l.type == Layer.Type.Min ||
l.type == Layer.Type.Max ||
l.type == Layer.Type.Mean ||
l.type == Layer.Type.Greater ||
l.type == Layer.Type.GreaterEqual ||
l.type == Layer.Type.Less ||
l.type == Layer.Type.LessEqual ||
l.type == Layer.Type.Equal ||
l.type == Layer.Type.LogicalOr ||
l.type == Layer.Type.LogicalAnd ||
l.type == Layer.Type.LogicalXor)
{
// gather shapes by names
var list = new List <TensorShape>(l.inputs.Length);
bool allShapesKnown = true;
foreach (var i in l.inputs)
{
if (shapesByName.TryGetValue(i, out TensorShape? shape) && shape != null)
{
list.Add(shape.Value);
}
else
{
allShapesKnown = false;
}
}
O = allShapesKnown ? TensorExtensions.Max(list.ToArray()) : default(TensorShape?);
}
else if (
l.type == Layer.Type.ReduceL1 ||
l.type == Layer.Type.ReduceL2 ||
l.type == Layer.Type.ReduceLogSum ||
l.type == Layer.Type.ReduceLogSumExp ||
l.type == Layer.Type.ReduceMax ||
l.type == Layer.Type.ReduceMean ||
l.type == Layer.Type.ReduceMin ||
l.type == Layer.Type.ReduceProd ||
l.type == Layer.Type.ReduceSum ||
l.type == Layer.Type.ReduceSumSquare ||
l.type == Layer.Type.ArgMax ||
l.type == Layer.Type.ArgMin)
{
O = X.Reduce(l.axis);
}
else if (
l.type == Layer.Type.Flatten)
{
O = X.Flatten();
}
else if (
l.type == Layer.Type.Reshape)
{
// pool size is treated as the shape, if not empty
var size = l.pool;
Assert.IsNotNull(size);
if (size.Length == 0 && l.inputs.Length > 1)
{
switch (l.axis)
{
// Legacy - use the shape of the input tensor as the shape
case -1:
if (shapesByName.TryGetValue(l.inputs[1], out TensorShape? shape))
{
size = shape.Value.ToArray();
}
break;
// Use the tensor values as the shape; Calculated at runtime
case 1:
O = null;
break;
}
if (O == null)
{
break;
}
}
Assert.IsTrue((size.Length == 4) || (size.Length == 8));
O = X.Reshape(size);
}
else if (
l.type == Layer.Type.Expand)
{
// pool size is treated as new shape
var newShape = l.pool;
Assert.IsNotNull(newShape);
Assert.IsTrue(newShape.Length == 8 || newShape.Length == 4);
O = new TensorShape(newShape);
}
else if (
l.type == Layer.Type.Transpose)
{
var permutations = l.pool;
if (permutations == null)
{
O = new TensorShape(X.flatWidth, X.flatHeight);
}
else
{
Assert.IsTrue(permutations.Length == 8 || permutations.Length == 4);
O = X.Permute(permutations);
}
}
else if (
l.type == Layer.Type.Gather)
{
if (!shapesByName.TryGetValue(l.inputs[0], out TensorShape? input0Shape) || input0Shape == null ||
!shapesByName.TryGetValue(l.inputs[1], out TensorShape? input1Shape) || input1Shape == null)
{
O = null;
break;
}
int[] shape = input0Shape.Value.ToArray();
shape[l.axis] = input1Shape.Value.length;
O = new TensorShape(shape);
}
else if (
l.type == Layer.Type.Squeeze ||
l.type == Layer.Type.Unsqueeze)
{
throw new NotImplementedException();
}
else if (
l.type == Layer.Type.Concat)
{
// gather shapes by names
var list = new List <TensorShape>(l.inputs.Length);
bool allShapesKnown = true;
foreach (var i in l.inputs)
{
if (!shapesByName.TryGetValue(i, out var shape) || shape == null)
{
allShapesKnown = false;
continue;
}
list.Add(shape.Value);
}
O = allShapesKnown ? TensorExtensions.Concat(list.ToArray(), l.axis) : default(TensorShape?);
}
else if (
l.type == Layer.Type.StridedSlice)
{
Assert.IsNotNull(l.pad);
Assert.IsNotNull(l.pool);
Assert.IsNotNull(l.stride);
O = X.ApplyStridedSlice(l.pad, l.pool, l.stride);
}
else if (
l.type == Layer.Type.Tile)
{
// pool size is treated as tiling coefficient here
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 4);
var scale = l.pool;
O = X.Scale(scale);
}
else if (
l.type == Layer.Type.Load)
{
O = l.datasets[0].shape;
}
else if (// elementwise operations
l.type == Layer.Type.Nop ||
l.type == Layer.Type.Activation ||
l.type == Layer.Type.ScaleBias ||
l.type == Layer.Type.Normalization ||
l.type == Layer.Type.LRN ||
l.type == Layer.Type.Dropout ||
l.type == Layer.Type.LogicalNot ||
l.activation == Layer.Activation.PRelu)
{
// works in place, keeps the same shape size
O = X;
}
else if (
l.type == Layer.Type.TopKIndices ||
l.type == Layer.Type.TopKValues ||
l.type == Layer.Type.NonMaxSuppression ||
l.type == Layer.Type.NonZero)
{
// Calculated at runtime
O = null;
}
else if (l.type == Layer.Type.Shape)
{
int shapeRank = l.axis > 0 ? 1 : X.length;
O = new TensorShape(shapeRank, 1, 1, 1);
}
else if (l.type == Layer.Type.Where)
{
O = X;
}
else if (
l.type == Layer.Type.Conv3D ||
l.type == Layer.Type.Conv3DTrans ||
l.type == Layer.Type.Upsample3D ||
l.type == Layer.Type.MaxPool3D ||
l.type == Layer.Type.AvgPool3D ||
l.type == Layer.Type.GlobalMaxPool3D ||
l.type == Layer.Type.GlobalAvgPool3D ||
l.type == Layer.Type.Border3D)
{
throw new NotImplementedException("3D operations are not implemented yet!");
}
else
{
throw new NotImplementedException($"Layer type {l.type} needs to be explicitly handled");
}
shapes.Add(O);
shapesByName.Add(l.name, O);
}
Profiler.EndSample();
return(shapes.ToArray());
}
static public TensorShape?[] ListTemporaryTensorShapes(Model model, IDictionary <string, TensorShape> inputShapes,
out IDictionary <string, TensorShape?> shapesByName)
{
Profiler.BeginSample("Barracuda.ListTemporaryTensorShapes");
var shapes = new List <TensorShape?>();
shapesByName = new Dictionary <string, TensorShape?>();
foreach (var entry in inputShapes)
{
shapesByName.Add(entry.Key, entry.Value);
}
TensorShape?Xn;
shapesByName.TryGetValue(GetDefaultInputName(model), out Xn); // default input
TensorShape?O = Xn;
foreach (var l in model.layers)
{
if (l.inputs.Length > 0 && shapesByName.ContainsKey(l.inputs[0]))
{
Xn = shapesByName[l.inputs[0]];
}
else
{
Xn = O; // previous output is used, if-and-only-if layer has no explicit inputs
}
if (Xn == null)
{
shapes.Add(Xn);
shapesByName.Add(l.name, Xn);
continue;
}
TensorShape X = Xn.Value;
if (l.type == Layer.Type.Dense)
{
Assert.IsNotNull(l.datasets);
var W = l.datasets[0].shape;
O = new TensorShape(X.flatHeight, W.flatWidth);
}
else if (
l.type == Layer.Type.Conv2D ||
l.type == Layer.Type.DepthwiseConv2D)
{
var K = l.datasets[0].shape;
Assert.IsNotNull(l.stride);
Assert.IsNotNull(l.pad);
var pad = X.AdjustPadToKernel(K, l.stride, l.pad);
O = X.ApplyKernel(K, l.stride, pad);
}
else if (
l.type == Layer.Type.Conv2DTrans)
{
var K = l.datasets[0].shape;
Assert.IsNotNull(l.stride);
Assert.IsNotNull(l.pad);
// pool size is treated as output_adjustment aka output_padding here
var outputAdjustment = l.pool;
var pad = X.AdjustPadToKernel(K, l.stride, l.pad);
O = X.ApplyKernelInverse(K, l.stride, pad, outputAdjustment);
}
else if (
l.type == Layer.Type.Upsample2D)
{
if (inputShapes.Count > 1)
{
O = null;
}
else
{
// pool size is treated as upsample coefficient here
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 2);
O = new TensorShape(X.batch, X.height * l.pool[1], X.width * l.pool[0], X.channels);
}
}
else if (
l.type == Layer.Type.MaxPool2D ||
l.type == Layer.Type.AvgPool2D)
{
Assert.IsNotNull(l.pool);
Assert.IsNotNull(l.stride);
Assert.IsNotNull(l.pad);
var pad = X.AdjustPadToPool(l.pool, l.stride, l.pad);
O = X.ApplyPool(l.pool, l.stride, pad);
}
else if (
l.type == Layer.Type.GlobalMaxPool2D ||
l.type == Layer.Type.GlobalAvgPool2D)
{
O = new TensorShape(X.batch, 1, 1, X.channels);
}
else if (
l.type == Layer.Type.Border2D ||
l.type == Layer.Type.Pad2DReflect ||
l.type == Layer.Type.Pad2DSymmetric ||
l.type == Layer.Type.Pad2DEdge)
{
Assert.IsNotNull(l.pad);
O = X.ApplyBorder(l.pad);
}
else if (
l.type == Layer.Type.Conv3D ||
l.type == Layer.Type.Conv3DTrans ||
l.type == Layer.Type.Upsample3D ||
l.type == Layer.Type.MaxPool3D ||
l.type == Layer.Type.AvgPool3D ||
l.type == Layer.Type.GlobalMaxPool3D ||
l.type == Layer.Type.GlobalAvgPool3D ||
l.type == Layer.Type.Border3D)
{
throw new NotImplementedException();
}
else if (
l.type == Layer.Type.RandomNormal ||
l.type == Layer.Type.RandomUniform)
{
Assert.IsNotNull(l.pool);
// pool size is treated as shape constant, if not empty
// otherwise shape of the previous tensor is used
if (l.pool.Length > 0)
{
O = new TensorShape(l.pool);
}
else
{
O = X;
}
}
else if (
l.type == Layer.Type.Multinomial)
{
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 1);
O = new TensorShape(X.batch, l.pool[0]);
}
else if (
l.type == Layer.Type.OneHot)
{
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 1);
int features = X.flatWidth;
int depth = l.pool[0];
O = new TensorShape(X.batch, 1, features, depth);
}
else if (
l.type == Layer.Type.Add ||
l.type == Layer.Type.Sub ||
l.type == Layer.Type.Mul ||
l.type == Layer.Type.Div ||
l.type == Layer.Type.Pow ||
l.type == Layer.Type.Min ||
l.type == Layer.Type.Max ||
l.type == Layer.Type.Mean ||
l.type == Layer.Type.Greater ||
l.type == Layer.Type.GreaterEqual ||
l.type == Layer.Type.Less ||
l.type == Layer.Type.LessEqual ||
l.type == Layer.Type.Equal ||
l.type == Layer.Type.LogicalOr ||
l.type == Layer.Type.LogicalAnd ||
l.type == Layer.Type.LogicalXor)
{
// gather shapes by names
var list = new List <TensorShape>(l.inputs.Length);
bool allShapesKnown = true;
foreach (var i in l.inputs)
{
if (!shapesByName.ContainsKey(i))
{
continue;
}
TensorShape?shape = shapesByName[i];
if (shape == null)
{
allShapesKnown = false;
continue;
}
list.Add(shapesByName[i].Value);
}
O = allShapesKnown ? TensorExtensions.Max(list.ToArray()) : default(TensorShape?);
}
else if (
l.type == Layer.Type.ReduceL1 ||
l.type == Layer.Type.ReduceL2 ||
l.type == Layer.Type.ReduceLogSum ||
l.type == Layer.Type.ReduceLogSumExp ||
l.type == Layer.Type.ReduceMax ||
l.type == Layer.Type.ReduceMean ||
l.type == Layer.Type.ReduceMin ||
l.type == Layer.Type.ReduceProd ||
l.type == Layer.Type.ReduceSum ||
l.type == Layer.Type.ReduceSumSquare)
{
O = X.Reduce(l.axis);
}
else if (
l.type == Layer.Type.Flatten)
{
O = X.Flatten();
}
else if (
l.type == Layer.Type.Reshape)
{
// pool size is treated as reshape coefficient, if not empty
// otherwise shape of the 2nd input tensor is used
var size = l.pool;
Assert.IsNotNull(size);
if (size.Length == 0 && l.inputs.Length > 1)
{
if (shapesByName[l.inputs[1]] == null)
{
O = null;
break;
}
size = shapesByName[l.inputs[1]].Value.ToArray();
}
Assert.AreEqual(size.Length, 4);
O = X.Reshape(size);
}
else if (
l.type == Layer.Type.Transpose)
{
O = new TensorShape(X.flatWidth, X.flatHeight);
}
else if (
l.type == Layer.Type.Gather)
{
if (shapesByName[l.inputs[0]] == null || shapesByName[l.inputs[1]] == null)
{
O = null;
break;
}
int[] shape = shapesByName[l.inputs[0]].Value.ToArray();
shape[l.axis] = shapesByName[l.inputs[1]].Value.flatWidth;
O = new TensorShape(shape);
}
else if (
l.type == Layer.Type.Squeeze ||
l.type == Layer.Type.Unsqueeze)
{
throw new NotImplementedException();
}
else if (
l.type == Layer.Type.Concat)
{
// gather shapes by names
var list = new List <TensorShape>(l.inputs.Length);
bool allShapesKnown = true;
foreach (var i in l.inputs)
{
if (!shapesByName.ContainsKey(i))
{
continue;
}
if (shapesByName[i] == null)
{
allShapesKnown = false;
continue;
}
list.Add(shapesByName[i].Value);
}
O = allShapesKnown ? TensorExtensions.Concat(list.ToArray(), l.axis) : default(TensorShape?);
}
else if (
l.type == Layer.Type.StridedSlice)
{
Assert.IsNotNull(l.pad);
Assert.IsNotNull(l.pool);
Assert.IsNotNull(l.stride);
O = X.ApplyStridedSlice(l.pad, l.pool, l.stride);
}
else if (
l.type == Layer.Type.Tile)
{
// pool size is treated as tiling coefficient here
Assert.IsNotNull(l.pool);
Assert.AreEqual(l.pool.Length, 4);
var scale = l.pool;
O = X.Scale(scale);
}
else if (
l.type == Layer.Type.Load)
{
O = l.datasets[0].shape;
}
else if (// elementwise operations
l.type == Layer.Type.Nop ||
l.type == Layer.Type.Activation ||
l.type == Layer.Type.ScaleBias ||
l.type == Layer.Type.Normalization ||
l.type == Layer.Type.LRN ||
l.type == Layer.Type.Dropout ||
l.type == Layer.Type.LogicalNot ||
l.activation == Layer.Activation.PRelu)
{
// works in place, keeps the same shape size
O = X;
}
else if (
l.type == Layer.Type.Conv3D ||
l.type == Layer.Type.Conv3DTrans ||
l.type == Layer.Type.Upsample3D ||
l.type == Layer.Type.MaxPool3D ||
l.type == Layer.Type.AvgPool3D ||
l.type == Layer.Type.GlobalMaxPool3D ||
l.type == Layer.Type.GlobalAvgPool3D ||
l.type == Layer.Type.Border3D)
{
throw new NotImplementedException("3D operations are not implemented yet!");
}
else
{
Assert.AreEqual(l.activation, Layer.Activation.None);
O = X;
}
shapes.Add(O);
shapesByName.Add(l.name, O);
}
Profiler.EndSample();
return(shapes.ToArray());
}
