mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-19 19:00:47 +00:00
[Web/JS] Add ConvTranspose support (#16433)
### Description Add ConvTranspose support for WebGPU ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
parent
5c6613875c
commit
00e8f2a2a9
8 changed files with 846 additions and 9 deletions
|
|
@ -23,6 +23,7 @@ Do not modify directly.*
|
|||
| Clip | ai.onnx(6-10,11,12,13+) | |
|
||||
| Concat | ai.onnx(1-3,4-10,11-12,13+) | |
|
||||
| Conv | ai.onnx(1-10,11+); com.ms.internal.nhwc(11+) | need perf optimization; conv3d not supported; need implementing activation |
|
||||
| ConvTranspose | ai.onnx(1-10,11+); com.ms.internal.nhwc(11+) | |
|
||||
| Cos | ai.onnx(7+) | |
|
||||
| Cosh | ai.onnx(9+) | |
|
||||
| Div | ai.onnx(7-12,13,14+) | |
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
import * as binaryOps from './ops/binary-op';
|
||||
import {concat, parseConcatAttributes} from './ops/concat';
|
||||
import {conv, parseConvAttributes} from './ops/conv';
|
||||
import {convTranspose, parseConvTransposeAttributes} from './ops/conv-transpose';
|
||||
import {gemm, parseGemmAttributes} from './ops/gemm';
|
||||
import {matMul} from './ops/matmul';
|
||||
import * as pool from './ops/pool';
|
||||
|
|
@ -33,6 +34,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map<string, OperatorImplementation> = new
|
|||
['Clip', [unaryOps.clip]],
|
||||
['Concat', [concat, parseConcatAttributes]],
|
||||
['Conv', [conv, parseConvAttributes]],
|
||||
['ConvTranspose', [convTranspose, parseConvTransposeAttributes]],
|
||||
['Cos', [unaryOps.cos]],
|
||||
['Cosh', [unaryOps.cosh]],
|
||||
['Div', [binaryOps.div]],
|
||||
|
|
|
|||
385
js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_webgpu.ts
vendored
Normal file
385
js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_webgpu.ts
vendored
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2021 Google LLC. All Rights Reserved.
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
* =============================================================================
|
||||
*/
|
||||
|
||||
// sampled from [@tensorflow/tfjs] tfjs-backend-webgpu/src/conv_backprop_webgpu.ts
|
||||
|
||||
import {LOG_DEBUG} from '../../../log';
|
||||
import {TensorView} from '../../../tensor';
|
||||
import {ShapeUtil} from '../../../util';
|
||||
import {GpuDataType, ProgramInfo, ProgramMetadata} from '../../types';
|
||||
import {createIndicesHelper, ShaderHelper} from '../common';
|
||||
import {ConvTransposeAttributes} from '../conv-transpose';
|
||||
|
||||
const createConvTranspose2DOpProgramShaderSource =
|
||||
(shaderHelper: ShaderHelper, inputs: readonly TensorView[], attributes: ConvTransposeAttributes,
|
||||
outputShape: readonly number[], hasBias: boolean, elementsPerThread: readonly number[]): string => {
|
||||
const isChannelsLast = attributes.format === 'NHWC';
|
||||
const rowDim = isChannelsLast ? 1 : 2;
|
||||
const colDim = isChannelsLast ? 2 : 3;
|
||||
const channelDim = isChannelsLast ? 3 : 1;
|
||||
const outputSize = ShapeUtil.size(outputShape);
|
||||
const outChannels = outputShape[isChannelsLast ? 3 : 1];
|
||||
const inChannels = inputs[0].dims[isChannelsLast ? 3 : 1];
|
||||
const isVec4 = inChannels % 4 === 0 && outChannels % 4 === 0;
|
||||
const workPerThread = isVec4 ? 2 : 1;
|
||||
|
||||
const innerElementSize = isVec4 ? (isChannelsLast && inChannels % 4 !== 0 ? 3 : 4) : elementsPerThread[0];
|
||||
|
||||
const declareInputs = [
|
||||
`@group(0) @binding(0) var<storage, read> Dy: array<${
|
||||
isVec4 && innerElementSize === 4 ? 'vec4<f32>' : 'f32'}>;`,
|
||||
`@group(0) @binding(1) var<storage, read> W: array<${isVec4 ? 'vec4<f32>' : 'f32'}>;`
|
||||
];
|
||||
let declareFunctions = `
|
||||
fn setOutputAtIndex(flatIndex : u32, value : ${isVec4 ? 'vec4<f32>' : 'f32'}) {
|
||||
result[flatIndex] = ${isVec4 ? 'vec4<f32>' : 'f32'}(value);
|
||||
}`;
|
||||
if (hasBias) {
|
||||
declareInputs.push(`@group(0) @binding(2) var<storage, read> bias: array<${isVec4 ? 'vec4<f32>' : 'f32'}>;`);
|
||||
declareFunctions += `
|
||||
fn getBiasByOutputCoords(coords : vec4<u32>) -> ${isVec4 ? 'vec4<f32>' : 'f32'} {
|
||||
return bias[coords.${isChannelsLast ? 'w' : 'y'}${isVec4 ? '/ 4' : ''}];
|
||||
}`;
|
||||
}
|
||||
const wIndicesHelper = createIndicesHelper('W', inputs[1].dims);
|
||||
const dyIndicesHelper = createIndicesHelper('Dy', inputs[0].dims);
|
||||
const outputIndicesHelper = createIndicesHelper('result', outputShape);
|
||||
const codeSnippet4 = `{
|
||||
let batch: u32 = global_id.z / outShape[1];
|
||||
let r = global_id.z % outShape[1];
|
||||
let c = global_id.y * ${workPerThread};
|
||||
let d1: u32 = global_id.x * 4;
|
||||
|
||||
let dyCorner = vec2<i32>(i32(r), i32(c)) - vec2<i32>(pads);
|
||||
|
||||
// Convolve dy(?, ?, d2) with w(:, :, d1, d2) to compute dx(xR, xC, d1).
|
||||
// ? = to be determined. : = across all values in that axis.
|
||||
var dotProd: array<vec4<f32>, ${workPerThread}>;
|
||||
for (var i = 0; i < ${workPerThread}; i++) {
|
||||
dotProd[i] = vec4<f32>(0.0);
|
||||
}
|
||||
for (var wR: u32 = 0; wR < filterDims[0]; wR = wR + 1) {
|
||||
var dyR = f32(dyCorner.x + wR) / f32(strides.x);
|
||||
let wRPerm: u32= filterDims[0] - 1 - wR;
|
||||
if (dyR < 0.0 || dyR >= f32(outBackprop[1]) ||
|
||||
fract(dyR) > 0.0) {
|
||||
continue;
|
||||
}
|
||||
let idyR: u32 = u32(dyR);
|
||||
|
||||
for (var wC: u32 = 0; wC < filterDims[1]; wC = wC + 1) {
|
||||
let dyC = f32(dyCorner.y + wC) / f32(strides.y);
|
||||
let dyC2 = f32(dyCorner.y + 1 + wC) / f32(strides.y);
|
||||
let wCPerm: u32 = filterDims[1] - 1 - wC;
|
||||
var bDyCVal = true;
|
||||
var bDyCVal2 = true;
|
||||
if (dyC < 0.0 || dyC >= f32(outBackprop[2]) ||
|
||||
fract(dyC) > 0.0) {
|
||||
bDyCVal = false;
|
||||
}
|
||||
if (dyC2 < 0.0 || dyC2 >= f32(outBackprop[2]) ||
|
||||
fract(dyC2) > 0.0) {
|
||||
bDyCVal2 = false;
|
||||
}
|
||||
|
||||
let idyC: u32 = u32(dyC);
|
||||
let idyC2: u32 = u32(dyC2);
|
||||
if (bDyCVal && bDyCVal2) {
|
||||
let d2Length = outBackprop[3];
|
||||
for (var d2 :u32 = 0; d2 < d2Length; d2 = d2 + 4) {
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices0',
|
||||
[
|
||||
'd2', 'd1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices1',
|
||||
[
|
||||
'd2', 'd1+1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices2',
|
||||
[
|
||||
'd2', 'd1+2', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices3',
|
||||
[
|
||||
'd2', 'd1+3', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
let wValue0 = W[${wIndicesHelper.i2oExpression('wIndices0')}];
|
||||
let wValue1 = W[${wIndicesHelper.i2oExpression('wIndices1')}];
|
||||
let wValue2 = W[${wIndicesHelper.i2oExpression('wIndices2')}];
|
||||
let wValue3 = W[${wIndicesHelper.i2oExpression('wIndices3')}];
|
||||
${
|
||||
dyIndicesHelper.indicesVariableDeclaration(
|
||||
'dyIndices',
|
||||
isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] :
|
||||
[
|
||||
'batch', 'd2', 'idyR', 'idyC'
|
||||
])};
|
||||
var xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}];
|
||||
let tmpval = vec4<f32>(xValue * wValue0,
|
||||
xValue * wValue1,
|
||||
xValue * wValue2,
|
||||
xValue * wValue3);
|
||||
dotProd[0] = dotProd[0] + tmpval;
|
||||
|
||||
${
|
||||
dyIndicesHelper.indicesVariableDeclaration(
|
||||
'dyIndices2',
|
||||
isChannelsLast ? ['batch', 'idyR', 'idyC2', 'd2'] :
|
||||
[
|
||||
'batch', 'd2', 'idyR', 'idyC2'
|
||||
])};
|
||||
xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}];
|
||||
|
||||
dotProd[1] = dotProd[1] + vec4<f32>(xValue * wValue0,
|
||||
xValue * wValue1,
|
||||
xValue * wValue2,
|
||||
xValue * wValue3);
|
||||
}
|
||||
} else if (bDyCVal) {
|
||||
let d2Length = outBackprop[3];
|
||||
for (var d2: u32 = 0; d2 < d2Length; d2 = d2 + 4) {
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices0',
|
||||
[
|
||||
'd2', 'd1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices1',
|
||||
[
|
||||
'd2', 'd1+1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices2',
|
||||
[
|
||||
'd2', 'd1+2', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices3',
|
||||
[
|
||||
'd2', 'd1+3', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
let wValue0 = W[${wIndicesHelper.i2oExpression('wIndices0')}];
|
||||
let wValue1 = W[${wIndicesHelper.i2oExpression('wIndices1')}];
|
||||
let wValue2 = W[${wIndicesHelper.i2oExpression('wIndices2')}];
|
||||
let wValue3 = W[${wIndicesHelper.i2oExpression('wIndices3')}];
|
||||
${
|
||||
dyIndicesHelper.indicesVariableDeclaration(
|
||||
'dyIndices',
|
||||
isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] :
|
||||
[
|
||||
'batch', 'd2', 'idyR', 'idyC'
|
||||
])};
|
||||
var xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}];
|
||||
let tmpval = vec4<f32>(xValue * wValue0,
|
||||
xValue * wValue1,
|
||||
xValue * wValue2,
|
||||
xValue * wValue3);
|
||||
dotProd[0] = dotProd[0] + tmpval;
|
||||
}
|
||||
} else if (bDyCVal2) {
|
||||
let d2Length = outBackprop[3];
|
||||
for (var d2: u32 = 0; d2 < d2Length; d2 = d2 + 4) {
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices0',
|
||||
[
|
||||
'd2', 'd1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices1',
|
||||
[
|
||||
'd2', 'd1+1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices2',
|
||||
[
|
||||
'd2', 'd1+2', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration(
|
||||
'wIndices3',
|
||||
[
|
||||
'd2', 'd1+3', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
let wValue0 = W[${wIndicesHelper.i2oExpression('wIndices0')}];
|
||||
let wValue1 = W[${wIndicesHelper.i2oExpression('wIndices1')}];
|
||||
let wValue2 = W[${wIndicesHelper.i2oExpression('wIndices2')}];
|
||||
let wValue3 = W[${wIndicesHelper.i2oExpression('wIndices3')}];
|
||||
${
|
||||
dyIndicesHelper.indicesVariableDeclaration(
|
||||
'dyIndices',
|
||||
isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] :
|
||||
[
|
||||
'batch', 'd2', 'idyR', 'idyC'
|
||||
])};
|
||||
var xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}];
|
||||
let tmpval = vec4<f32>(xValue * wValue0,
|
||||
xValue * wValue1,
|
||||
xValue * wValue2,
|
||||
xValue * wValue3);
|
||||
dotProd[1] = dotProd[1] + tmpval;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (var i: u32 = 0; i < ${workPerThread}; i = i + 1) {
|
||||
${
|
||||
outputIndicesHelper.indicesVariableDeclaration('outputIndices', [
|
||||
'batch', 'r', 'c+i', 'd1'
|
||||
])};
|
||||
result[${outputIndicesHelper.i2oExpression('outputIndices')}] = dotProd[i];
|
||||
}
|
||||
}`;
|
||||
const codeSnippet = `
|
||||
${outputIndicesHelper.o2iCall('global_idx', 'outputIndices')}
|
||||
let batch = outputIndices[0];
|
||||
let d1 = outputIndices[${channelDim}];
|
||||
let dyCorner = vec2<i32>(i32(outputIndices[${rowDim}]), i32(outputIndices[${colDim}])) - pads;
|
||||
let dyRCorner = dyCorner.x;
|
||||
let dyCCorner = dyCorner.y;
|
||||
// Convolve dy(?, ?, d2) with w(:, :, d1, d2) to compute dx(xR, xC, d1).
|
||||
// ? = to be determined. : = across all values in that axis.
|
||||
var dotProd = 0.0;
|
||||
for (var wR: u32 = 0; wR < effectiveFilterDims.x; wR = wR + 1) {
|
||||
if (wR % dilations.x != 0) {
|
||||
continue;
|
||||
}
|
||||
let dyR = (f32(dyRCorner) + f32(wR)) / f32(strides[0]);
|
||||
let wRPerm = filterDims.x - 1 - wR / dilations.x;
|
||||
if (dyR < 0.0 || dyR >= f32(outBackprop[1]) || fract(dyR) > 0.0 ||
|
||||
wRPerm < 0) {
|
||||
continue;
|
||||
}
|
||||
let idyR: u32 = u32(dyR);
|
||||
|
||||
for (var wC: u32 = 0; wC < effectiveFilterDims.y; wC = wC + 1) {
|
||||
if (wC % dilations.y != 0) {
|
||||
continue;
|
||||
}
|
||||
let dyC = (f32(dyCCorner) + f32(wC)) / f32(strides.y);
|
||||
let wCPerm = filterDims.y - 1 - wC / dilations.y;
|
||||
if (dyC < 0.0 || dyC >= f32(outBackprop[2]) ||
|
||||
fract(dyC) > 0.0 || wCPerm < 0) {
|
||||
continue;
|
||||
}
|
||||
let idyC: u32 = u32(dyC);
|
||||
|
||||
for (var d2: u32 = 0; d2 < outBackprop[3]; d2 = d2 + 1) {
|
||||
${
|
||||
dyIndicesHelper.indicesVariableDeclaration(
|
||||
'dyIndices',
|
||||
isChannelsLast ? ['batch', 'idyR', 'idyC', 'd2'] :
|
||||
[
|
||||
'batch', 'd2', 'idyR', 'idyC'
|
||||
])};
|
||||
let xValue = Dy[${dyIndicesHelper.i2oExpression('dyIndices')}];
|
||||
${
|
||||
wIndicesHelper.indicesVariableDeclaration('wIndices', [
|
||||
'd2', 'd1', 'wRPerm', 'wCPerm'
|
||||
])};
|
||||
|
||||
let wValue = W[${wIndicesHelper.i2oExpression('wIndices')}];
|
||||
dotProd = dotProd + xValue * wValue;
|
||||
}
|
||||
}
|
||||
}
|
||||
result[global_idx] = dotProd;
|
||||
`;
|
||||
|
||||
return `
|
||||
${wIndicesHelper.i2oImpl}
|
||||
${dyIndicesHelper.i2oImpl}
|
||||
${outputIndicesHelper.o2iImpl}
|
||||
${declareFunctions}
|
||||
${declareInputs.join('\n')}
|
||||
@group(0) @binding(${declareInputs.length}) var<storage, read_write> result: array<${isVec4 ? 'vec4<f32>' : 'f32'}>;
|
||||
const outShape : vec4<u32> = vec4<u32>(${outputShape.join(',')});
|
||||
const outBackprop : vec4<u32> = vec4<u32>(${inputs[0].dims.join(',')});
|
||||
const strides : vec2<u32> = vec2<u32>(${attributes.strides[0]}, ${attributes.strides[1]});
|
||||
const filterDims : vec2<u32> = vec2<u32>(${attributes.kernelShape[isChannelsLast ? 1 : 2]}, ${
|
||||
attributes.kernelShape[isChannelsLast ? 2 : 3]});
|
||||
const dilations : vec2<u32> = vec2<u32>(${attributes.dilations[0]}, ${attributes.dilations[1]});
|
||||
const effectiveFilterDims : vec2<u32> = filterDims + vec2<u32>(
|
||||
${
|
||||
attributes.dilations[0] <= 1 ?
|
||||
0 :
|
||||
(attributes.kernelShape[isChannelsLast ? 1 : 2] - 1) * (attributes.dilations[0] - 1)},
|
||||
${
|
||||
attributes.dilations[1] <= 1 ?
|
||||
0 :
|
||||
(attributes.kernelShape[isChannelsLast ? 2 : 3] - 1) * (attributes.dilations[1] - 1)});
|
||||
const pads : vec2<i32> = vec2<i32>(i32(effectiveFilterDims[0]) - 1 - (${attributes.pads[0] + attributes.pads[2]})/2,
|
||||
i32(effectiveFilterDims[1]) - 1 - (${attributes.pads[1] + attributes.pads[3]})/2);
|
||||
${shaderHelper.mainStart()}
|
||||
${outputIndicesHelper.indicesVariableDeclaration('outputIndices')}
|
||||
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)};
|
||||
${isVec4 ? codeSnippet4 : codeSnippet}}`;
|
||||
};
|
||||
|
||||
export const createConvTranspose2DProgramInfo =
|
||||
(inputs: readonly TensorView[], metadata: ProgramMetadata, attributes: ConvTransposeAttributes,
|
||||
squeezeOutputShapeFunction?: (shape: readonly number[]) => number[]): ProgramInfo => {
|
||||
const hasBias = inputs.length > 2;
|
||||
const isChannelsLast = attributes.format === 'NHWC';
|
||||
const outputShape = attributes.outputShape;
|
||||
const batchSize = outputShape[0];
|
||||
const outWidth = outputShape[isChannelsLast ? 1 : 2];
|
||||
const outHeight = outputShape[isChannelsLast ? 2 : 3];
|
||||
const outChannels = outputShape[isChannelsLast ? 3 : 1];
|
||||
const inChannels = inputs[0].dims[isChannelsLast ? 3 : 1];
|
||||
const isVec4 = inChannels % 4 === 0 && outChannels % 4 === 0;
|
||||
|
||||
const dispatchX = isChannelsLast ? outChannels : outWidth * outHeight;
|
||||
const dispatchY = isChannelsLast ? outWidth * outHeight : outChannels;
|
||||
const workGroupSize: [number, number, number] =
|
||||
isVec4 ? [8, 8, 1] : [dispatchX <= 4 ? 4 : 16, dispatchX > 4 && dispatchY <= 4 ? 4 : 16, 1];
|
||||
const elementsPerThread =
|
||||
isVec4 ? [4, 4, 1] : [dispatchX <= 4 ? 1 : 2, dispatchX > 4 && dispatchY <= 4 ? 1 : 2, 1];
|
||||
const dispatch = [
|
||||
Math.ceil(dispatchX / workGroupSize[0] / elementsPerThread[0]),
|
||||
Math.ceil(dispatchY / workGroupSize[1] / elementsPerThread[1]),
|
||||
Math.ceil(batchSize / workGroupSize[2] / elementsPerThread[1])
|
||||
];
|
||||
LOG_DEBUG('verbose', () => `[conv2d_backprop_webgpu] dispatch = ${dispatch}`);
|
||||
|
||||
return {
|
||||
...metadata,
|
||||
outputs: [{
|
||||
dims: squeezeOutputShapeFunction ? squeezeOutputShapeFunction(outputShape) : outputShape,
|
||||
dataType: inputs[0].dataType,
|
||||
gpuDataType: GpuDataType.default
|
||||
}],
|
||||
dispatchGroup: () => ({x: dispatch[0], y: dispatch[1], z: dispatch[2]}),
|
||||
getShaderSource: (shaderHelper: ShaderHelper) => createConvTranspose2DOpProgramShaderSource(
|
||||
shaderHelper, inputs, attributes, outputShape, hasBias, elementsPerThread),
|
||||
};
|
||||
};
|
||||
287
js/web/lib/wasm/jsep/webgpu/ops/conv-transpose.ts
Normal file
287
js/web/lib/wasm/jsep/webgpu/ops/conv-transpose.ts
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
import {DataType} from '../../../wasm-common';
|
||||
import {TensorView} from '../../tensor';
|
||||
import {createAttributeWithCacheKey} from '../attribute-with-cache-key';
|
||||
import {ComputeContext, GpuDataType, ProgramInfoLoader, ProgramMetadata} from '../types';
|
||||
|
||||
import {createConvTranspose2DProgramInfo} from './3rd-party/conv_backprop_webgpu';
|
||||
import {ConvAttributes} from './conv';
|
||||
import {parseInternalActivationAttributes} from './fuse-utils';
|
||||
|
||||
const computeTotalPad =
|
||||
(inDim: number, stride: number, adj: number, kernel: number, dilation: number, outSize: number) =>
|
||||
(inDim - 1) * stride + adj + (kernel - 1) * dilation + 1 - outSize;
|
||||
|
||||
const distributePadding = (totalPad: number, autoPad: string, pads: number[], head: number, tail: number) => {
|
||||
const smallPad = Math.floor(totalPad / 2);
|
||||
if (autoPad === 'SAME_UPPER') {
|
||||
pads[head] = smallPad;
|
||||
pads[tail] = totalPad - smallPad;
|
||||
} else if (autoPad === 'SAME_LOWER') {
|
||||
pads[head] = totalPad - smallPad;
|
||||
pads[tail] = smallPad;
|
||||
}
|
||||
};
|
||||
|
||||
const calculateOutputShapeAndPads =
|
||||
(inputShape: readonly number[], kernelShape: readonly number[], dilations: readonly number[], autoPad: string,
|
||||
group: number, pads: number[], strides: readonly number[], isChannelLast: boolean, outputPadding: number[],
|
||||
outputShape: number[]) => {
|
||||
const spatialRank = inputShape.length - 2;
|
||||
const updateOutputShape = outputShape.length === 0;
|
||||
if (outputPadding.length === 0) {
|
||||
for (let i = 0; i < spatialRank; ++i) {
|
||||
outputPadding.push(0);
|
||||
}
|
||||
}
|
||||
const batchSize = inputShape[0];
|
||||
const outChannels = kernelShape[isChannelLast ? 3 : 1] * group;
|
||||
for (let i = 0, j = inputShape.length - spatialRank - (isChannelLast ? 1 : 0); i < spatialRank; ++i, ++j) {
|
||||
const inSize = inputShape[j];
|
||||
const outSize = updateOutputShape ? inSize * strides[i] : outputShape[i];
|
||||
const totalPad = computeTotalPad(inSize, strides[i], pads[i], kernelShape[j], dilations[i], outSize);
|
||||
distributePadding(totalPad, autoPad, pads, i, i + spatialRank);
|
||||
if (updateOutputShape) {
|
||||
outputShape.push(
|
||||
strides[i] * (inSize - 1) + outputPadding[i] + (kernelShape[j] - 1) * dilations[i] + 1 - pads[i] -
|
||||
pads[i + spatialRank]);
|
||||
}
|
||||
}
|
||||
outputShape.splice(0, 0, batchSize);
|
||||
outputShape.splice(isChannelLast ? 3 : 1, 0, outChannels);
|
||||
};
|
||||
|
||||
export interface ConvTransposeAttributes extends ConvAttributes {
|
||||
readonly outputPadding: readonly number[];
|
||||
readonly outputShape: readonly number[];
|
||||
}
|
||||
|
||||
|
||||
const getAdjustedConvTransposeAttributes =
|
||||
<T extends ConvTransposeAttributes>(attributes: T, inputs: readonly TensorView[]): T => {
|
||||
const kernelShape = attributes.kernelShape.slice();
|
||||
// if kernelShape is not specified in the attributes of this op, infer it from the weight tensor dims
|
||||
if (attributes.kernelShape.length === 0 || attributes.kernelShape.reduce((a, b) => a * b, 0) === 0) {
|
||||
kernelShape.length = 0;
|
||||
for (let i = 2; i < inputs[1].dims.length; ++i) {
|
||||
kernelShape.push(inputs[1].dims[i]);
|
||||
}
|
||||
}
|
||||
const isChannelsLast = attributes.format === 'NHWC';
|
||||
kernelShape.splice(0, 0, inputs[1].dims[0]);
|
||||
kernelShape.splice(isChannelsLast ? 3 : 1, 0, inputs[1].dims[1]);
|
||||
|
||||
const pads = attributes.pads.slice();
|
||||
const outputShape = attributes.outputShape.slice();
|
||||
const outputPadding = attributes.outputPadding.slice();
|
||||
const inputShape = inputs[0].dims;
|
||||
let dilations = attributes.dilations.slice();
|
||||
if (dilations.reduce((a, b) => a + b, 0) === 0) {
|
||||
const spatialRank = inputs[0].dims.length - 2;
|
||||
dilations = new Array(spatialRank).fill(1);
|
||||
}
|
||||
let strides = attributes.strides.slice();
|
||||
if (strides.reduce((a, b) => a + b, 0) === 0) {
|
||||
const spatialRank = inputs[0].dims.length - 2;
|
||||
strides = new Array(spatialRank).fill(1);
|
||||
}
|
||||
// If outputShape is not specified in the attributes of this op, infer it from the parameters
|
||||
// Similarly, automatically infer pads if not specified
|
||||
calculateOutputShapeAndPads(
|
||||
inputShape, kernelShape, dilations, attributes.autoPad, attributes.group, pads, strides, isChannelsLast,
|
||||
outputPadding, outputShape);
|
||||
|
||||
// always return a new object so does not modify the original attributes
|
||||
const newAttributes: T = Object.assign({}, attributes);
|
||||
Object.assign(
|
||||
newAttributes,
|
||||
{kernelShape, pads, outputPadding, outputShape, dilations, strides, cacheKey: attributes.cacheKey});
|
||||
return newAttributes;
|
||||
};
|
||||
|
||||
export const parseConvTransposeAttributes = (attributes: Record<string, unknown>): ConvTransposeAttributes => {
|
||||
const activationAttributes = parseInternalActivationAttributes(attributes);
|
||||
// TODO : Make this generic enough to compute default attributes for multi-dimensional conv
|
||||
const format = attributes.format as 'NHWC' | 'NCHW';
|
||||
const autoPad =
|
||||
['NOTSET', 'VALID', 'SAME_UPPER',
|
||||
'SAME_LOWER'][typeof attributes.autoPad == 'undefined' ? 0 : attributes.autoPad as number];
|
||||
const dilations = attributes.dilations as [number, number];
|
||||
const group = attributes.group as number;
|
||||
const kernelShape = attributes.kernelShape as [number, number];
|
||||
const pads = attributes.pads as [number, number, number, number];
|
||||
const strides = attributes.strides as [number, number];
|
||||
const wIsConst = (attributes.wIsConst as () => boolean)();
|
||||
const outputPadding = attributes.outputPadding as [number, number, number, number];
|
||||
const outputShape = attributes.outputShape as [number, number];
|
||||
return createAttributeWithCacheKey({
|
||||
autoPad,
|
||||
format,
|
||||
dilations,
|
||||
group,
|
||||
kernelShape,
|
||||
outputPadding,
|
||||
outputShape,
|
||||
pads,
|
||||
strides,
|
||||
wIsConst,
|
||||
...activationAttributes
|
||||
});
|
||||
};
|
||||
|
||||
const validateInputs = (inputs: readonly TensorView[], attributes: ConvTransposeAttributes): void => {
|
||||
// Refer to the below link for all input checks
|
||||
// https://github.com/onnx/onnx/blob/main/docs/Operators.md#ConvTranspose
|
||||
if (!inputs || (inputs.length !== 2 && inputs.length !== 3)) {
|
||||
throw new Error('Conv requires 2 or 3 inputs');
|
||||
}
|
||||
|
||||
// TODO : Need to add support for multi-dimensional conv
|
||||
if (inputs[0].dims.length !== 4 && inputs[0].dims.length !== 3) {
|
||||
throw new Error('currently only support 2-dimensional conv');
|
||||
}
|
||||
|
||||
if (inputs[0].dims.length !== inputs[1].dims.length) {
|
||||
throw new Error('filter does not have same dimension as input');
|
||||
}
|
||||
|
||||
// FILTER_IN_CHANNEL should be equal to DATA_CHANNEL
|
||||
const dataChannel = inputs[0].dims[attributes.format === 'NHWC' ? inputs[0].dims.length - 1 : 1];
|
||||
const filterInChannel = inputs[1].dims[0];
|
||||
if (dataChannel !== filterInChannel) {
|
||||
throw new Error('FILTER_IN_CHANNEL should be equal to DATA_CHANNEL');
|
||||
}
|
||||
|
||||
const featureMaps = inputs[1].dims[1] * attributes.group;
|
||||
|
||||
// if bias is provided it should be 1D and the number of elements should be equal to the number of feature maps
|
||||
if (inputs.length === 3 && (inputs[2].dims.length !== 1 || inputs[2].dims[0] !== featureMaps)) {
|
||||
throw new Error('invalid bias');
|
||||
}
|
||||
|
||||
const spatialRank = inputs[0].dims.length - 2;
|
||||
const dilationsSet = attributes.dilations.reduce((a, b) => a + b, 0) > 0;
|
||||
// wrong dilations dimension
|
||||
if (dilationsSet && attributes.dilations.length !== spatialRank) {
|
||||
throw new Error(`dilations should be ${spatialRank}D`);
|
||||
}
|
||||
|
||||
const stridesSet = attributes.strides.reduce((a, b) => a + b, 0) > 0;
|
||||
// Wrong strides dimension
|
||||
if (stridesSet && attributes.strides.length !== spatialRank) {
|
||||
throw new Error(`strides should be ${spatialRank}D`);
|
||||
}
|
||||
|
||||
// Wrong pads dimension
|
||||
const padsSet = attributes.pads.reduce((a, b) => a + b, 0) > 0;
|
||||
if (padsSet && attributes.pads.length !== spatialRank * 2) {
|
||||
throw new Error(`pads should be ${spatialRank * 2}D`);
|
||||
}
|
||||
|
||||
// Wrong output padding dimension
|
||||
if (attributes.outputPadding.length !== spatialRank && attributes.outputPadding.length !== 0) {
|
||||
throw new Error(`output_padding should be ${spatialRank}D`);
|
||||
}
|
||||
|
||||
// if kernelShape is specified, it's data length must be 2 less than dims length of the weights tensor
|
||||
// (the first 2 dims are batch_size and channels)
|
||||
const kernelShapeSet = attributes.kernelShape.reduce((a, b) => a + b, 0) > 0;
|
||||
if (kernelShapeSet && attributes.kernelShape.length !== 0 &&
|
||||
attributes.kernelShape.length !== inputs[1].dims.length - 2) {
|
||||
throw new Error('invalid kernel shape');
|
||||
}
|
||||
|
||||
// as with kernelShape, must have same number of spatial dims as input
|
||||
if (attributes.outputShape.length !== 0 && attributes.outputShape.length !== inputs[0].dims.length - 2) {
|
||||
throw new Error('invalid output shape');
|
||||
}
|
||||
|
||||
// TODO : Need to add support for float64
|
||||
if (inputs[0].dataType !== DataType.float || inputs[1].dataType !== DataType.float) {
|
||||
throw new Error('ConvTranspose input(X,W) should be float tensor');
|
||||
}
|
||||
|
||||
if (inputs.length === 3 && inputs[2].dataType !== DataType.float) {
|
||||
throw new Error('ConvTranspose input(bias) should be float tensor');
|
||||
}
|
||||
};
|
||||
|
||||
const createConvTranspose2DProgramMetadata = (hasBias: boolean, cacheHint: string): ProgramMetadata => ({
|
||||
name: 'ConvTranspose2D',
|
||||
inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
|
||||
[GpuDataType.default, GpuDataType.default],
|
||||
cacheHint
|
||||
});
|
||||
|
||||
const createConvTranspose2DProgramInfoLoader =
|
||||
(inputs: readonly TensorView[], attributes: ConvTransposeAttributes,
|
||||
squeezeOutputShapeFunction?: (shape: readonly number[]) => number[]): ProgramInfoLoader => {
|
||||
const hasBias = inputs.length === 3;
|
||||
const metadata = createConvTranspose2DProgramMetadata(hasBias, attributes.cacheKey);
|
||||
return {
|
||||
...metadata,
|
||||
get: () => createConvTranspose2DProgramInfo(inputs, metadata, attributes, squeezeOutputShapeFunction)
|
||||
};
|
||||
};
|
||||
|
||||
const convTranspose2d =
|
||||
(context: ComputeContext, inputs: readonly TensorView[], attributes: ConvTransposeAttributes): void => {
|
||||
const adjustedAttributes = getAdjustedConvTransposeAttributes(attributes, inputs);
|
||||
|
||||
context.compute(createConvTranspose2DProgramInfoLoader(inputs, adjustedAttributes));
|
||||
};
|
||||
const convTranspose1d = (context: ComputeContext, attributes: ConvTransposeAttributes): void => {
|
||||
// extend the input to 2D by adding H dimension
|
||||
const isChannelLast = attributes.format === 'NHWC';
|
||||
|
||||
const inputs = [
|
||||
context.inputs[0].reshape(
|
||||
isChannelLast ?
|
||||
// [N, W, C] -> [N, H=1, W, C]
|
||||
[context.inputs[0].dims[0], 1, context.inputs[0].dims[1], context.inputs[0].dims[2]] :
|
||||
// [N, C, W] -> [N, C, H=1, W]
|
||||
[context.inputs[0].dims[0], context.inputs[0].dims[1], 1, context.inputs[0].dims[2]]),
|
||||
//[FILTER_OUT_CHANNEL, FILTER_IN_CHANNEL, kW] -> [FILTER_OUT_CHANNEL, FILTER_IN_CHANNEL, kH=1, kW]
|
||||
context.inputs[1].reshape([context.inputs[1].dims[0], context.inputs[1].dims[1], 1, context.inputs[1].dims[2]])
|
||||
];
|
||||
if (inputs.length === 3) {
|
||||
inputs.push(context.inputs[2]);
|
||||
}
|
||||
let kernelShape = attributes.kernelShape;
|
||||
if (kernelShape.length === 0 || kernelShape[0] === 0) {
|
||||
kernelShape = [context.inputs[1].dims[2]];
|
||||
}
|
||||
let dilations = attributes.dilations;
|
||||
if (dilations.length === 0 || dilations[0] === 0) {
|
||||
dilations = [1];
|
||||
}
|
||||
let strides = attributes.strides;
|
||||
if (strides.length === 0 || strides[0] === 0) {
|
||||
strides = [1];
|
||||
}
|
||||
let pads = attributes.pads;
|
||||
if (pads.length === 0) {
|
||||
pads = [0, 0];
|
||||
}
|
||||
pads = [0, pads[0], 0, pads[1]];
|
||||
strides = [1].concat(strides);
|
||||
dilations = [1].concat(dilations);
|
||||
kernelShape = [1].concat(kernelShape);
|
||||
const adjustedAttributes =
|
||||
getAdjustedConvTransposeAttributes({...attributes, pads, strides, dilations, kernelShape}, inputs);
|
||||
context.compute(createConvTranspose2DProgramInfoLoader(
|
||||
inputs, adjustedAttributes,
|
||||
outputShape => isChannelLast ? [outputShape[0], outputShape[2], outputShape[3]] :
|
||||
[outputShape[0], outputShape[1], outputShape[3]]));
|
||||
};
|
||||
|
||||
export const convTranspose = (context: ComputeContext, attributes: ConvTransposeAttributes): void => {
|
||||
validateInputs(context.inputs, attributes);
|
||||
if (context.inputs[0].dims.length === 3) {
|
||||
convTranspose1d(context, attributes);
|
||||
} else {
|
||||
convTranspose2d(context, context.inputs, attributes);
|
||||
}
|
||||
};
|
||||
|
|
@ -450,16 +450,16 @@
|
|||
"test_conv_with_strides_padding",
|
||||
// // "test_convinteger_with_padding",
|
||||
// // "test_convinteger_without_padding",
|
||||
// // "test_convtranspose_1d",
|
||||
"test_convtranspose_1d",
|
||||
// // "test_convtranspose_3d",
|
||||
// // "test_convtranspose_autopad_same",
|
||||
// // "test_convtranspose_dilations",
|
||||
// // "test_convtranspose_kernel_shape",
|
||||
// // "test_convtranspose_output_shape",
|
||||
// // "test_convtranspose_pad",
|
||||
// // "test_convtranspose_pads",
|
||||
// // "test_convtranspose_with_kernel",
|
||||
// // "test_convtranspose",
|
||||
"test_convtranspose_autopad_same",
|
||||
"test_convtranspose_dilations",
|
||||
"test_convtranspose_kernel_shape",
|
||||
"opset{9,17}/test_convtranspose_output_shape",
|
||||
"test_convtranspose_pad",
|
||||
"test_convtranspose_pads",
|
||||
"test_convtranspose_with_kernel",
|
||||
"test_convtranspose",
|
||||
"test_cos_example",
|
||||
"test_cos",
|
||||
"test_cosh_example",
|
||||
|
|
|
|||
|
|
@ -205,6 +205,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomai
|
|||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Transpose);
|
||||
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, Conv);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, ConvTranspose);
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, 11, float, MaxPool);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 12, float, MaxPool);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, AveragePool);
|
||||
|
|
@ -213,6 +214,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHW
|
|||
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, Conv);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Conv);
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, ConvTranspose);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, ConvTranspose);
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 7, 8, float, Gemm);
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 10, float, Gemm);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Gemm);
|
||||
|
|
@ -383,6 +386,7 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Transpose)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, Conv)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, ConvTranspose)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, 11, float, MaxPool)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 12, float, MaxPool)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 11, float, AveragePool)>,
|
||||
|
|
@ -391,6 +395,8 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
|
|||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, Conv)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Conv)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 10, float, ConvTranspose)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, ConvTranspose)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 7, 8, float, Gemm)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 10, float, Gemm)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, float, Gemm)>,
|
||||
|
|
|
|||
39
onnxruntime/core/providers/js/operators/conv_transpose.cc
Normal file
39
onnxruntime/core/providers/js/operators/conv_transpose.cc
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/providers/js/js_kernel.h"
|
||||
#include "core/providers/cpu/nn/conv_attributes.h"
|
||||
|
||||
#include "conv_transpose.h"
|
||||
namespace onnxruntime {
|
||||
namespace js {
|
||||
#define REGISTER_KERNEL_TYPED(T) \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
ConvTranspose, \
|
||||
kMSInternalNHWCDomain, \
|
||||
11, \
|
||||
T, \
|
||||
kJsExecutionProvider, \
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
|
||||
ConvTranspose<T, true>); \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
ConvTranspose, \
|
||||
kOnnxDomain, \
|
||||
11, \
|
||||
T, \
|
||||
kJsExecutionProvider, \
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
|
||||
ConvTranspose<T, false>); \
|
||||
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
|
||||
ConvTranspose, \
|
||||
kOnnxDomain, \
|
||||
1, 10, \
|
||||
T, \
|
||||
kJsExecutionProvider, \
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
|
||||
ConvTranspose<T, false>);
|
||||
|
||||
REGISTER_KERNEL_TYPED(float)
|
||||
|
||||
} // namespace js
|
||||
} // namespace onnxruntime
|
||||
117
onnxruntime/core/providers/js/operators/conv_transpose.h
Normal file
117
onnxruntime/core/providers/js/operators/conv_transpose.h
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include "core/common/gsl.h"
|
||||
#include "core/providers/cpu/nn/conv_transpose_attributes.h"
|
||||
#include "core/providers/js/js_kernel.h"
|
||||
namespace onnxruntime {
|
||||
namespace js {
|
||||
template <typename T, bool is_channels_last>
|
||||
class ConvTranspose : public JsKernel {
|
||||
public:
|
||||
ConvTranspose(const OpKernelInfo& info) : JsKernel(info), conv_transpose_attrs_(info), w_is_const_(false) {
|
||||
TensorShapeVector kernel_shape;
|
||||
if (conv_transpose_attrs_.kernel_shape_specified) {
|
||||
ORT_ENFORCE(info.GetAttrs("kernel_shape", kernel_shape).IsOK());
|
||||
}
|
||||
|
||||
int64_t channels_last = is_channels_last ? 1 : info.GetAttrOrDefault<int64_t>("channels_last", 0);
|
||||
|
||||
// currently only support Conv 1D/2D. TODO: support Conv3D and other
|
||||
if (conv_transpose_attrs_.dilations.size() == 1 ||
|
||||
(conv_transpose_attrs_.kernel_shape_specified && kernel_shape.size() == 1) ||
|
||||
conv_transpose_attrs_.strides.size() == 1) {
|
||||
JSEP_INIT_KERNEL_ATTRIBUTE(ConvTranspose, ({
|
||||
"format" : $8 ? "NHWC" : "NCHW",
|
||||
"autoPad" : $1,
|
||||
"dilations" : [$2],
|
||||
"group" : $3,
|
||||
"kernel_shape" : [$4],
|
||||
"pads" : [ $5, $6 ],
|
||||
"strides" : [$7],
|
||||
"wIsConst" : () JS_ARROW(!!HEAP8[$9]),
|
||||
"outputPadding" : $10 ? Array.from(HEAP32.subarray($11, $11 + $10)) : [],
|
||||
"outputShape" : $12 ? Array.from(HEAP32.subarray($12, $13 + $12)) : []
|
||||
}),
|
||||
static_cast<int32_t>(conv_transpose_attrs_.auto_pad),
|
||||
static_cast<int32_t>(conv_transpose_attrs_.dilations.size() > 0 ? conv_transpose_attrs_.dilations[0] : 0),
|
||||
static_cast<int32_t>(conv_transpose_attrs_.group),
|
||||
static_cast<int32_t>(conv_transpose_attrs_.kernel_shape_specified && kernel_shape.size() > 0) ? kernel_shape[0] : 0,
|
||||
static_cast<int32_t>(conv_transpose_attrs_.pads.size()),
|
||||
static_cast<int32_t>(conv_transpose_attrs_.pads.size() > 1) ? conv_transpose_attrs_.pads[1] : 0,
|
||||
static_cast<int32_t>(conv_transpose_attrs_.strides.size() > 0) ? conv_transpose_attrs_.strides[0] : 0,
|
||||
static_cast<int32_t>(channels_last),
|
||||
reinterpret_cast<int32_t>(&w_is_const_),
|
||||
gsl::narrow_cast<int32_t>(conv_transpose_attrs_.output_shape.size()),
|
||||
reinterpret_cast<int32_t>(conv_transpose_attrs_.output_padding.size() > 0 ? conv_transpose_attrs_.output_padding.data() : nullptr) >> 2,
|
||||
gsl::narrow_cast<int32_t>(conv_transpose_attrs_.output_shape.size()),
|
||||
reinterpret_cast<int32_t>(conv_transpose_attrs_.output_shape.size() > 0 ? conv_transpose_attrs_.output_shape.data() : nullptr) >> 2);
|
||||
} else {
|
||||
constexpr size_t pads_vec_size = 4;
|
||||
constexpr size_t strides_vec_size = 2;
|
||||
constexpr size_t dialations_vec_size = 2;
|
||||
constexpr size_t kernel_shape_vec_size = 2;
|
||||
// First set default values for pads, strides and dialations
|
||||
std::vector<int32_t> local_pads(pads_vec_size, 0);
|
||||
std::vector<int32_t> local_strides(strides_vec_size, 0);
|
||||
std::vector<int32_t> local_dilations(dialations_vec_size, 0);
|
||||
std::vector<int32_t> local_kernel_shape;
|
||||
std::vector<int32_t> local_output_shape(conv_transpose_attrs_.output_shape.begin(), conv_transpose_attrs_.output_shape.end());
|
||||
std::vector<int32_t> local_output_padding(conv_transpose_attrs_.output_padding.begin(), conv_transpose_attrs_.output_padding.end());
|
||||
if (conv_transpose_attrs_.kernel_shape_specified) {
|
||||
for (size_t i = 0; i < kernel_shape.size() && i < kernel_shape_vec_size; ++i) {
|
||||
local_kernel_shape.push_back(kernel_shape[i]);
|
||||
}
|
||||
} else {
|
||||
for (size_t i = 0; i < kernel_shape_vec_size; ++i) {
|
||||
local_kernel_shape.push_back(0);
|
||||
}
|
||||
}
|
||||
for (size_t i = 0; i < conv_transpose_attrs_.pads.size() && i < pads_vec_size; ++i) {
|
||||
local_pads[i] = conv_transpose_attrs_.pads[i];
|
||||
}
|
||||
for (size_t i = 0; i < conv_transpose_attrs_.dilations.size() && i < dialations_vec_size; ++i) {
|
||||
local_dilations[i] = conv_transpose_attrs_.dilations[i];
|
||||
}
|
||||
for (size_t i = 0; i < conv_transpose_attrs_.strides.size() && i < strides_vec_size; ++i) {
|
||||
local_strides[i] = conv_transpose_attrs_.strides[i];
|
||||
}
|
||||
LOGS_DEFAULT(VERBOSE) << "output_shape = " << conv_transpose_attrs_.output_shape << std::endl;
|
||||
LOGS_DEFAULT(VERBOSE) << "output_padding = " << conv_transpose_attrs_.output_padding << std::endl;
|
||||
JSEP_INIT_KERNEL_ATTRIBUTE(ConvTranspose, ({
|
||||
"format" : $7 ? "NHWC" : "NCHW",
|
||||
"autoPad" : $1,
|
||||
"dilations" : Array.from(HEAP32.subarray($2, $2 + /* dialations_vec_size */ 2)),
|
||||
"group" : $3,
|
||||
"kernelShape" : Array.from(HEAP32.subarray($4, $4 + /* kernel_shape_vec_size */ 2)),
|
||||
"pads" : Array.from(HEAP32.subarray($5, $5 + /* pads_vec_size */ 4)),
|
||||
"strides" : Array.from(HEAP32.subarray($6, $6 + /* strides_vec_size */ 2)),
|
||||
"wIsConst" : () JS_ARROW(!!HEAP8[$8]),
|
||||
"outputPadding" : ($9 > 0) ? Array.from(HEAP32.subarray($10, $10 + $9)) : [],
|
||||
"outputShape" : ($11 > 0) ? Array.from(HEAP32.subarray($12, $12 + $11)) : []
|
||||
}),
|
||||
static_cast<int32_t>(conv_transpose_attrs_.auto_pad),
|
||||
reinterpret_cast<int32_t>(local_dilations.data()) >> 2,
|
||||
static_cast<int32_t>(conv_transpose_attrs_.group),
|
||||
reinterpret_cast<int32_t>(local_kernel_shape.data()) >> 2,
|
||||
reinterpret_cast<int32_t>(local_pads.data()) >> 2,
|
||||
reinterpret_cast<int32_t>(local_strides.data()) >> 2,
|
||||
static_cast<int32_t>(channels_last),
|
||||
reinterpret_cast<int32_t>(&w_is_const_),
|
||||
gsl::narrow_cast<int32_t>(local_output_padding.size()),
|
||||
reinterpret_cast<int32_t>(local_output_padding.size() > 0 ? local_output_padding.data() : nullptr) >> 2,
|
||||
gsl::narrow_cast<int32_t>(local_output_shape.size()),
|
||||
reinterpret_cast<int32_t>(local_output_shape.size() > 0 ? local_output_shape.data() : nullptr) >> 2);
|
||||
}
|
||||
}
|
||||
|
||||
protected:
|
||||
ConvTransposeAttributes conv_transpose_attrs_;
|
||||
bool w_is_const_;
|
||||
};
|
||||
|
||||
} // namespace js
|
||||
} // namespace onnxruntime
|
||||
Loading…
Reference in a new issue