void CorrectConstantsForBroadCast(ref Model nhwc) { List <Layer> correctedConstants = new List <Layer>(); for (int l = 0; l < nhwc.layers.Count; l++) { Layer layer = nhwc.layers[l]; for (int i = 0; i < layer.inputs.Length; i++) { var input = layer.inputs[i]; if (!ModelAnalyzer.IsLayerBroacastable(layer)) { continue; } if (!m_RanksByName.ContainsKey(input) || !m_RanksByName.ContainsKey(layer.name)) { continue; } Layer inputLayer; bool found = ModelAnalyzer.FindLayerByName(nhwc, input, out inputLayer); if (!found) { continue; } if (!ModelOptimizer.IsLayerConstant(inputLayer)) { continue; } if (m_RanksByName[input] < 1 || m_RanksByName[input] == m_RanksByName[layer.name]) { continue; } if (inputLayer.weights.Length == 1) { continue; } if (m_RanksByName[input] > m_RanksByName[layer.name]) { throw new Exception($"constant must be lower rank than input for broadcast to work, TODO add transpose before input"); } Layer correctedConstLayer = new Layer("c_" + inputLayer.name + "For_" + layer.name, Layer.Type.Load); // transpose dataset correctedConstLayer.datasets = new Layer.DataSet[1]; Array.Copy(inputLayer.datasets, correctedConstLayer.datasets, inputLayer.datasets.Length); correctedConstLayer.datasets[0].name = correctedConstLayer.name; correctedConstLayer.weights = new float[inputLayer.weights.Length]; var X = inputLayer.DataSetToTensor(0); int[] permutations = new[] { 0, 1, 2, 3 }; var rank = m_RanksByName[layer.name].Value; switch (rank) { case 2: // ONNX: 5,7 + 7 // Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7 permutations = new[] { 1, 2, 3, 0 }; break; case 3: // ONNX: 5,7,3 + 3 // Barracuda: 5,_,3,7 + 3,_,_,_ => _,_,3,_ if (m_RanksByName[input] == 1) { permutations = new[] { 1, 2, 0, 3 } } ; // ONNX: 5,7,3 + 7,3 // Barracuda: 5,_,3,7 + 7,_,_,3 => _,_,3,7 else if (m_RanksByName[input] == 2) { permutations = new[] { 1, 2, 3, 0 } } ; break; case 4: // ONNX: 2,5,7,3 + 3 // Barracuda: 2,7,3,5 + 3,_,_,_ => _,_,3,_ if (m_RanksByName[input] == 1) { permutations = new[] { 1, 2, 0, 3 } } ; // ONNX: 2,5,7,3 + 7,3 // Barracuda: 2,7,3,5 + 7,_,_,3 => _,7,3,_ else if (m_RanksByName[input] == 2) { permutations = new[] { 2, 0, 1, 3 } } ; // ONNX: 2,5,7,3 + 5,7,3 // Barracuda: 2,7,3,5 + 5,_,3,7 => _,7,3,5 else if (m_RanksByName[input] == 3) { permutations = new[] { 1, 3, 2, 0 } } ; break; } if (m_isModelExportedFromNCHW && (m_layersChannelOrder[layer.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC)) { switch (rank) { case 2: // ONNX: 5,7 + 7 // Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7 permutations = new[] { 1, 2, 3, 0 }; break; case 3: // ONNX: 5,7,3 + 3 // Barracuda: 5,_,7,3 + 3,_,_,_ => _,_,_,3 if (m_RanksByName[input] == 1) { permutations = new[] { 1, 2, 3, 0 } } ; // ONNX: 5,7,3 + 7,3 // Barracuda: 5,_,7,3 + 7,_,_,3 => _,_,7,3 else if (m_RanksByName[input] == 2) { permutations = new[] { 2, 3, 0, 1 } } ; break; case 4: // ONNX: 2,5,7,3 + 3 // Barracuda: 2,5,7,3 + 3,_,_,_ => _,_,_,3 if (m_RanksByName[input] == 1) { permutations = new[] { 1, 2, 3, 0 } } ; // ONNX: 2,5,7,3 + 7,3 // Barracuda: 2,5,7,3 + 7,_,_,3 => _,_,7,3, else if (m_RanksByName[input] == 2) { permutations = new[] { 2, 3, 0, 1 } } ; // ONNX: 2,5,7,3 + 5,7,3 // Barracuda: 2,5,7,3 + 5,_,7,3 => _,5,7,3 else if (m_RanksByName[input] == 3) { permutations = new[] { 1, 0, 2, 3 } } ; break; } } var O = m_Ops.Transpose(X, permutations); correctedConstLayer.ApplyTensorToDataSet(O, 0); correctedConstants.Add(correctedConstLayer); layer.inputs[i] = correctedConstLayer.name; } nhwc.layers[l] = layer; } foreach (var l in correctedConstants) { nhwc.layers.Insert(0, l); } } } }
// 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); }
int[] GetPermutationForBroadcast(int targetRank, int rank, bool isNHWC = false) { int[] permutations = new[] { 0, 1, 2, 3 }; if (rank == 0 || targetRank == 1) { return(permutations); } switch (targetRank) { case 2: // ONNX: 5,7 + 7 // Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7 permutations = new[] { 1, 2, 3, 0 }; break; case 3: // ONNX: 5,7,3 + 3 // Barracuda: 5,_,3,7 + 3,_,_,_ => _,_,3,_ if (rank == 1) { permutations = new[] { 1, 2, 0, 3 } } ; // ONNX: 5,7,3 + 7,3 // Barracuda: 5,_,3,7 + 7,_,_,3 => _,_,3,7 else if (rank == 2) { permutations = new[] { 1, 2, 3, 0 } } ; break; case 4: // ONNX: 2,5,7,3 + 3 // Barracuda: 2,7,3,5 + 3,_,_,_ => _,_,3,_ if (rank == 1) { permutations = new[] { 1, 2, 0, 3 } } ; // ONNX: 2,5,7,3 + 7,3 // Barracuda: 2,7,3,5 + 7,_,_,3 => _,7,3,_ else if (rank == 2) { permutations = new[] { 1, 0, 3, 2 } } ; // ONNX: 2,5,7,3 + 5,7,3 // Barracuda: 2,7,3,5 + 5,_,3,7 => _,7,3,5 else if (rank == 3) { permutations = new[] { 1, 3, 2, 0 } } ; break; } if (isNHWC) { switch (targetRank) { case 2: // ONNX: 5,7 + 7 // Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7 permutations = new[] { 1, 2, 3, 0 }; break; case 3: // ONNX: 5,7,3 + 3 // Barracuda: 5,_,7,3 + 3,_,_,_ => _,_,_,3 if (rank == 1) { permutations = new[] { 1, 2, 3, 0 } } ; // ONNX: 5,7,3 + 7,3 // Barracuda: 5,_,7,3 + 7,_,_,3 => _,_,7,3 else if (rank == 2) { permutations = new[] { 1, 2, 0, 3 } } ; break; case 4: // ONNX: 2,5,7,3 + 3 // Barracuda: 2,5,7,3 + 3,_,_,_ => _,_,_,3 if (rank == 1) { permutations = new[] { 1, 2, 3, 0 } } ; // ONNX: 2,5,7,3 + 7,3 // Barracuda: 2,5,7,3 + 7,_,_,3 => _,_,7,3, else if (rank == 2) { permutations = new[] { 1, 2, 0, 3 } } ; // ONNX: 2,5,7,3 + 5,7,3 // Barracuda: 2,5,7,3 + 5,_,7,3 => _,5,7,3 else if (rank == 3) { permutations = new[] { 1, 0, 2, 3 } } ; break; } } return(permutations); } void CorrectConstantsForBroadCast(ref Model nhwc) { List <Layer> correctedConstants = new List <Layer>(); for (int l = 0; l < nhwc.layers.Count; l++) { Layer layer = nhwc.layers[l]; for (int i = 0; i < layer.inputs.Length; i++) { var input = layer.inputs[i]; if (!ModelAnalyzer.IsLayerBroacastable(layer)) { continue; } if (!m_RanksByName.ContainsKey(input) || !m_RanksByName.ContainsKey(layer.name)) { continue; } Layer inputLayer; bool found = ModelAnalyzer.FindLayerByName(nhwc, input, out inputLayer); if (!found) { continue; } if (!ModelOptimizer.IsLayerConstant(inputLayer)) { continue; } if (m_RanksByName[input] < 1 || m_RanksByName[input] == m_RanksByName[layer.name]) { continue; } if (inputLayer.weights.Length == 1) { continue; } if (m_RanksByName[input] > m_RanksByName[layer.name]) { throw new Exception($"constant must be lower rank than input for broadcast to work, TODO add transpose before input"); } Layer correctedConstLayer = new Layer("c_" + inputLayer.name + "For_" + layer.name, Layer.Type.Load); // transpose dataset correctedConstLayer.datasets = new Layer.DataSet[1]; Array.Copy(inputLayer.datasets, correctedConstLayer.datasets, inputLayer.datasets.Length); correctedConstLayer.datasets[0].name = correctedConstLayer.name; correctedConstLayer.weights = new BarracudaArray(inputLayer.weights.Length); var X = inputLayer.DataSetToTensor(0); var rank = m_RanksByName[layer.name].Value; var inputRank = m_RanksByName[input].Value; int[] permutations = GetPermutationForBroadcast(rank, inputRank, (m_isModelExportedFromNHWC && (m_layersChannelOrder[layer.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC))); var O = m_Ops.Transpose(X, permutations); correctedConstLayer.ApplyTensorToDataSet(O, 0); O.Dispose(); X.Dispose(); correctedConstants.Add(correctedConstLayer); layer.inputs[i] = correctedConstLayer.name; } nhwc.layers[l] = layer; } foreach (var l in correctedConstants) { nhwc.layers.Insert(0, l); } } void CorrectDynamicInputsForBroadCast(ref Model nhwc) { // insert transposes before broadcastalbe ops for (int l = 0; l < nhwc.layers.Count; l++) { Layer layer = nhwc.layers[l]; if (!ModelAnalyzer.IsLayerBroacastable(layer)) { continue; } if (!m_RanksByName.ContainsKey(layer.name) || m_RanksByName[layer.name] == null) { continue; } int maxRank = m_RanksByName[layer.name].Value; if (maxRank <= 1) { continue; } for (int i = 0; i < layer.inputs.Length; i++) { string input = layer.inputs[i]; if (!m_RanksByName.ContainsKey(input) || m_RanksByName[input] == null) { continue; } int inputRank = m_RanksByName[input].Value; if (inputRank < 1 || inputRank == maxRank) { continue; } int[] permutations = GetPermutationForBroadcast(maxRank, inputRank, (m_isModelExportedFromNHWC && (m_layersChannelOrder[layer.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC))); Layer transpose = new Layer("transpose_forbroadcast_" + layer.name + "_" + input, Layer.Type.Transpose); transpose.inputs = new[] { input }; transpose.pool = permutations; nhwc.layers[l].inputs[i] = transpose.name; nhwc.layers.Insert(l, transpose); l += 1; } } } } }