/// <summary> /// Static constructor, initializes and caches data /// </summary> static ComputeInfo() { string[] args = System.Environment.GetCommandLineArgs(); for (int i = 0; i < args.Length; i++) { if (args [i] == "-barracuda-compute-use-nchw") { channelsOrder = ChannelsOrder.NCHW; } } supportsCompute = SystemInfo.supportsComputeShaders; graphicsDeviceVendor = SystemInfo.graphicsDeviceVendor; // TODO switch to SystemInfo.maxComputeWorkGroupSize when we bump min spec to 2019.3 if (Application.platform == RuntimePlatform.Android) { maxComputeWorkGroupSize = (SystemInfo.graphicsDeviceType == GraphicsDeviceType.Vulkan) ? 256u : 128u; var gpuName = SystemInfo.graphicsDeviceName ?? ""; var osName = SystemInfo.operatingSystem ?? ""; // Known issue with Adreno Vulkan drivers on Android 8.x if (gpuName.Contains("Adreno") && osName.StartsWith("Android OS 8") && SystemInfo.graphicsDeviceType == GraphicsDeviceType.Vulkan) { maxComputeWorkGroupSize = 128u; } } else if (Application.platform == RuntimePlatform.IPhonePlayer || Application.platform == RuntimePlatform.tvOS) { var gpuName = SystemInfo.graphicsDeviceName; if (gpuName != null && gpuName.StartsWith("Apple A")) { int gpuNumber = 0, idx = "Apple A".Length; while (idx < gpuName.Length && '0' <= gpuName[idx] && gpuName[idx] <= '9') { gpuNumber = gpuNumber * 10 + gpuName[idx++] - '0'; } // TODO check on lower end iOS devices maxComputeWorkGroupSize = (gpuNumber <= 10) ? 224u : 256u; } else { maxComputeWorkGroupSize = 256u; } } }
// works on IRModel public bool InferAllLayersChannelOrder(Model model, out Dictionary <string, ChannelsOrder> layerChannelOrder) { layerChannelOrder = new Dictionary <string, ChannelsOrder>(); IDictionary <string, TensorShape?> shapesByName = new Dictionary <string, TensorShape?>(); IDictionary <string, int?> ranksByName = new Dictionary <string, int?>(); foreach (var i in model.inputs) { ranksByName[i.name] = i.rank; if (!ModelAnalyzer.IsInputShapeAcceptablyKnowForShapeInference(i)) { continue; } shapesByName[i.name] = new TensorShape(i.shape); } IRShapeInferenceAndConstantFusing shapeInferencePass = new IRShapeInferenceAndConstantFusing(); shapeInferencePass.InferAllShapes(model, ref shapesByName, ref ranksByName); // flood-fill approach: NCHW layout is propagated from NCHW ops // * onnx-nchw ops are flagged as being native nchw // * nchw layout is propagated to upstream and downstream nodes // foreach node: // take layout being propagated to // if T or T-1 flip layout depending on upstream/downstream direction // - stop if layout is the same as previously propagated // - native nchw layout has priority Queue <(string, ChannelsOrder, FlowDirection)> layersToInferLayout = new Queue <(string, ChannelsOrder, FlowDirection)>(); for (int l = 0; l < model.layers.Count; l++) { var layer = model.layers[l]; if (!IsLayerNecessarilyNCHWOnnx(layer)) { continue; } layersToInferLayout.Enqueue((layer.name, ChannelsOrder.NativeNCHW, FlowDirection.Seed)); } while (layersToInferLayout.Any()) { (string, ChannelsOrder, FlowDirection)layerData = layersToInferLayout.Dequeue(); string name = layerData.Item1; ChannelsOrder deducedChannelOrder = layerData.Item2; // 0: in-place native // 1: downstream // 2: upstream FlowDirection flowDirection = layerData.Item3; if (!layerChannelOrder.ContainsKey(name)) { layerChannelOrder[name] = deducedChannelOrder; } else if (deducedChannelOrder == layerChannelOrder[name]) { continue; } else if (layerChannelOrder[name] == ChannelsOrder.NativeNCHW) { continue; } // heuristic to stop ping-pong loop, prioritize NHWC over NCHW as it implies less transposes // if incoming is NativeNCHW always propagate that // TODO: count # of transpose swaps else if (layerChannelOrder[name] == ChannelsOrder.NHWC && deducedChannelOrder != ChannelsOrder.NativeNCHW) { continue; } Layer layer; bool found = ModelAnalyzer.FindLayerByName(model, name, out layer); if (IsLayerChangingLayoutToNHWC(layer, shapesByName, ranksByName)) { // NCHW -> T -> NHWC if (((deducedChannelOrder == ChannelsOrder.NCHW) || (deducedChannelOrder == ChannelsOrder.NativeNCHW)) && (flowDirection == FlowDirection.Downstream)) { deducedChannelOrder = ChannelsOrder.TransposeToNHWC; } // NCHW <- T <- NHWC else if ((deducedChannelOrder == ChannelsOrder.NHWC) && (flowDirection == FlowDirection.Upstream)) { deducedChannelOrder = ChannelsOrder.TransposeToNHWC; } } else if (IsLayerChangingLayoutToNCHW(layer, shapesByName, ranksByName)) { // NHWC -> T-1 -> NCHW if ((deducedChannelOrder == ChannelsOrder.NHWC) && (flowDirection == FlowDirection.Downstream)) { deducedChannelOrder = ChannelsOrder.TransposeToNCHW; } // NHWC <- T-1 <- NCHW else if (((deducedChannelOrder == ChannelsOrder.NCHW) || (deducedChannelOrder == ChannelsOrder.NativeNCHW)) && (flowDirection == FlowDirection.Upstream)) { deducedChannelOrder = ChannelsOrder.TransposeToNCHW; } } if ((deducedChannelOrder == ChannelsOrder.TransposeToNCHW || deducedChannelOrder == ChannelsOrder.TransposeToNHWC) && (deducedChannelOrder == layerChannelOrder[name])) { continue; } layerChannelOrder[name] = deducedChannelOrder; foreach (var input in layer.inputs) { if (deducedChannelOrder == ChannelsOrder.TransposeToNCHW) { layersToInferLayout.Enqueue((input, ChannelsOrder.NHWC, FlowDirection.Upstream)); } else if (deducedChannelOrder == ChannelsOrder.TransposeToNHWC) { layersToInferLayout.Enqueue((input, ChannelsOrder.NCHW, FlowDirection.Upstream)); } else { layersToInferLayout.Enqueue((input, deducedChannelOrder, FlowDirection.Upstream)); } } var outputs = ModelAnalyzer.FindLayerOutputs(model, layer.name); foreach (var output in outputs) { if (deducedChannelOrder == ChannelsOrder.TransposeToNCHW) { layersToInferLayout.Enqueue((output, ChannelsOrder.NCHW, FlowDirection.Downstream)); } else if (deducedChannelOrder == ChannelsOrder.TransposeToNHWC) { layersToInferLayout.Enqueue((output, ChannelsOrder.NHWC, FlowDirection.Downstream)); } else { layersToInferLayout.Enqueue((output, deducedChannelOrder, FlowDirection.Downstream)); } } } bool modelExportedASNHWC = false; foreach (string key in layerChannelOrder.Keys.ToList()) { var value = layerChannelOrder[key]; if (value == ChannelsOrder.NativeNCHW) { layerChannelOrder[key] = ChannelsOrder.NCHW; } if (value == ChannelsOrder.NHWC) { modelExportedASNHWC = true; } } return(modelExportedASNHWC); }
// works on IRModel public bool InferAllLayersChannelOrder(Model model, out Dictionary <string, ChannelsOrder> layerChannelOrder) { // TF2Onnx : pattern T (.* Conv .*) T-1 // * being transpose commutative layer layerChannelOrder = new Dictionary <string, ChannelsOrder>(); IDictionary <string, TensorShape?> shapesByName = new Dictionary <string, TensorShape?>(); IDictionary <string, int?> ranksByName = new Dictionary <string, int?>(); foreach (var i in model.inputs) { ranksByName[i.name] = i.rank; if (!ModelAnalyzer.IsInputShapeAcceptablyKnowForShapeInference(i)) { continue; } shapesByName[i.name] = new TensorShape(i.shape); } IRShapeInferenceAndConstantFusing shapeInferencePass = new IRShapeInferenceAndConstantFusing(); shapeInferencePass.InferAllShapes(model, ref shapesByName, ref ranksByName); bool inputsNHWC = false; bool inputsNHWCExportedInputsAsNCHW = false; bool patternMatchStart = false; bool patternMatchConv = false; // tf to onnx does not swizzle axis, need to match * Conv * T-1 ... bool patternMatchStartInputsAsNCHWConv = false; for (int l = 0; l < model.layers.Count; l++) { var layer = model.layers[l]; if (!patternMatchStart && IsLayerTranpose(layer) && Enumerable.SequenceEqual(layer.pool, new[] { 0, 3, 1, 2 }) || IsLayerReshape(layer) && (shapesByName[layer.inputs[0]] != null) && IsReshapeTransposeToNCHW(layer, shapesByName[layer.inputs[0]].Value)) { patternMatchStart = true; } else if (patternMatchStart && patternMatchConv && ((IsLayerTranpose(layer) && Enumerable.SequenceEqual(layer.pool, new[] { 0, 2, 3, 1 })) || (IsLayerReshape(layer) && (shapesByName[layer.inputs[0]] != null) && IsReshapeTransposeToNHWC(layer, shapesByName[layer.inputs[0]].Value)) || (IsLayerSqueeze(layer) && (ranksByName[layer.inputs[0]] != null) && IsSqueezeTransposeToNHWC(layer, ranksByName[layer.inputs[0]].Value)) || (IsLayerFlatten(layer) && (ranksByName[layer.inputs[0]] != null) && IsFlattenTransposeToNHWC(layer, ranksByName[layer.inputs[0]].Value)))) { inputsNHWC = true; } else if (patternMatchStart && IsLayerConv(layer)) { patternMatchConv = true; } if (!inputsNHWCExportedInputsAsNCHW && patternMatchStartInputsAsNCHWConv && ((IsLayerTranpose(layer) && Enumerable.SequenceEqual(layer.pool, new[] { 0, 2, 3, 1 })) || (IsLayerReshape(layer) && (shapesByName[layer.inputs[0]] != null) && IsReshapeTransposeToNHWC(layer, shapesByName[layer.inputs[0]].Value)))) { inputsNHWCExportedInputsAsNCHW = true; } else if (!patternMatchStartInputsAsNCHWConv && !patternMatchStart && IsLayerConv(layer)) { patternMatchStartInputsAsNCHWConv = true; } } // flag each layer as being NHWC or NCHW for (int i = 0; i < model.inputs.Count; i++) { Model.Input input = model.inputs[i]; if (!inputsNHWCExportedInputsAsNCHW) { layerChannelOrder[input.name] = inputsNHWC ? ChannelsOrder.NHWC : ChannelsOrder.NCHW; } else { layerChannelOrder[input.name] = ChannelsOrder.NCHW; } } for (int l = 0; l < model.layers.Count; l++) { var layer = model.layers[l]; if (IsLayerTranpose(layer) && Enumerable.SequenceEqual(layer.pool, new[] { 0, 3, 1, 2 }) || IsLayerReshape(layer) && (shapesByName[layer.inputs[0]] != null) && IsReshapeTransposeToNCHW(layer, shapesByName[layer.inputs[0]].Value) && layerChannelOrder[layer.inputs[0]] == ChannelsOrder.NHWC) { layerChannelOrder[layer.name] = ChannelsOrder.TransposeToNCHW; } else if (IsLayerTranpose(layer) && Enumerable.SequenceEqual(layer.pool, new[] { 0, 2, 3, 1 }) || IsLayerReshape(layer) && (shapesByName[layer.inputs[0]] != null) && IsReshapeTransposeToNHWC(layer, shapesByName[layer.inputs[0]].Value) && layerChannelOrder[layer.inputs[0]] == ChannelsOrder.NCHW) { layerChannelOrder[layer.name] = ChannelsOrder.TransposeToNHWC; } else { string inputWithKnownOrder = null; for (int i = 0; i < layer.inputs.Length; i++) { var input = layer.inputs[i]; if (layerChannelOrder.ContainsKey(input)) { inputWithKnownOrder = input; break; } } if (inputWithKnownOrder == null) { continue; } Assert.IsNotNull(inputWithKnownOrder); ChannelsOrder inputOrder = layerChannelOrder[inputWithKnownOrder]; if (inputOrder == ChannelsOrder.TransposeToNCHW) { inputOrder = ChannelsOrder.NCHW; } else if (inputOrder == ChannelsOrder.TransposeToNHWC) { inputOrder = ChannelsOrder.NHWC; } // all layers with unknown layout are const for (int i = 0; i < layer.inputs.Length; i++) { var input = layer.inputs[i]; if (!layerChannelOrder.ContainsKey(input)) { layerChannelOrder[input] = inputOrder; } } layerChannelOrder[layer.name] = inputOrder; } } // TODO Assert that all layers have a channel order // Assert that all layers are NHWC if inputsNHWC return(inputsNHWC || inputsNHWCExportedInputsAsNCHW); }