license import { backend_util , env , FusedDepthwiseConv2D , FusedDepthwiseConv2DAttrs , FusedDepthwiseConv2DInputs , KernelConfig , KernelFunc , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { DepthwiseConv2DProgram } from '../conv_gpu_depthwise' ; import { DepthwiseConvPacked2DProgram } from '../conv_packed_gpu_depthwise' ; import { mapActivationToShaderProgram } from '../kernel_utils/kernel_funcs_utils' ; export function fusedDepthwiseConv2D ( args ) { const { inputs , backend , attrs } = args ; const { x , filter , bias , preluActivationWeights } = inputs ; const { strides , pad , dilations , dimRoundingMode , activation , leakyreluAlpha } = attrs ; const intermediates = [ ] ; let $dilations = dilations ; if ( $dilations == null ) { $dilations = [ 1 , 1 ] ; } util . assert ( backend_util . eitherStridesOrDilationsAreOne ( strides , $dilations ) , ( ) => 'Error in depthwiseConv2d: Either strides or dilations must be ' + `1. Got strides ${ strides } and dilations '${ $dilations }'` ) ; const convInfo = backend_util . computeConv2DInfo ( x . shape as [ number , number , number , number ] , filter . shape as [ number , number , number , number ] , strides , $dilations , pad , dimRoundingMode , true ) ; const shouldPackDepthwiseConv = env ( ) . getBool ( 'WEBGL_PACK_DEPTHWISECONV' ) && convInfo . strideWidth <= 2 && convInfo . outChannels / convInfo . inChannels === 1 ; const fusedActivation = activation ? mapActivationToShaderProgram ( activation , shouldPackDepthwiseConv ) : null ; const programInputs = [ x , filter ] ; const hasBias = bias != null ; const hasPreluActivationWeights = preluActivationWeights != null ; const hasLeakyreluAlpha = activation === 'leakyrelu' ; if ( hasBias ) { programInputs . push ( bias ) ; } if ( hasPreluActivationWeights ) { programInputs . push ( preluActivationWeights ) ; } if ( hasLeakyreluAlpha ) { const $leakyreluAlpha = backend . makeTensorInfo ( [ ] , 'float32' , util . createScalarValue ( leakyreluAlpha as { } as 'float32' , 'float32' ) ) ; programInputs . push ( $leakyreluAlpha ) ; intermediates . push ( $leakyreluAlpha ) ; } let program ; if ( shouldPackDepthwiseConv ) { program = new DepthwiseConvPacked2DProgram ( convInfo , hasBias , fusedActivation , hasPreluActivationWeights , hasLeakyreluAlpha ) ; } else { program = new DepthwiseConv2DProgram ( convInfo , hasBias , fusedActivation , hasPreluActivationWeights , hasLeakyreluAlpha ) ; } const customValues = [ [ convInfo . padInfo . top , convInfo . padInfo . left ] , [ convInfo . strideHeight , convInfo . strideWidth ] , [ convInfo . dilationHeight , convInfo . dilationWidth ] , [ convInfo . inHeight , convInfo . inWidth ] ] ; const result = backend . runWebGLProgram ( program , programInputs , 'float32' , customValues ) ; intermediates . forEach ( t => backend . disposeIntermediateTensorInfo ( t ) ) ; return result ; } export const fusedDepthwiseConv2DConfig = { kernelName : FusedDepthwiseConv2D , backendName : 'webgl' , kernelFunc : fusedDepthwiseConv2D as { } as KernelFunc , } ;