Files
unity-application/Packages/com.unity.barracuda/Runtime/Core/TensorExtensions.cs
2023-03-18 19:53:17 +00:00

1196 lines
42 KiB
C#

using System;
using System.Collections.Generic;
using System.Linq;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Assertions;
namespace Unity.Barracuda {
/// <summary>
/// Tensor extension methods
/// </summary>
public static class TensorExtensions
{
static internal void TestInit(this Tensor X, int n = -1, int modulus = -1)
{
if (n < 0)
n = X.length;
n = Math.Min(n, X.length);
for (int i = 0; i < n; ++i)
{
if (modulus > 1)
X[i] = i % modulus;
else
X[i] = i;
}
}
static internal void TestInitCos(this Tensor X, int n = -1, float offset = 0.0f)
{
if (n < 0)
n = X.length;
n = Math.Min(n, X.length);
for (int i = 0; i < n; ++i)
X[i] = Mathf.Cos(i + offset);
}
static internal void TestInitRandom(this Tensor X, int n = -1)
{
if (n < 0)
n = X.length;
n = Math.Min(n, X.length);
for (int i = 0; i < n; ++i)
X[i] = UnityEngine.Random.value;
}
static internal void TestInitValue(this Tensor X, float value=0.1f, int n = -1)
{
if (n < 0)
n = X.length;
n = Math.Min(n, X.length);
for (int i = 0; i < n; ++i)
X[i] = value;
}
/// <summary>
/// Return Tensor data as float array, this will create a blocking read operation
/// </summary>
/// <param name="x">Tensor</param>
/// <returns>Tensor data as float array</returns>
static public float[] AsFloats(this Tensor x)
{
return x.ToReadOnlyArray();
}
/// <summary>
/// Return Tensor data as int array (slow operation), this will create a blocking read operation
/// </summary>
/// <param name="x">Tensor</param>
/// <returns>Tensor data as int array</returns>
static public int[] AsInts(this Tensor x)
{
return Array.ConvertAll(x.ToReadOnlyArray(), v => v <= (float)int.MinValue ? int.MinValue : v >= (float)int.MaxValue ? int.MaxValue : (int)v);
}
/// <summary>
/// Return Tensor data as string, limits number of elements to `size`
/// </summary>
/// <param name="X">Tensor</param>
/// <param name="size">element number limit</param>
/// <returns>Returns Tensor data as string</returns>
static public string DataToString(this Tensor X, int size = 32)
{
var str = "";
for (int i = 0; i < X.length && i < size; ++i)
{
str += X[i];
str += " ";
}
if (X.length > size)
str += "...";
return str;
}
/// <summary>
/// Print Tensor metadata to console
/// </summary>
/// <param name="X">Tensor</param>
/// <param name="msg">message prefix</param>
static public void Print(this Tensor X, string msg = "")
{
if (msg.Length > 0)
msg += " ";
D.Log($"{msg}{X.name} {X.shape}");
}
/// <summary>
/// Print Tensor data to console
/// </summary>
/// <param name="X">Tensor</param>
/// <param name="size">element number limit</param>
/// <param name="msg">message prefix</param>
static public void PrintDataPart(this Tensor X, int size, string msg = "")
{
if (msg.Length > 0)
msg += " ";
D.Log($"{msg}{X.DataToString(size)}");
}
/// <summary>
/// Compare Tensor contents
/// </summary>
/// <param name="X">left Tensor</param>
/// <param name="Y">right Tensor</param>
/// <returns>`true` if shape and data content matches</returns>
static public bool Equals(this Tensor X, Tensor Y)
{
if (X.shape != Y.shape)
return false;
if (X.length != Y.length)
return false;
for (int i = 0; i < X.length; ++i)
{
if (X[i] != Y[i])
return false;
}
return true;
}
/// <summary>
/// Compare Tensor contents approximately
/// </summary>
/// <param name="X">left Tensor</param>
/// <param name="Y">right Tensor</param>
/// <param name="epsilon">comparison threshold</param>
/// <param name="count">limit number of elements to compare</param>
/// <returns>`true` if shape match and while data content matches approximately</returns>
static public bool Approximately(this Tensor X, Tensor Y, float epsilon = 1e-4f, int count = -1)
{
if (X.shape != Y.shape)
return false;
if (X.length != Y.length)
return false;
if (count < 0)
count = X.length;
for (int i = 0; i < count; ++i)
{
// If one of the values is NaN, the comparison against epislon will return false.
// But if tensor has NaN and the other doesn't, they shouldn't be considered "close".
if (Mathf.Abs(X[i] - Y[i]) > epsilon || float.IsNaN(X[i]) != float.IsNaN(Y[i]))
{
// @TODO: move logging into dedicated function
D.Log("First mismatch @ [" + i + "]: " + X[i] + " != " + Y[i]);
return false;
}
}
return true;
}
/// <summary>
/// Calculate max difference between two tensors
/// </summary>
/// <param name="X">first Tensor</param>
/// <param name="Y">second Tensor</param>
/// <returns></returns>
static public float MaxDifference(this Tensor X, Tensor Y)
{
float maxD = 0f;
for (int i = 0; i < X.length; ++i)
maxD = Mathf.Max(Mathf.Abs(X[i] - Y[i]), maxD);
return maxD;
}
/// <summary>
/// Reshape Tensor
/// </summary>
/// <param name="X">Tensor</param>
/// <param name="size">new shape as array of int (expected as size 4 for NHWC or size 8 for SRNTDHWC)</param>
/// <returns>reshaped Tensor</returns>
static public Tensor Reshape(this Tensor X, int[] size)
{
var newShape = X.shape.Reshape(size);
return X.Reshape(newShape);
}
/// <summary>
/// Calculate max value index
/// </summary>
/// <param name="X">Tensor</param>
/// <returns>max value index</returns>
static public int[] ArgMax(this Tensor X)
{
Assert.AreEqual(TensorShape.DataChannel, TensorShape.MaxRank - 1); // expects channels last layout
Assert.IsTrue(X.channels != 0);
Assert.AreEqual(X.length % X.channels, 0);
// reduce over the last dimension - channels
var innerLength = X.channels;
var outterLength = X.length / innerLength;
int[] result = new int[outterLength];
for (var n = 0; n < outterLength; ++n)
{
float maxV = Mathf.NegativeInfinity;
for (int c = 0; c < innerLength; ++c)
{
var v = X[n * innerLength + c];
if (maxV >= v)
continue;
maxV = v;
result[n] = c;
}
}
return result;
}
/// <summary>
/// Return indices in order that would produce sorted Tensor values
/// </summary>
/// <param name="X">Tensor</param>
/// <returns>indices in order that would produce sorted Tensor values</returns>
static public int[][] ArgSort(this Tensor X)
{
Assert.AreEqual(TensorShape.DataChannel, TensorShape.MaxRank - 1); // expects channels last layout
Assert.IsTrue(X.channels != 0);
Assert.AreEqual(X.length % X.channels, 0);
// reduce over the last dimension - channels
var innerLength = X.channels;
var outterLength = X.length / innerLength;
var result = new List<int[]>();
for (var n = 0; n < outterLength; ++n)
{
int[] indices = Enumerable.Range(0, innerLength).ToArray<int>();
var sliceOffset = n * innerLength;
Array.Sort<int>(indices, (a, b) => X[sliceOffset + a].CompareTo(X[sliceOffset + b]));
result.Add(indices);
}
return result.ToArray();
}
/// <summary>
/// Fill Tensor with `value`
/// </summary>
/// <param name="X">Tensor</param>
/// <param name="value">value</param>
public static void Fill(this Tensor X, float value)
{
for (int i = 0; i < X.length; ++i)
X[i] = value;
}
/// <summary>
/// Calculate output shape for Gather operation
/// </summary>
/// <param name="shapes">input shapes</param>
/// <param name="axis">axis</param>
/// <returns>output shape</returns>
static public TensorShape Gather(TensorShape[] shapes, int axis)
{
TensorShape shape = shapes[0];
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
TensorShape indices = shapes[1];
if (!indices.hasNamedDimensions)
indices = indices.AsNamed();
shape[axis] = indices.length;
return shape;
}
/// <summary>
/// Concatenate `Tensor` array along `axis` and calculate output shape
/// </summary>
/// <param name="tensors">Tensor array</param>
/// <param name="axis">axis</param>
/// <returns>new `TensorShape`</returns>
/// <exception cref="ArgumentException">Off-axis dimension mismatch</exception>
static public TensorShape Concat(Tensor[] tensors, int axis)
{
if (tensors.Length == 0)
return new TensorShape();
var a = tensors[0].shape;
if (!a.hasNamedDimensions)
a = a.AsNamed();
var aAxis = a.Axis(axis);
// validate that off axis dimensions are equal
for (var i = 1; i < tensors.Length; ++i)
{
var b = tensors[i].shape;
if (!b.hasNamedDimensions)
b = b.AsNamed();
var bAxis = b.Axis(axis);
a[aAxis] = 0; b[bAxis] = 0;
if (a != b)
{
foreach (var s in tensors)
D.Log(s.shape);
throw new ArgumentException("Off-axis dimensions must match");
}
}
var shape = tensors[0].shape;
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
var dstAxis = tensors[0].shape.Axis(axis);
for (var i = 1; i < tensors.Length; ++i)
{
var otherShape = tensors[i].shape;
if (!otherShape.hasNamedDimensions)
otherShape = otherShape.AsNamed();
shape[dstAxis] += otherShape[axis];
}
return shape;
}
/// <summary>
/// Calculate concatenation output shape
/// </summary>
/// <param name="shapes">input shapes</param>
/// <param name="axis">concatenation axis</param>
/// <returns>output shape</returns>
/// <exception cref="ArgumentException">Off-axis dimension mismatch</exception>
static public TensorShape Concat(TensorShape[] shapes, int axis)
{
if (shapes.Length == 0)
return new TensorShape();
var a = shapes[0];
if (!a.hasNamedDimensions)
a = a.AsNamed();
var aAxis = a.Axis(axis);
// validate that off axis dimensions are equal
for (var i = 1; i < shapes.Length; ++i)
{
var b = shapes[i];
if (!b.hasNamedDimensions)
b = b.AsNamed();
var bAxis = b.Axis(axis);
a[aAxis] = 0; b[bAxis] = 0;
if (a != b)
{
foreach (var s in shapes)
D.Log(s);
throw new ArgumentException("Off-axis dimensions must match");
}
}
var shape = shapes[0];
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
var dstAxis = shape.Axis(axis);
for (var i = 1; i < shapes.Length; ++i)
{
var otherShape = shapes[i];
if (!otherShape.hasNamedDimensions)
otherShape = otherShape.AsNamed();
shape[dstAxis] += otherShape[axis];
}
return shape;
}
/// <summary>
/// Calculate maximum shape that would cover all input shapes
/// </summary>
/// <param name="shapes">input shapes</param>
/// <returns>output shape</returns>
static public TensorShape Max(TensorShape[] shapes)
{
Assert.IsTrue(shapes.Length > 0);
var shape = shapes[0];
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
for (var i = 1; i < shapes.Length; ++i)
{
var otherShape = shapes[i];
if (!otherShape.hasNamedDimensions)
otherShape = otherShape.AsNamed();
for (var axis = 0; axis < TensorShape.MaxRank; axis++)
{
shape[axis] = Math.Max(shape[axis], otherShape[axis]);
}
}
return shape;
}
/// <summary>
/// Calculate maximum shape that would cover all input tensors
/// </summary>
/// <param name="tensors">input tensors</param>
/// <returns>output shape</returns>
static public TensorShape MaxShape(Tensor[] tensors)
{
Assert.IsTrue(tensors.Length > 0);
var shape = tensors[0].shape;
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
for (var i = 1; i < tensors.Length; ++i)
{
for (var axis = 0; axis < TensorShape.MaxRank; axis++)
{
var otherShape = tensors[i].shape;
if (!otherShape.hasNamedDimensions)
otherShape = otherShape.AsNamed();
shape[axis] = Math.Max(shape[axis], otherShape[axis]);
}
}
return shape;
}
/// <summary>
/// Scale TensorShape by the `scale` factor
/// </summary>
/// <param name="shape">TensorShape</param>
/// <param name="scale">scale</param>
/// <returns>output shape</returns>
static public TensorShape Scale(this TensorShape shape, TensorShape scale)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
var newShape = shape;
for (var axis = 0; axis < TensorShape.MaxRank; axis++)
newShape[axis] *= scale[axis];
return newShape;
}
/// <summary>
/// Scale TensorShape by the `scale` factor
/// </summary>
/// <param name="shape">TensorShape</param>
/// <param name="scale">scale</param>
/// <returns>output shape</returns>
static public TensorShape Scale(this TensorShape shape, int[] scale)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (scale.Length == TensorShape.MaxRank)
{
for (var axis = 0; axis < TensorShape.MaxRank; axis++)
shape[axis] *= scale[axis];
}
else
{
Assert.AreEqual(4, scale.Length);
shape[TensorShape.DataBatch] *= scale[0];
shape[5] *= scale[1];
shape[6] *= scale[2];
shape[7] *= scale[3];
}
return shape;
}
/// <summary>
/// Reduce TensorShape across specified `axis`
/// </summary>
/// <param name="shape">TensorShape</param>
/// <param name="axis">axis</param>
/// <returns>output shape</returns>
static public TensorShape Reduce(this TensorShape shape, int axis)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
axis = shape.Axis(axis);
var newShapeArray = shape;
newShapeArray[axis] = 1;
return newShapeArray;
}
/// <summary>
/// Reshape TensorShape into new shape specified by `size`. At most one dimension of the new shape can be -1.
/// See: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Reshape
/// </summary>
/// <param name="shape">TensorShape</param>
/// <param name="size4Dor8D">new shape</param>
/// <returns>output shape</returns>
/// <exception cref="ArgumentException">more than one dimension is unspecified</exception>
static public TensorShape Reshape(this TensorShape shape, int[] size4Dor8D)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
unsafe
{
int* size = stackalloc int[TensorShape.MaxRank];
int* newShapeArray = stackalloc int[TensorShape.MaxRank];
Get8DParametersNoAlloc(shape, size4Dor8D, size, 1);
for (int d = 0; d < TensorShape.MaxRank; ++d)
newShapeArray[d] = shape[d];
// From: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Reshape
//
// At most one dimension of the new shape can be -1.
// In this case, the value is inferred from the size of the tensor and the remaining dimensions.
//
// A dimension could also be 0,
// in which case the actual dimension value is unchanged (i.e. taken from the input tensor).
var multipleOf = 1;
var unknownIndex = -1;
for (int q = 0; q < TensorShape.MaxRank; ++q)
{
if (size[q] > 0)
{
multipleOf *= size[q];
newShapeArray[q] = size[q];
}
else if (size[q] == 0)
multipleOf *= newShapeArray[q];
else if (unknownIndex == -1)
unknownIndex = q;
else
throw new ArgumentException("Can only specify one unknown dimension");
}
if (unknownIndex == -1)
{
// all dimensions are given
var newShape = new TensorShape(newShapeArray[0], newShapeArray[1], newShapeArray[2], newShapeArray[3],
newShapeArray[4], newShapeArray[5], newShapeArray[6], newShapeArray[7]);
if (shape.length != newShape.length)
throw new ArgumentException("Cannot reshape array of size " + shape.length +
" into shape " + newShape);
return newShape;
}
var solveForIndex = shape.length / multipleOf;
bool remainderLeft = shape.length % multipleOf != 0;
if (remainderLeft)
throw new ArgumentException("Cannot reshape array of size " + shape.length +
" into shape with multiple of " + multipleOf + " elements");
newShapeArray[unknownIndex] = solveForIndex;
return new TensorShape(newShapeArray[0], newShapeArray[1], newShapeArray[2], newShapeArray[3],
newShapeArray[4], newShapeArray[5], newShapeArray[6], newShapeArray[7]);
}
}
/// <summary>
/// Calculate new shape after applying border to current TensorShape
/// </summary>
/// <param name="shape">TensorShape</param>
/// <param name="border">border</param>
/// <returns>new TensorShape</returns>
static public TensorShape ApplyBorder(this TensorShape shape, int[] border)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
Assert.IsTrue(border.Length == 6 || border.Length == 8);
if(border.Length == 6)
{
shape[TensorShape.H] += border[1] + border[4];
shape[TensorShape.W] += border[0] + border[3];
shape[TensorShape.C] += border[2] + border[5];
}
else if (border.Length == 8)
{
shape[TensorShape.D] += border[2] + border[6];
shape[TensorShape.H] += border[1] + border[5];
shape[TensorShape.W] += border[0] + border[4];
shape[TensorShape.C] += border[3] + border[7];
}
return shape;
}
static internal int[] AdjustPadToKernel(this Tensor tensor, Tensor kernel, int[] stride, int[] pad)
{
return AdjustPadToKernel(tensor.shape, kernel.shape, stride, pad);
}
static internal int[] AdjustPadToKernel(this TensorShape shape, TensorShape kernel, int[] stride, int[] pad)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
Assert.IsTrue(stride.Length==2 || stride.Length==3);
unsafe
{
int* kernelDims = stackalloc int[stride.Length == 2 ? 2 : 3];
kernelDims[0] = kernel.kernelWidth;
kernelDims[1] = kernel.kernelHeight;
if (stride.Length > 2)
kernelDims[2] = kernel.kernelSpatialDepth;
return AdjustPadToPool(shape, kernelDims, stride, pad);
}
}
static internal int[] AdjustPadToPool(this Tensor tensor, int[] pool, int[] stride, int[] pad)
{
return AdjustPadToPool(tensor.shape, pool, stride, pad);
}
static internal unsafe int[] AdjustPadToPool(this Tensor tensor, int* pool, int[] stride, int[] pad)
{
return AdjustPadToPool(tensor.shape, pool, stride, pad);
}
static internal int[] AdjustPadToPool(this TensorShape shape, int[] pool, int[] stride, int[] pad)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
unsafe
{
fixed (int* pPool = pool)
{
return AdjustPadToPool(shape, pPool, stride, pad);
}
}
}
static internal unsafe int[] AdjustPadToPool(this TensorShape shape, int* pool, int[] stride, int[] pad)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
Assert.IsTrue(stride.Length > 0);
int featureCount = stride.Length;
Assert.IsTrue(featureCount <= TensorShape.DataFeatures.Length);
// negative pad values mean auto_pad type is used
if (pad[0] >= 0)
return pad;
var type = (Layer.AutoPad)pad[0];
if (type == Layer.AutoPad.SameUpper || type == Layer.AutoPad.SameLower)
{
// Based on ONNX (AveragePool & MaxPool)
// https://github.com/onnx/onnx/blob/master/docs/Operators.md
// and TensorFlow docs:
// https://www.tensorflow.org/api_guides/python/nn#Notes_on_SAME_Convolution_Padding
var adjustedPad = new int [featureCount*2];
for (var i = 0; i < featureCount; ++i)
{
var featureModStride = shape.width % stride[i];
if (featureModStride == 0)
featureModStride = stride[i];
var padAlongFeature = Math.Max(pool[i] - featureModStride, 0);
// Code above (based on TensorFlow docs) is equivalent to (based on ONNX docs):
// padAlongWidth = (Mathf.Ceil(shape.width/stride[0]) - 1) * stride[0] + pool[0] - shape.width;
// padAlongHeight = (Mathf.Ceil(shape.height/stride[1]) - 1) * stride[1] + pool[1] - shape.height;
var featureSmall = padAlongFeature / 2;
var featureLarge = padAlongFeature - featureSmall;
if (type == Layer.AutoPad.SameUpper) {
adjustedPad[i] = featureSmall;
adjustedPad[i+featureCount] = featureLarge;
} else {
adjustedPad[i] = featureLarge;
adjustedPad[i+featureCount] = featureSmall;
}
}
return adjustedPad;
}
else
throw new NotImplementedException("This padding type is not implemented yet!");
}
static internal TensorShape ApplyPool(this TensorShape shape, int[] pool, int[] stride, int[] pad,
bool ceilMode = false)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
Assert.IsTrue(stride.Length == pool.Length);
unsafe
{
fixed (int* pPool = pool)
{
return ApplyPool(shape, pPool, stride, pad, ceilMode);
}
}
}
static internal unsafe TensorShape ApplyPool(this TensorShape shape, int* pool, int[] stride, int[] pad, bool ceilMode = false)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
Assert.IsTrue(stride.Length > 0);
Assert.IsTrue(stride.Length*2 == pad.Length);
int featureCount = stride.Length;
Assert.IsTrue(featureCount <= TensorShape.DataFeatures.Length);
// Based on ONNX (AveragePool & MaxPool)
// https://github.com/onnx/onnx/blob/master/docs/Operators.md
// Theano "Convolution arithmetic tutorial"
// http://deeplearning.net/software/theano/tutorial/conv_arithmetic.html#quick-reference
// and TensorFlow docs:
// https://www.tensorflow.org/api_guides/python/nn#Convolution
// https://www.tensorflow.org/api_guides/python/nn#Notes_on_SAME_Convolution_Padding
//
// output_size = (input_size + pad_left + pad_right - kernel_size) / stride + 1
var newShape = shape;
for (var i = 0; i < featureCount; ++i)
{
// C# automatically rounds down
// https://docs.microsoft.com/en-us/dotnet/csharp/language-reference/operators/arithmetic-operators
if (ceilMode)
newShape[TensorShape.DataFeatures[i]] = (shape[TensorShape.DataFeatures[i]] + (pad[i]+pad[i+featureCount]) - pool[i] + stride[i] - 1) / stride[i] + 1;
else
newShape[TensorShape.DataFeatures[i]] = (shape[TensorShape.DataFeatures[i]] + (pad[i]+pad[i+featureCount]) - pool[i]) / stride[i] + 1;
}
return newShape;
}
static internal TensorShape ApplyKernel(this TensorShape shape, TensorShape kernel, int[] stride, int[] pad)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
unsafe
{
Assert.IsTrue(stride.Length==2 || stride.Length==3);
int* kernelDims = stackalloc int[stride.Length == 2 ? 2 : 3];
kernelDims[0] = kernel.kernelWidth;
kernelDims[1] = kernel.kernelHeight;
if (stride.Length > 2)
kernelDims[2] = kernel.kernelSpatialDepth;
var outShape = ApplyPool(shape, kernelDims, stride, pad);
outShape[7] = kernel.kernelCount;
return outShape;
}
}
static internal TensorShape ApplyKernelInverse(this TensorShape shape, TensorShape kernel, int[] stride, int[] pad, int[] outputAdjustment)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
Assert.IsTrue(stride.Length > 0);
Assert.IsTrue(stride.Length * 2 == pad.Length);
Assert.IsTrue(stride.Length <= TensorShape.KernelSpatials.Length);
Assert.IsTrue(stride.Length <= TensorShape.DataFeatures.Length);
// Based on ONNX (ConvTranspose)
// https://github.com/onnx/onnx/blob/master/docs/Operators.md
// and Theano "Convolution arithmetic tutorial"
// http://deeplearning.net/software/theano/tutorial/conv_arithmetic.html#transposed-convolution-arithmetic
//
// Inverse of:
// output_size = (input_size + pad_left + pad_right - kernel_size) / stride + 1
// Resulting in:
// output_size = (input_size - 1 ) * stride - (pad_left + pad_right) + kernel_size + output_adj
// output_adj = (input_size + (pad_left + pad_right) - kernel_size) % stride
//
if (outputAdjustment == null || outputAdjustment.Length == 0)
{
outputAdjustment = new int[stride.Length];
for (var i = 0; i < stride.Length; ++i)
{
var featureAxis = TensorShape.DataFeatures[i];
var kernelAxis = TensorShape.KernelSpatials[i];
var padding = pad[i] + pad[stride.Length+i];
outputAdjustment[i] = (shape[featureAxis] + padding - kernel[kernelAxis]) % stride[i];
}
}
var newShape = shape;
for (var i = 0; i < stride.Length; ++i)
{
var featureAxis = TensorShape.DataFeatures[i];
var kernelAxis = TensorShape.KernelSpatials[i];
var padding = pad[i] + pad[stride.Length+i];
newShape[featureAxis] = (shape[featureAxis] - 1) * stride[i] - padding + kernel[kernelAxis] + outputAdjustment[i];
}
newShape[TensorShape.KernelOutChannel] = kernel.kernelCount;
return newShape;
}
/// <summary>
/// Wrap index (emulate Python array index behavior)
/// </summary>
/// <param name="i">index</param>
/// <param name="length">array length</param>
/// <returns>wrapped around index</returns>
static public int WrapIndex(int i, int length)
{
// allow index to be equal to length
// in order to enable iteration over [i,end) range
if (i >= length)
return length;
// in C# modulo of negative is negative
// to emulate Python array behavior, we use: https://stackoverflow.com/questions/1082917/mod-of-negative-number-is-melting-my-brain/1082938
var v = i % length;
return v < 0 ? (v + length): v;
}
static internal bool IsNDHWC(this TensorShape shape)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
return shape.sequenceLength == 1 &&
shape.numberOfDirections == 1 &&
shape.extraDimension == 1;
}
static internal bool Is4D(this TensorShape shape)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
return shape.sequenceLength == 1 &&
shape.numberOfDirections == 1 &&
shape.extraDimension == 1 &&
shape.depth == 1;
}
// Works for NCHW or NHWC
static internal int Convert4DTo8DAxis(int axis)
{
Assert.IsTrue(axis < 4);
Assert.IsTrue(axis > -4);
if (axis < 0) //backward indexing
{
return axis;
}
else if (axis == 0) //batch
return TensorShape.DataBatch;
else //H,W,C
return axis + TensorShape.D;
}
static internal int FirstNotIdentityFeatureDimensionIndex(this TensorShape shape)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
for (int dimIndex = TensorShape.DataFeature3; dimIndex < TensorShape.MaxRank; ++dimIndex)
{
if (shape[dimIndex] > 1)
return dimIndex;
}
return TensorShape.MaxRank;
}
static internal bool Is8DAxisConvertibleTo4D(int axis)
{
Assert.IsTrue(axis > -4);
Assert.IsTrue(axis < TensorShape.MaxRank);
return axis < 0 || axis == TensorShape.DataBatch || axis > TensorShape.D;
}
/// <summary>
/// Check if all tensors are convertible to 4D tensors
/// </summary>
/// <param name="tensors">tensors</param>
/// <returns>`true` if all tensors are 4D (or less)</returns>
static public bool AreAllTensorsConvertibleTo4D(Tensor[] tensors)
{
for (int i = 0; i < tensors.Length; ++i)
{
if (!tensors[i].shape.Is4D())
return false;
}
return true;
}
static internal int Convert8DAxisTo4D(int axis)
{
Assert.IsTrue(Is8DAxisConvertibleTo4D(axis));
if (axis < 0) //backward indexing
{
return axis;
}
else if (axis == TensorShape.DataBatch) //batch
return 0;
else //H,W,C
return axis - TensorShape.D;
}
static internal unsafe void Get8DParametersNoAlloc(this TensorShape shape, int[] parameters, int* parameters8D, int defaultValue)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (parameters.Length == TensorShape.MaxRank)
{
for (int i = 0; i < TensorShape.MaxRank; ++i)
parameters8D[i] = parameters[i];
}
else
{
Assert.AreEqual(4, parameters.Length);
if (!shape.Is4D()) Assert.IsTrue(false, $"4D Parameters {parameters} can't be used with a tensor of shape {shape} as it contains other dimensions, please use 8D parameters for this shape.");
parameters8D[0] = defaultValue;
parameters8D[1] = defaultValue;
parameters8D[2] = parameters[0];
parameters8D[3] = defaultValue;
parameters8D[4] = defaultValue;
parameters8D[5] = parameters[1];
parameters8D[6] = parameters[2];
parameters8D[7] = parameters[3];
}
}
/// <summary>
/// Calculate 8D permutations from 4D
/// </summary>
/// <param name="shape">shape</param>
/// <param name="permutations">permutations</param>
/// <returns>8D permutations</returns>
static public int[] Get8DPermutationsForNHWCPermutationsAndShape(this TensorShape shape, int[] permutations)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (permutations.Length == TensorShape.MaxRank)
return permutations;
Assert.AreEqual(4, permutations.Length);
if (!shape.Is4D()) Assert.IsTrue(false, $"4D Permutation {permutations} can't be used with a tensor of shape {shape} as it contains other dimensions, please use an 8D permutation for this shape.");
int batchOldAxis = Convert4DTo8DAxis(permutations[0]);
int heighOldAxis = Convert4DTo8DAxis(permutations[1]);
int widthOldIndex = Convert4DTo8DAxis(permutations[2]);
int channeOldIndex = Convert4DTo8DAxis(permutations[3]);
return new int[] {0, 1, batchOldAxis, 3, 4, heighOldAxis, widthOldIndex, channeOldIndex };
}
static internal NativeArray<int> Get8DPermutationsForNHWCPermutationsAndShape(this TensorShape shape, NativeArray<int> inPermutations)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (inPermutations.Length == TensorShape.MaxRank)
return inPermutations;
Assert.AreEqual(4, inPermutations.Length);
if (!shape.Is4D()) Assert.IsTrue(false, $"4D Permutation {inPermutations.ToString()} can't be used with a tensor of shape {shape} as it contains other dimensions, please use an 8D permutation for this shape.");
int batchOldAxis = Convert4DTo8DAxis(inPermutations[0]);
int heighOldAxis = Convert4DTo8DAxis(inPermutations[1]);
int widthOldIndex = Convert4DTo8DAxis(inPermutations[2]);
int channeOldIndex = Convert4DTo8DAxis(inPermutations[3]);
// Valid only for single frame
NativeArray<int> outPermutations = new NativeArray<int>(8, Allocator.Temp);
outPermutations[0] = 0;
outPermutations[1] = 1;
outPermutations[2] = batchOldAxis;
outPermutations[3] = 3;
outPermutations[4] = 4;
outPermutations[5] = heighOldAxis;
outPermutations[6] = widthOldIndex;
outPermutations[7] = channeOldIndex;
return outPermutations;
}
static internal int[] Get8DPermutationsForNCHWPermutationsAndShape(this TensorShape shape, int[] permutations)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (permutations.Length == TensorShape.MaxRank)
return permutations;
Assert.AreEqual(4, permutations.Length);
if (!shape.Is4D()) Assert.IsTrue(false, $"4D Permutation {permutations} can't be used with a tensor of shape {shape} as it contains other dimensions, please use an 8D permutation for this shape.");
int batchOldAxis = Convert4DTo8DAxis(permutations[0]);
int channelOldIndex = Convert4DTo8DAxis(permutations[1]);
int heightOldIndex = Convert4DTo8DAxis(permutations[2]);
int widthOldIndex = Convert4DTo8DAxis(permutations[3]);
return new int[] {0, 1, batchOldAxis, 3, 4, channelOldIndex, heightOldIndex, widthOldIndex };
}
static internal NativeArray<int> Get8DPermutationsForNCHWPermutationsAndShape(this TensorShape shape, NativeArray<int> inPermutations)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (inPermutations.Length == TensorShape.MaxRank)
return inPermutations;
Assert.AreEqual(4, inPermutations.Length);
if (!shape.Is4D()) Assert.IsTrue(false, $"4D Permutation {inPermutations.ToString()} can't be used with a tensor of shape {shape} as it contains other dimensions, please use an 8D permutation for this shape.");
int batchOldAxis = Convert4DTo8DAxis(inPermutations[0]);
int channelOldIndex = Convert4DTo8DAxis(inPermutations[1]);
int heightOldIndex = Convert4DTo8DAxis(inPermutations[2]);
int widthOldIndex = Convert4DTo8DAxis(inPermutations[3]);
// Valid only for single frame
NativeArray<int> outPermutations = new NativeArray<int>(8, Allocator.Temp);
outPermutations[0] = 0;
outPermutations[1] = 1;
outPermutations[2] = batchOldAxis;
outPermutations[3] = 3;
outPermutations[4] = 4;
outPermutations[5] = channelOldIndex;
outPermutations[6] = heightOldIndex;
outPermutations[7] = widthOldIndex;
return outPermutations;
}
static internal unsafe TensorShape ApplyStridedSlice8DUnsafeNoAlloc(this TensorShape shape, int* starts, int* ends,
int* stride)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
TensorShape sliced = shape;
for (int i = 0; i < shape.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
// take + 1 is si > shape[i]
int ei = TensorExtensions.WrapIndex(ends[i], shape[i]);
int si = TensorExtensions.WrapIndex(starts[i], shape[i]);
// Barracuda convetion (non ONNX), t[0:0] => t[:]
if (si == 0 && ei == 0)
ei = shape[i];
if (stride[i] > 0)
sliced[i] = (int)Math.Round((double)(Math.Min(ei, shape[i]) - Math.Min(si, shape[i] - 1)) / (double)(Mathf.Abs(stride[i])), MidpointRounding.AwayFromZero);
else if (stride[i] < 0)
{
bool inclusive = ends[i] < -shape[i]; // edge case when ends is negative and bigger than nchwShape
sliced[i] = (int)Math.Round((double)(Math.Min(si, shape[i] - 1) - Math.Min(ei, shape[i]) + (inclusive ? 1 : 0)) / (double)(Mathf.Abs(stride[i])), MidpointRounding.AwayFromZero);
}
else
{
// Assert.IsTrue(stride[i] != 0); // 0 strides not allowed
// breaks legacy implementations
D.LogWarning("StridedSlice with 0 strides, not supported! Slicing to 1D dimension");
sliced[i] = 1;
}
}
return sliced;
}
static internal TensorShape ApplyStridedSlice(this TensorShape shape, int[] starts, int[] ends, int[] stride)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
unsafe
{
int* starts8Dbuffer = stackalloc int[TensorShape.MaxRank];
int* ends8Dbuffer = stackalloc int[TensorShape.MaxRank];
int* stride8Dbuffer = stackalloc int[TensorShape.MaxRank];
Get8DParametersNoAlloc(shape, starts, starts8Dbuffer, 0);
Get8DParametersNoAlloc(shape, ends, ends8Dbuffer, 1);
Get8DParametersNoAlloc(shape, stride, stride8Dbuffer, 1);
return shape.ApplyStridedSlice8DUnsafeNoAlloc(starts8Dbuffer, ends8Dbuffer, stride8Dbuffer);
}
}
/// <summary>
/// Calculate shape after applying permutations
/// </summary>
/// <param name="shape">shape</param>
/// <param name="permutations">permutations</param>
/// <returns>new shape</returns>
static public int[] Permute(int[] shape, int[] permutations)
{
Assert.AreEqual(shape.Length, permutations.Length);
var output = new int[shape.Length];
for (var i = 0; i < permutations.Length; ++i)
output[i] = permutations[i] >= 0 ? shape[permutations[i]] : 1;
return output;
}
/// <summary>
/// Calculate TensorShape after applying permutations
/// </summary>
/// <param name="shape">shape</param>
/// <param name="permutations">permutations</param>
/// <returns>new TensorShape</returns>
static public TensorShape Permute(this TensorShape shape, int[] permutations)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (permutations.Length == 4)
permutations = Get8DPermutationsForNHWCPermutationsAndShape(shape, permutations);
var permutedShape = new int[TensorShape.MaxRank];
for (var i = 0; i < permutations.Length; ++i)
permutedShape[i] = permutations[i] >= 0 ? shape[permutations[i]] : 1;
var output = new TensorShape(permutedShape);
return output;
}
static internal TensorShape Permute(this TensorShape shape, NativeArray<int> permutations)
{
if (!shape.hasNamedDimensions)
shape = shape.AsNamed();
if (permutations.Length == 4)
permutations = Get8DPermutationsForNHWCPermutationsAndShape(shape, permutations);
var permutedShape = new int[TensorShape.MaxRank];
for (var i = 0; i < permutations.Length; ++i)
permutedShape[i] = permutations[i] >= 0 ? shape[permutations[i]] : 1;
var output = new TensorShape(permutedShape);
return output;
}
/// <summary>
/// Create ITensorData from Texture
/// </summary>
/// <param name="tex">Texture</param>
/// <param name="shape">shape</param>
/// <returns>created ITensorData</returns>
/// <exception cref="NotImplementedException">thrown if unsupported texture type is supplied</exception>
static public ITensorData CreateFromTexture(Texture tex, TensorShape shape)
{
Assert.AreEqual(tex.width, shape.width);
Assert.AreEqual(tex.height, shape.height);
Assert.IsTrue(shape.channels < 4);
// @TODO: implement proper GPU storage
var data = new ArrayTensorData(shape);
if (tex is Texture2D)
{
Texture2D tex2d = tex as Texture2D;
var pixels = tex2d.GetPixels();
for (int i = 0; i < data.array.Length && i < pixels.Length * shape.channels; ++i)
data.array[i] = pixels[i / shape.channels][i % shape.channels];
}
else
throw new NotImplementedException();
return data;
}
}
} // namespace Unity.Barracuda