diff --git a/js/web/docs/webgpu-operators.md b/js/web/docs/webgpu-operators.md index b246e19137888..00c27fe3ab034 100644 --- a/js/web/docs/webgpu-operators.md +++ b/js/web/docs/webgpu-operators.md @@ -22,6 +22,7 @@ Do not modify directly.* | Atanh | ai.onnx(9+) | | | Attention | com.microsoft(1+) | need implementing mask and past/present | | AveragePool | ai.onnx(7-9,10,11+); com.ms.internal.nhwc(7-9,10,11+) | need perf optimization; need implementing activation | +| BatchNormalization | ai.onnx(7-8,9-13,14,15+); com.ms.internal.nhwc(7-8,9-13,14,15+) | | | BiasAdd | com.microsoft(1+) | | | BiasSplitGelu | com.microsoft(1+) | | | Cast | ai.onnx(6-8,9-12,13-18,19+) | | diff --git a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts index bac44328d8f44..80f6e3bc11195 100644 --- a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts +++ b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts @@ -3,6 +3,7 @@ import {argMax, argMin, parseArgMinMaxAttributes} from './ops/argminmax'; import {attention, parseAttentionAttributes} from './ops/attention'; +import {batchNorm} from './ops/batch-norm'; import {biasAdd} from './ops/bias-add'; import {biasSplitGelu} from './ops/bias-split-gelu'; import * as binaryOps from './ops/binary-op'; @@ -51,6 +52,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map = new ['Attention', [attention, parseAttentionAttributes]], // TODO: support new attributes for AveragePool-10 ['AveragePool', [pool.averagePool, pool.parseAveragePoolAttributes]], + ['BatchNormalization', [batchNorm]], ['BiasAdd', [biasAdd]], ['BiasSplitGelu', [biasSplitGelu]], ['Cast', [unaryOps.cast, unaryOps.parseCastAttributes]], diff --git a/js/web/lib/wasm/jsep/webgpu/ops/batch-norm.ts b/js/web/lib/wasm/jsep/webgpu/ops/batch-norm.ts new file mode 100644 index 0000000000000..ec9da2613f406 --- /dev/null +++ b/js/web/lib/wasm/jsep/webgpu/ops/batch-norm.ts @@ -0,0 +1,150 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import {env} from 'onnxruntime-common'; + +import {TensorView} from '../../tensor-view'; +import {ShapeUtil} from '../../util'; +import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; +import {ComputeContext, ProgramInfo} from '../types'; + +import {createTensorShapeVariables, enableShapesUniforms, getMaxComponents, inputVariable, outputVariable, ShaderHelper} from './common'; + +export interface BatchNormAttributes extends AttributeWithCacheKey { + readonly epsilon: number; + readonly momentum: number; + readonly spatial: boolean; + readonly trainingMode: boolean; + readonly format: 'NHWC'|'NCHW'; + readonly outputCount: number; +} + +const validateInputs = (inputs: readonly TensorView[], attributes: BatchNormAttributes): void => { + if (!inputs || inputs.length !== 5) { + throw new Error('BatchNormalization requires 5 inputs'); + } + + const checkShapeEqual = (actual: readonly number[], expected: readonly number[], message: string) => { + const r = expected.length; + if (r !== actual.length) { + throw new Error(`${message}: num dimensions != ${r}`); + } + expected.forEach((v, i) => { + if (v !== actual[i]) { + throw new Error(`${message}: dim[${i}] do not match`); + } + }); + }; + + if (inputs[0].dims.length > 1) { + const shape = attributes.format === 'NHWC' ? + (attributes.spatial ? inputs[0].dims.slice(-1) : + inputs[0].dims.slice(-1).concat(inputs[0].dims.slice(1, inputs[0].dims.length - 1))) : + inputs[0].dims.slice(1, attributes.spatial ? 2 : undefined); + checkShapeEqual(inputs[1].dims, shape, 'Invalid input scale'); + checkShapeEqual(inputs[2].dims, shape, 'Invalid input B'); + checkShapeEqual(inputs[3].dims, shape, 'Invalid input mean'); + checkShapeEqual(inputs[4].dims, shape, 'Invalid input var'); + } else { + checkShapeEqual(inputs[1].dims, [1], 'Invalid input scale'); + checkShapeEqual(inputs[2].dims, [1], 'Invalid input B'); + checkShapeEqual(inputs[3].dims, [1], 'Invalid input mean'); + checkShapeEqual(inputs[4].dims, [1], 'Invalid input var'); + } +}; + +const createBatchNormInferenceProgramInfo = + (inputs: readonly TensorView[], attributes: BatchNormAttributes): ProgramInfo => { + const {epsilon, spatial, format} = attributes; + const yShape = inputs[0].dims; + const components = spatial ? getMaxComponents(yShape[yShape.length - 1]) : 1; + const cComponents = format === 'NHWC' && yShape.length > 1 ? components : 1; + const outputSize = ShapeUtil.size(yShape) / components; + // Only support uniforms for opset version >= 9 (spatial = true). + const useShapesUniforms = enableShapesUniforms(yShape.length) && spatial; + const shapeOrRank = useShapesUniforms ? yShape.length : yShape; + const x = inputVariable('x', inputs[0].dataType, inputs[0].dims, components); + const scale = inputVariable('scale', inputs[1].dataType, inputs[1].dims, cComponents); + const bias = inputVariable('bias', inputs[2].dataType, inputs[2].dims, cComponents); + const inputMean = inputVariable('inputMean', inputs[3].dataType, inputs[3].dims, cComponents); + const inputVar = inputVariable('inputVar', inputs[4].dataType, inputs[4].dims, cComponents); + const y = outputVariable('y', inputs[0].dataType, shapeOrRank, components); + // TODO: support inputs with different data type. Current we need to make sure all inputs have the same data type. + // Otherwise, the shader compilation will fail. + const calcCOffset = (): string => { + let cOffset = ''; + if (spatial) { + cOffset = `let cOffset = ${ + yShape.length === 1 ? '0u' : + format === 'NHWC' ? `outputIndices[${yShape.length - 1}] / ${components}` : + 'outputIndices[1]'};`; + } else { + if (format === 'NCHW') { + cOffset = ` + ${y.indicesSet('outputIndices', '0', '0')} + let cOffset = ${y.indicesToOffset('outputIndices')};`; + } else { + // update C channel. + cOffset = `var cIndices = ${scale.type.indices}(0); + cIndices[0] = outputIndices[${yShape.length - 1}];`; + // update D1 x ... x Dn channels. + for (let i = 1; i < scale.rank; i++) { + cOffset += `cIndices[${i}] = outputIndices[${i}];`; + } + cOffset += `let cOffset = ${scale.indicesToOffset('cIndices')};`; + } + } + return cOffset; + }; + const getInferenceModeShaderSource = (helper: ShaderHelper) => ` + const epsilon = ${epsilon}; + ${helper.registerUniform('outputSize', 'u32').declareVariables(x, scale, bias, inputMean, inputVar, y)} + ${helper.mainStart()} + ${helper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')} + var outputIndices = ${y.offsetToIndices(`global_idx * ${components}`)}; + ${calcCOffset()} + let scale = ${scale.getByOffset('cOffset')}; + let bias = ${bias.getByOffset('cOffset')}; + let inputMean = ${inputMean.getByOffset('cOffset')}; + let inputVar = ${inputVar.getByOffset('cOffset')}; + let x = ${x.getByOffset('global_idx')}; + let value = (x - inputMean) / sqrt(inputVar + epsilon) * scale + bias; + ${y.setByOffset('global_idx', 'value')} + }`; + return { + name: 'BatchNormalization', + shaderCache: { + hint: `${attributes.epsilon}_${attributes.format}_${spatial}_${components}`, + inputDependencies: useShapesUniforms ? ['rank', 'type', 'type', 'type', 'type'] : undefined, + }, + getShaderSource: getInferenceModeShaderSource, + getRunData: () => ({ + outputs: [{dims: inputs[0].dims, dataType: inputs[0].dataType}], + dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}, + programUniforms: useShapesUniforms ? + [ + {type: 'uint32', data: outputSize}, + ...createTensorShapeVariables(yShape), + ] : + [ + {type: 'uint32', data: outputSize}, + ], + }), + }; + }; + +export const parseBatchNormAttributes = (attributes: Record): BatchNormAttributes => + createAttributeWithCacheKey(attributes as Omit); + +export const batchNorm = (context: ComputeContext, attributes: Record): void => { + const {inputs, outputCount} = context; + const updatedAttributes = parseBatchNormAttributes({...attributes, outputCount}); + if (env.webgpu.validateInputContent) { + validateInputs(inputs, updatedAttributes); + } + if (attributes.trainingMode) { + throw new Error('BatchNormalization trainingMode is not supported yet.'); + } else { + context.compute(createBatchNormInferenceProgramInfo(inputs, updatedAttributes)); + } +}; diff --git a/js/web/test/data/ops/batch-norm.jsonc b/js/web/test/data/ops/batch-norm.jsonc new file mode 100644 index 0000000000000..4ea16f290dc8f --- /dev/null +++ b/js/web/test/data/ops/batch-norm.jsonc @@ -0,0 +1,446 @@ +[ + { + "name": "BatchNormalization with no attributes", + "operator": "BatchNormalization", + "attributes": [], + "cases": [ + { + "name": "T[64]", + "inputs": [ + { + "data": [ + 2.02384, -0.935186, 0.488569, -0.513934, -1.27082, -0.131913, -1.806, -0.37904, 0.667796, -1.14826, + 1.2522, 0.0300339, 2.4758, 1.55511, 0.385341, 1.46645, -1.09355, -2.56309, 0.976015, -1.47036, 0.89486, + 0.580989, -1.12418, -0.339189, 1.3314, 0.418893, -0.301401, -1.2983, -0.839063, 0.170261, 1.15486, + -0.255735, -0.589851, -0.416289, -0.952648, -0.360487, 0.253287, 0.437195, 0.32023, 0.209606, -0.279519, + -0.546527, 0.265286, -1.07383, -1.65879, 1.1222, 0.946612, 0.822549, 0.64689, -0.292639, -0.73995, + -0.694949, 1.33899, -0.0652476, 1.61791, 1.49692, -0.761145, -0.201874, -1.15431, -1.83111, -0.705267, + -0.143026, -0.129819, -0.799425 + ], + "dims": [64], + "type": "float32" + }, + { + "data": [0.241661], + "dims": [1], + "type": "float32" + }, + { + "data": [0], + "dims": [1], + "type": "float32" + }, + { + "data": [0], + "dims": [1], + "type": "float32" + }, + { + "data": [1], + "dims": [1], + "type": "float32" + } + ], + "outputs": [ + { + "data": [ + 0.489082, -0.225997, 0.118068, -0.124197, -0.307105, -0.031878, -0.436439, -0.0915989, 0.16138, -0.277489, + 0.302606, 0.007258, 0.598301, 0.375807, 0.0931215, 0.354382, -0.264267, -0.619395, 0.235864, -0.355328, + 0.216252, 0.140402, -0.271669, -0.0819684, 0.321747, 0.10123, -0.0728365, -0.313746, -0.202768, 0.0411454, + 0.279085, -0.0618009, -0.142543, -0.1006, -0.230217, -0.0871152, 0.0612094, 0.105652, 0.0773867, + 0.0506533, -0.0675486, -0.132074, 0.064109, -0.259501, -0.400863, 0.271191, 0.228758, 0.198777, 0.156327, + -0.0707191, -0.178816, -0.167941, 0.323581, -0.0157677, 0.390985, 0.361745, -0.183938, -0.0487849, + -0.27895, -0.442507, -0.170435, -0.0345637, -0.031372, -0.193189 + ], + "dims": [64], + "type": "float32" + } + ] + }, + { + "name": "T[2,3,4,4,4]", + "inputs": [ + { + "data": [ + 2.02384, -0.935186, 0.488569, -0.513934, -1.27082, -0.131913, -1.806, -0.37904, 0.667796, -1.14826, + 1.2522, 0.0300339, 2.4758, 1.55511, 0.385341, 1.46645, -1.09355, -2.56309, 0.976015, -1.47036, 0.89486, + 0.580989, -1.12418, -0.339189, 1.3314, 0.418893, -0.301401, -1.2983, -0.839063, 0.170261, 1.15486, + -0.255735, -0.589851, -0.416289, -0.952648, -0.360487, 0.253287, 0.437195, 0.32023, 0.209606, -0.279519, + -0.546527, 0.265286, -1.07383, -1.65879, 1.1222, 0.946612, 0.822549, 0.64689, -0.292639, -0.73995, + -0.694949, 1.33899, -0.0652476, 1.61791, 1.49692, -0.761145, -0.201874, -1.15431, -1.83111, -0.705267, + -0.143026, -0.129819, -0.799425, 0.168795, 0.740422, -0.377683, 0.432598, -2.07414, -2.85251, 0.273531, + 0.0532606, 1.31052, -0.769382, 0.9976, 0.850536, -1.53812, -0.00496016, 0.931242, 0.0517056, -0.497829, + 0.275869, 0.860001, 1.23747, 0.179686, 1.5914, 0.740327, 0.798208, 2.12478, 1.74205, -0.322054, + -0.0112451, 0.204525, -0.431252, -1.3114, 0.186204, 0.780569, -1.42994, 1.63344, -0.00839034, -0.187035, + 1.8406, 1.32053, -0.636963, 0.408944, -1.50846, -1.2076, -0.129118, -0.0441307, 1.47558, 1.07251, 1.05295, + -0.420297, -1.13402, -0.524053, 3.20754, -0.588935, -0.527549, 0.591928, -1.10529, 0.520412, 0.19404, + -1.21229, -0.399594, -0.280935, -0.363324, -0.00804771, 1.43102, -0.523222, 1.17608, -0.53195, 0.914993, + 2.69308, -0.517211, 0.472273, -0.464725, -0.929768, -0.631145, 0.919709, -0.27391, 1.76689, 0.894897, + 0.235798, 1.2544, 0.858985, -0.139707, 0.354544, 0.200878, 0.353255, 0.0722632, -1.56074, 1.03685, + 1.73434, 0.193269, -0.864609, 0.842739, -0.372717, 0.584484, 0.16315, 1.60674, -0.0611289, -1.24544, + 1.33361, -0.961942, -0.15732, -0.348637, 0.361842, 0.7386, 0.517256, 1.20406, -2.07277, -1.01983, -1.9163, + 0.239934, 0.177979, 0.464564, 0.988822, 0.284607, -1.56099, -0.429143, 0.111043, -0.0853688, -0.319176, + -0.279777, 0.520971, -1.078, -0.670242, 0.065652, 0.468538, -0.825062, 0.370068, 1.68751, -1.16928, + -0.411782, 1.61624, -0.973004, 2.64703, -0.220014, -1.43954, -0.018692, 1.34982, -0.95197, -1.72586, + 1.32725, 0.280984, 0.00847463, 0.512869, 0.0378154, 0.13898, 0.35758, -0.084558, 1.04045, -1.79933, + 1.3002, 0.390457, 1.22267, 0.959344, -0.964296, -0.0935597, 0.288953, -0.158046, 0.532672, -0.500988, + 0.25187, -2.14384, -0.633315, 1.24612, -1.41525, 0.36494, -0.00714732, -0.608963, 0.508496, 0.995365, + 1.21159, -0.169055, -0.968783, 1.52779, -0.082381, 2.2049, 0.928655, 0.120245, 0.911429, -0.885258, + -1.2072, 0.770694, 2.36621, 1.08456, -1.60069, 0.0345025, 0.359559, -0.785411, 0.466532, -0.78543, + 0.024879, 1.59337, 1.13718, -1.27073, -0.263788, -1.7702, 0.203263, 1.34631, 1.11914, -2.04911, -0.804137, + 0.466763, 2.18386, 1.4689, 0.898297, -0.648948, 0.252202, 1.12501, -0.204563, 0.124608, 0.377214, + 0.894327, -0.249118, 0.709188, 0.999397, -1.4079, 0.193873, 0.657753, -0.709732, 1.09897, -0.145793, + 0.779199, 0.88378, -1.2676, 1.15709, 0.62295, -0.370894, -0.103268, -1.55949, -0.470747, 0.100394, + 0.422334, -0.0685312, -0.434488, -0.568974, -0.256987, 2.01276, -0.923322, -0.613144, 1.50676, 0.65756, + 1.20524, 1.10395, -0.975241, 2.44035, 1.08276, 0.330393, -0.508918, -1.25545, 0.189815, -0.156263, + -0.960866, 1.0859, -0.674478, 2.76743, 1.21399, 1.71666, -1.73198, -1.1062, 0.951285, -0.713336, 1.61586, + 1.96514, 0.002603, 0.0953297, 0.949256, -1.76552, 0.372816, -0.781229, 1.50532, 1.28462, 1.31116, + 0.731908, 1.54835, 0.371081, 0.409244, -0.106938, -1.79396, -1.61198, -0.80869, -1.10381, 1.1872, + -0.832439, 0.0755941, -1.09553, 0.960059, 1.44252, -0.196482, -1.07364, 0.165547, 0.630078, 1.56569, + -0.669592, 1.15974, 0.0953399, -0.202313, 0.812631, -0.318567, -0.16644, 0.887062, -0.0264821, -0.740725, + 0.0797577, -1.1037, 0.90236, 1.13427, 0.364186, -2.01043, -0.415748, 0.116046, 0.369949, 0.317886, + 0.530332, 1.48341, 0.74666, -1.64142, 0.22569, 1.18015, 1.31827, -1.33904, -0.101125 + ], + "dims": [2, 3, 4, 4, 4], + "type": "float32" + }, + { + "data": [0.241661, 0.960798, 0.474727], + "dims": [3], + "type": "float32" + }, + { + "data": [0, 0, 0], + "dims": [3], + "type": "float32" + }, + { + "data": [0, 0, 0], + "dims": [3], + "type": "float32" + }, + { + "data": [1, 1, 1], + "dims": [3], + "type": "float32" + } + ], + "outputs": [ + { + "data": [ + 0.489082, -0.225997, 0.118068, -0.124197, -0.307105, -0.031878, -0.436439, -0.0915989, 0.16138, -0.277489, + 0.302606, 0.007258, 0.598301, 0.375807, 0.0931215, 0.354382, -0.264267, -0.619395, 0.235864, -0.355328, + 0.216252, 0.140402, -0.271669, -0.0819684, 0.321747, 0.10123, -0.0728365, -0.313746, -0.202768, 0.0411454, + 0.279085, -0.0618009, -0.142543, -0.1006, -0.230217, -0.0871152, 0.0612094, 0.105652, 0.0773867, + 0.0506533, -0.0675486, -0.132074, 0.064109, -0.259501, -0.400863, 0.271191, 0.228758, 0.198777, 0.156327, + -0.0707191, -0.178816, -0.167941, 0.323581, -0.0157677, 0.390985, 0.361745, -0.183938, -0.0487849, + -0.27895, -0.442507, -0.170435, -0.0345637, -0.031372, -0.193189, 0.162177, 0.711393, -0.362876, 0.415637, + -1.99282, -2.74067, 0.262807, 0.0511725, 1.25914, -0.739217, 0.958488, 0.817189, -1.47782, -0.00476569, + 0.894731, 0.0496784, -0.478311, 0.265053, 0.826283, 1.18895, 0.172641, 1.52901, 0.711301, 0.766913, + 2.04147, 1.67375, -0.309427, -0.0108042, 0.196507, -0.414344, -1.25999, 0.178903, 0.749965, -1.37387, + 1.5694, -0.00806138, -0.179702, 1.76844, 1.26875, -0.61199, 0.392911, -1.44932, -1.16025, -0.124055, + -0.0424004, 1.41773, 1.03046, 1.01167, -0.403818, -1.08956, -0.503507, 3.08178, -0.565845, -0.506866, + 0.56872, -1.06196, 0.500008, 0.186433, -1.16476, -0.383928, -0.269921, -0.349079, -0.00773219, 1.37492, + -0.248386, 0.558316, -0.25253, 0.43437, 1.27847, -0.245533, 0.2242, -0.220617, -0.441384, -0.29962, + 0.436609, -0.130032, 0.838785, 0.424829, 0.111939, 0.595496, 0.407781, -0.0663221, 0.168311, 0.0953618, + 0.167699, 0.0343051, -0.74092, 0.492219, 0.823334, 0.0917494, -0.410451, 0.400069, -0.176938, 0.277469, + 0.0774512, 0.762761, -0.0290194, -0.59124, 0.6331, -0.456657, -0.0746837, -0.165507, 0.171775, 0.350631, + 0.245554, 0.571595, -0.983996, -0.484139, -0.909715, 0.113902, 0.0844908, 0.22054, 0.469418, 0.13511, + -0.741041, -0.203725, 0.0527148, -0.0405267, -0.151521, -0.132817, 0.247318, -0.511752, -0.31818, + 0.0311666, 0.222426, -0.391677, 0.17568, 0.801104, -0.282569, -0.0995112, 0.39058, -0.235136, 0.639682, + -0.0531687, -0.347878, -0.0045171, 0.326198, -0.230053, -0.41707, 0.320744, 0.0679025, 0.00204798, + 0.12394, 0.00913847, 0.0335859, 0.0864127, -0.0204343, 0.251436, -0.434827, 0.314206, 0.0943579, 0.295471, + 0.231835, -0.233032, -0.0226096, 0.0698283, -0.0381934, 0.128725, -0.121069, 0.060867, -0.51808, + -0.153047, 0.301137, -0.342009, 0.0881915, -0.00172722, -0.147162, 0.122883, 0.24054, 0.292792, + -0.0408538, -0.234116, 0.369206, -0.0199082, 0.532835, 0.224419, 0.0290583, 0.220256, -0.213931, + -0.291733, 0.186246, 0.571817, 0.262095, -0.386822, 0.00833788, 0.086891, -0.189802, 0.112742, -0.189807, + 0.00601226, 0.385054, 0.274811, -1.22091, -0.253445, -1.7008, 0.195294, 1.29353, 1.07526, -1.96877, + -0.772609, 0.448463, 2.09824, 1.4113, 0.863078, -0.623505, 0.242314, 1.0809, -0.196543, 0.119722, + 0.362425, 0.859263, -0.239351, 0.681383, 0.960214, -1.3527, 0.186272, 0.631964, -0.681905, 1.05588, + -0.140077, 0.748649, 0.84913, -1.2179, 1.11172, 0.598526, -0.356353, -0.099219, -1.49835, -0.452291, + 0.0964582, 0.405776, -0.0658444, -0.417454, -0.546667, -0.246911, 1.93385, -0.887121, -0.589104, 1.44769, + 0.631779, 1.15798, 1.06067, -0.937005, 2.34467, 1.04031, 0.31744, -0.488965, -1.20623, 0.182373, + -0.150136, -0.923194, 1.04332, -0.648034, 2.65893, 1.1664, 1.64935, -0.822216, -0.525139, 0.451599, + -0.338638, 0.767087, 0.932899, 0.00123571, 0.0452554, 0.450635, -0.838136, 0.176985, -0.370868, 0.714614, + 0.60984, 0.622438, 0.347455, 0.73504, 0.176161, 0.194278, -0.0507662, -0.851639, -0.765246, -0.383905, + -0.524005, 0.563593, -0.395179, 0.0358864, -0.520076, 0.455763, 0.684801, -0.093275, -0.509682, 0.0785892, + 0.299113, 0.743272, -0.317872, 0.550556, 0.0452602, -0.0960432, 0.385776, -0.151232, -0.079013, 0.42111, + -0.0125717, -0.35164, 0.0378629, -0.523955, 0.428372, 0.538468, 0.172888, -0.954402, -0.197366, 0.0550898, + 0.175624, 0.150908, 0.251761, 0.704209, 0.354458, -0.779221, 0.107141, 0.560244, 0.625814, -0.635675, + -0.0480064 + ], + "dims": [2, 3, 4, 4, 4], + "type": "float32" + } + ] + } + ] + }, + { + "name": "BatchNormalization with no attributes - NHWC", + "operator": "BatchNormalization", + "opset": { "domain": "com.ms.internal.nhwc", "version": 12 }, + "attributes": [], + "cases": [ + { + "name": "T[64]", + "inputs": [ + { + "data": [ + 2.02384, -0.935186, 0.488569, -0.513934, -1.27082, -0.131913, -1.806, -0.37904, 0.667796, -1.14826, + 1.2522, 0.0300339, 2.4758, 1.55511, 0.385341, 1.46645, -1.09355, -2.56309, 0.976015, -1.47036, 0.89486, + 0.580989, -1.12418, -0.339189, 1.3314, 0.418893, -0.301401, -1.2983, -0.839063, 0.170261, 1.15486, + -0.255735, -0.589851, -0.416289, -0.952648, -0.360487, 0.253287, 0.437195, 0.32023, 0.209606, -0.279519, + -0.546527, 0.265286, -1.07383, -1.65879, 1.1222, 0.946612, 0.822549, 0.64689, -0.292639, -0.73995, + -0.694949, 1.33899, -0.0652476, 1.61791, 1.49692, -0.761145, -0.201874, -1.15431, -1.83111, -0.705267, + -0.143026, -0.129819, -0.799425 + ], + "dims": [64], + "type": "float32" + }, + { + "data": [0.241661], + "dims": [1], + "type": "float32" + }, + { + "data": [0], + "dims": [1], + "type": "float32" + }, + { + "data": [0], + "dims": [1], + "type": "float32" + }, + { + "data": [1], + "dims": [1], + "type": "float32" + } + ], + "outputs": [ + { + "data": [ + 0.489082, -0.225997, 0.118068, -0.124197, -0.307105, -0.031878, -0.436439, -0.0915989, 0.16138, -0.277489, + 0.302606, 0.007258, 0.598301, 0.375807, 0.0931215, 0.354382, -0.264267, -0.619395, 0.235864, -0.355328, + 0.216252, 0.140402, -0.271669, -0.0819684, 0.321747, 0.10123, -0.0728365, -0.313746, -0.202768, 0.0411454, + 0.279085, -0.0618009, -0.142543, -0.1006, -0.230217, -0.0871152, 0.0612094, 0.105652, 0.0773867, + 0.0506533, -0.0675486, -0.132074, 0.064109, -0.259501, -0.400863, 0.271191, 0.228758, 0.198777, 0.156327, + -0.0707191, -0.178816, -0.167941, 0.323581, -0.0157677, 0.390985, 0.361745, -0.183938, -0.0487849, + -0.27895, -0.442507, -0.170435, -0.0345637, -0.031372, -0.193189 + ], + "dims": [64], + "type": "float32" + } + ] + }, + { + "name": "T[2,4,4,4,3]", + "inputs": [ + { + "data": [ + 2.02384, 0.168795, -0.523222, -0.935186, 0.740422, 1.17608, 0.488569, -0.377683, -0.53195, -0.513934, + 0.432598, 0.914993, -1.27082, -2.07414, 2.69308, -0.131913, -2.85251, -0.517211, -1.806, 0.273531, + 0.472273, -0.37904, 0.0532606, -0.464725, 0.667796, 1.31052, -0.929768, -1.14826, -0.769382, -0.631145, + 1.2522, 0.9976, 0.919709, 0.0300339, 0.850536, -0.27391, 2.4758, -1.53812, 1.76689, 1.55511, -0.00496016, + 0.894897, 0.385341, 0.931242, 0.235798, 1.46645, 0.0517056, 1.2544, -1.09355, -0.497829, 0.858985, + -2.56309, 0.275869, -0.139707, 0.976015, 0.860001, 0.354544, -1.47036, 1.23747, 0.200878, 0.89486, + 0.179686, 0.353255, 0.580989, 1.5914, 0.0722632, -1.12418, 0.740327, -1.56074, -0.339189, 0.798208, + 1.03685, 1.3314, 2.12478, 1.73434, 0.418893, 1.74205, 0.193269, -0.301401, -0.322054, -0.864609, -1.2983, + -0.0112451, 0.842739, -0.839063, 0.204525, -0.372717, 0.170261, -0.431252, 0.584484, 1.15486, -1.3114, + 0.16315, -0.255735, 0.186204, 1.60674, -0.589851, 0.780569, -0.0611289, -0.416289, -1.42994, -1.24544, + -0.952648, 1.63344, 1.33361, -0.360487, -0.00839034, -0.961942, 0.253287, -0.187035, -0.15732, 0.437195, + 1.8406, -0.348637, 0.32023, 1.32053, 0.361842, 0.209606, -0.636963, 0.7386, -0.279519, 0.408944, 0.517256, + -0.546527, -1.50846, 1.20406, 0.265286, -1.2076, -2.07277, -1.07383, -0.129118, -1.01983, -1.65879, + -0.0441307, -1.9163, 1.1222, 1.47558, 0.239934, 0.946612, 1.07251, 0.177979, 0.822549, 1.05295, 0.464564, + 0.64689, -0.420297, 0.988822, -0.292639, -1.13402, 0.284607, -0.73995, -0.524053, -1.56099, -0.694949, + 3.20754, -0.429143, 1.33899, -0.588935, 0.111043, -0.0652476, -0.527549, -0.0853688, 1.61791, 0.591928, + -0.319176, 1.49692, -1.10529, -0.279777, -0.761145, 0.520412, 0.520971, -0.201874, 0.19404, -1.078, + -1.15431, -1.21229, -0.670242, -1.83111, -0.399594, 0.065652, -0.705267, -0.280935, 0.468538, -0.143026, + -0.363324, -0.825062, -0.129819, -0.00804771, 0.370068, -0.799425, 1.43102, 1.68751, -1.16928, -1.27073, + -1.73198, -0.411782, -0.263788, -1.1062, 1.61624, -1.7702, 0.951285, -0.973004, 0.203263, -0.713336, + 2.64703, 1.34631, 1.61586, -0.220014, 1.11914, 1.96514, -1.43954, -2.04911, 0.002603, -0.018692, + -0.804137, 0.0953297, 1.34982, 0.466763, 0.949256, -0.95197, 2.18386, -1.76552, -1.72586, 1.4689, + 0.372816, 1.32725, 0.898297, -0.781229, 0.280984, -0.648948, 1.50532, 0.00847463, 0.252202, 1.28462, + 0.512869, 1.12501, 1.31116, 0.0378154, -0.204563, 0.731908, 0.13898, 0.124608, 1.54835, 0.35758, 0.377214, + 0.371081, -0.084558, 0.894327, 0.409244, 1.04045, -0.249118, -0.106938, -1.79933, 0.709188, -1.79396, + 1.3002, 0.999397, -1.61198, 0.390457, -1.4079, -0.80869, 1.22267, 0.193873, -1.10381, 0.959344, 0.657753, + 1.1872, -0.964296, -0.709732, -0.832439, -0.0935597, 1.09897, 0.0755941, 0.288953, -0.145793, -1.09553, + -0.158046, 0.779199, 0.960059, 0.532672, 0.88378, 1.44252, -0.500988, -1.2676, -0.196482, 0.25187, + 1.15709, -1.07364, -2.14384, 0.62295, 0.165547, -0.633315, -0.370894, 0.630078, 1.24612, -0.103268, + 1.56569, -1.41525, -1.55949, -0.669592, 0.36494, -0.470747, 1.15974, -0.00714732, 0.100394, 0.0953399, + -0.608963, 0.422334, -0.202313, 0.508496, -0.0685312, 0.812631, 0.995365, -0.434488, -0.318567, 1.21159, + -0.568974, -0.16644, -0.169055, -0.256987, 0.887062, -0.968783, 2.01276, -0.0264821, 1.52779, -0.923322, + -0.740725, -0.082381, -0.613144, 0.0797577, 2.2049, 1.50676, -1.1037, 0.928655, 0.65756, 0.90236, + 0.120245, 1.20524, 1.13427, 0.911429, 1.10395, 0.364186, -0.885258, -0.975241, -2.01043, -1.2072, 2.44035, + -0.415748, 0.770694, 1.08276, 0.116046, 2.36621, 0.330393, 0.369949, 1.08456, -0.508918, 0.317886, + -1.60069, -1.25545, 0.530332, 0.0345025, 0.189815, 1.48341, 0.359559, -0.156263, 0.74666, -0.785411, + -0.960866, -1.64142, 0.466532, 1.0859, 0.22569, -0.78543, -0.674478, 1.18015, 0.024879, 2.76743, 1.31827, + 1.59337, 1.21399, -1.33904, 1.13718, 1.71666, -0.101125 + ], + "dims": [2, 4, 4, 4, 3], + "type": "float32" + }, + { + "data": [0.241661, 0.960798, 0.474727], + "dims": [3], + "type": "float32" + }, + { + "data": [0, 0, 0], + "dims": [3], + "type": "float32" + }, + { + "data": [0, 0, 0], + "dims": [3], + "type": "float32" + }, + { + "data": [1, 1, 1], + "dims": [3], + "type": "float32" + } + ], + "outputs": [ + { + "data": [ + 0.489082, 0.162177, -0.248386, -0.225997, 0.711393, 0.558316, 0.118068, -0.362876, -0.25253, -0.124197, + 0.415637, 0.43437, -0.307105, -1.99282, 1.27847, -0.031878, -2.74067, -0.245533, -0.436439, 0.262807, + 0.2242, -0.0915989, 0.0511725, -0.220617, 0.16138, 1.25914, -0.441384, -0.277489, -0.739217, -0.29962, + 0.302606, 0.958488, 0.436609, 0.007258, 0.817189, -0.130032, 0.598301, -1.47782, 0.838785, 0.375807, + -0.00476569, 0.424829, 0.0931215, 0.894731, 0.111939, 0.354382, 0.0496784, 0.595496, -0.264267, -0.478311, + 0.407781, -0.619395, 0.265053, -0.0663221, 0.235864, 0.826283, 0.168311, -0.355328, 1.18895, 0.0953618, + 0.216252, 0.172641, 0.167699, 0.140402, 1.52901, 0.0343051, -0.271669, 0.711301, -0.74092, -0.0819684, + 0.766913, 0.492219, 0.321747, 2.04147, 0.823334, 0.10123, 1.67375, 0.0917494, -0.0728365, -0.309427, + -0.410451, -0.313746, -0.0108042, 0.400069, -0.202768, 0.196507, -0.176938, 0.0411454, -0.414344, + 0.277469, 0.279085, -1.25999, 0.0774512, -0.0618009, 0.178903, 0.762761, -0.142543, 0.749965, -0.0290194, + -0.1006, -1.37387, -0.59124, -0.230217, 1.5694, 0.6331, -0.0871152, -0.00806138, -0.456657, 0.0612094, + -0.179702, -0.0746837, 0.105652, 1.76844, -0.165507, 0.0773867, 1.26875, 0.171775, 0.0506533, -0.61199, + 0.350631, -0.0675486, 0.392911, 0.245554, -0.132074, -1.44932, 0.571595, 0.064109, -1.16025, -0.983996, + -0.259501, -0.124055, -0.484139, -0.400863, -0.0424004, -0.909715, 0.271191, 1.41773, 0.113902, 0.228758, + 1.03046, 0.0844908, 0.198777, 1.01167, 0.22054, 0.156327, -0.403818, 0.469418, -0.0707191, -1.08956, + 0.13511, -0.178816, -0.503507, -0.741041, -0.167941, 3.08178, -0.203725, 0.323581, -0.565845, 0.0527148, + -0.0157677, -0.506866, -0.0405267, 0.390985, 0.56872, -0.151521, 0.361745, -1.06196, -0.132817, -0.183938, + 0.500008, 0.247318, -0.0487849, 0.186433, -0.511752, -0.27895, -1.16476, -0.31818, -0.442507, -0.383928, + 0.0311666, -0.170435, -0.269921, 0.222426, -0.0345637, -0.349079, -0.391677, -0.031372, -0.00773219, + 0.17568, -0.193189, 1.37492, 0.801104, -0.282569, -1.22091, -0.822216, -0.0995112, -0.253445, -0.525139, + 0.39058, -1.7008, 0.451599, -0.235136, 0.195294, -0.338638, 0.639682, 1.29353, 0.767087, -0.0531687, + 1.07526, 0.932899, -0.347878, -1.96877, 0.00123571, -0.0045171, -0.772609, 0.0452554, 0.326198, 0.448463, + 0.450635, -0.230053, 2.09824, -0.838136, -0.41707, 1.4113, 0.176985, 0.320744, 0.863078, -0.370868, + 0.0679025, -0.623505, 0.714614, 0.00204798, 0.242314, 0.60984, 0.12394, 1.0809, 0.622438, 0.00913847, + -0.196543, 0.347455, 0.0335859, 0.119722, 0.73504, 0.0864127, 0.362425, 0.176161, -0.0204343, 0.859263, + 0.194278, 0.251436, -0.239351, -0.0507662, -0.434827, 0.681383, -0.851639, 0.314206, 0.960214, -0.765246, + 0.0943579, -1.3527, -0.383905, 0.295471, 0.186272, -0.524005, 0.231835, 0.631964, 0.563593, -0.233032, + -0.681905, -0.395179, -0.0226096, 1.05588, 0.0358864, 0.0698283, -0.140077, -0.520076, -0.0381934, + 0.748649, 0.455763, 0.128725, 0.84913, 0.684801, -0.121069, -1.2179, -0.093275, 0.060867, 1.11172, + -0.509682, -0.51808, 0.598526, 0.0785892, -0.153047, -0.356353, 0.299113, 0.301137, -0.099219, 0.743272, + -0.342009, -1.49835, -0.317872, 0.0881915, -0.452291, 0.550556, -0.00172722, 0.0964582, 0.0452602, + -0.147162, 0.405776, -0.0960432, 0.122883, -0.0658444, 0.385776, 0.24054, -0.417454, -0.151232, 0.292792, + -0.546667, -0.079013, -0.0408538, -0.246911, 0.42111, -0.234116, 1.93385, -0.0125717, 0.369206, -0.887121, + -0.35164, -0.0199082, -0.589104, 0.0378629, 0.532835, 1.44769, -0.523955, 0.224419, 0.631779, 0.428372, + 0.0290583, 1.15798, 0.538468, 0.220256, 1.06067, 0.172888, -0.213931, -0.937005, -0.954402, -0.291733, + 2.34467, -0.197366, 0.186246, 1.04031, 0.0550898, 0.571817, 0.31744, 0.175624, 0.262095, -0.488965, + 0.150908, -0.386822, -1.20623, 0.251761, 0.00833788, 0.182373, 0.704209, 0.086891, -0.150136, 0.354458, + -0.189802, -0.923194, -0.779221, 0.112742, 1.04332, 0.107141, -0.189807, -0.648034, 0.560244, 0.00601226, + 2.65893, 0.625814, 0.385054, 1.1664, -0.635675, 0.274811, 1.64935, -0.0480064 + ], + "dims": [2, 4, 4, 4, 3], + "type": "float32" + } + ] + } + ] + }, + { + "name": "BatchNormalization non-spatial mode", + "operator": "BatchNormalization", + "opset": { "domain": "", "version": 7 }, + "attributes": [{ "name": "spatial", "data": 0, "type": "int" }], + "cases": [ + { + "name": "T[3,1,2]", + "inputs": [ + { + "data": [0.2134, 0.32434, 0.5644, 0.3234, 0.4545, 0.3445], + "dims": [3, 1, 2], + "type": "float32" + }, + { + "data": [0.5, 0.6], + "dims": [1, 2], + "type": "float32" + }, + { + "data": [0.2, 0.1], + "dims": [1, 2], + "type": "float32" + }, + { + "data": [0.034, 0.342], + "dims": [1, 2], + "type": "float32" + }, + { + "data": [1, 1], + "dims": [1, 2], + "type": "float32" + } + ], + "outputs": [ + { + "data": [0.2897, 0.089404, 0.4652, 0.08884, 0.41025, 0.1015], + "dims": [3, 1, 2], + "type": "float32" + } + ] + } + ] + }, + { + "name": "BatchNormalization non-spatial mode - NHWC", + "operator": "BatchNormalization", + "opset": { "domain": "com.ms.internal.nhwc", "version": 7 }, + "attributes": [{ "name": "spatial", "data": 0, "type": "int" }], + "cases": [ + { + "name": "T[3,2,1]", + "inputs": [ + { + "data": [0.2134, 0.32434, 0.5644, 0.3234, 0.4545, 0.3445], + "dims": [3, 2, 1], + "type": "float32" + }, + { + "data": [0.5, 0.6], + "dims": [1, 2], + "type": "float32" + }, + { + "data": [0.2, 0.1], + "dims": [1, 2], + "type": "float32" + }, + { + "data": [0.034, 0.342], + "dims": [1, 2], + "type": "float32" + }, + { + "data": [1, 1], + "dims": [1, 2], + "type": "float32" + } + ], + "outputs": [ + { + "data": [0.2897, 0.089404, 0.4652, 0.08884, 0.41025, 0.1015], + "dims": [3, 2, 1], + "type": "float32" + } + ] + } + ] + } +] diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc index 37aa9394c7f96..a313adef7151b 100644 --- a/js/web/test/suite-test-list.jsonc +++ b/js/web/test/suite-test-list.jsonc @@ -1337,6 +1337,7 @@ //"and.jsonc", "asin.jsonc", "attention.jsonc", + "batch-norm.jsonc", "bias-add.jsonc", "bias-split-gelu.jsonc", "ceil.jsonc", diff --git a/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc b/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc index 03ad95260c0ad..c8960578f9e3d 100644 --- a/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc +++ b/onnxruntime/core/graph/contrib_ops/internal_nhwc_onnx_schemas.cc @@ -101,6 +101,7 @@ void OpSet_Internal_NHWC_ONNX::ForEachSchema(const std::function RegisterKernels() { auto kernel_registry = std::make_unique(); @@ -636,6 +645,15 @@ std::unique_ptr RegisterKernels() { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/js/operators/batch_norm.cc b/onnxruntime/core/providers/js/operators/batch_norm.cc new file mode 100644 index 0000000000000..e18ad835792f7 --- /dev/null +++ b/onnxruntime/core/providers/js/operators/batch_norm.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "batch_norm.h" + +namespace onnxruntime { +namespace js { + +#define REGISTER_BATCHNORM_KERNEL(OP_TYPE, DOMAIN, KERNEL_CLASS) \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + OP_TYPE, DOMAIN, 7, 8, kJsExecutionProvider, \ + KernelDefBuilder().TypeConstraint("T", JsepSupportedFloatTypes()), KERNEL_CLASS); \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + OP_TYPE, DOMAIN, 9, 13, kJsExecutionProvider, \ + KernelDefBuilder().TypeConstraint("T", JsepSupportedFloatTypes()), KERNEL_CLASS); \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX(OP_TYPE, DOMAIN, 14, 14, kJsExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", JsepSupportedFloatTypes()) \ + .TypeConstraint("U", JsepSupportedFloatTypes()), \ + KERNEL_CLASS); \ + ONNX_OPERATOR_KERNEL_EX(OP_TYPE, DOMAIN, 15, kJsExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", JsepSupportedFloatTypes()) \ + .TypeConstraint("T1", JsepSupportedFloatTypes()) \ + .TypeConstraint("T2", JsepSupportedFloatTypes()), \ + KERNEL_CLASS); + +REGISTER_BATCHNORM_KERNEL(BatchNormalization, kMSInternalNHWCDomain, BatchNorm); +REGISTER_BATCHNORM_KERNEL(BatchNormalization, kOnnxDomain, BatchNorm); + +} // namespace js +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/js/operators/batch_norm.h b/onnxruntime/core/providers/js/operators/batch_norm.h new file mode 100644 index 0000000000000..bb987a8aeab44 --- /dev/null +++ b/onnxruntime/core/providers/js/operators/batch_norm.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/js/js_kernel.h" + +namespace onnxruntime { +namespace js { + +template +class BatchNorm final : public JsKernel { + public: + explicit BatchNorm(const OpKernelInfo& info) : JsKernel(info) { + float epsilon = info.GetAttrOrDefault("epsilon", 1e-5); + float momentum = info.GetAttrOrDefault("momentum", 0.9); + int64_t spatial = info.GetAttrOrDefault("spatial", 1); + + const auto& node = info.node(); + int opset = node.SinceVersion(); + int64_t training_mode = opset <= 9 ? info.GetOutputCount() > 1 : info.GetAttrOrDefault("training_mode", 0); + + JSEP_INIT_KERNEL_ATTRIBUTE(BatchNormalization, ({ + "epsilon" : $1, + "momentum" : $2, + "spatial" : !!$4, + "trainingMode" : !!$3, + "format" : $5 ? "NHWC" : "NCHW", + }), + static_cast(epsilon), static_cast(momentum), + static_cast(training_mode), static_cast(spatial), + static_cast(is_channels_last)); + } +}; + +} // namespace js +} // namespace onnxruntime