def local_bnGrad_mkl(node):
if not mkl_available():
return
if not isinstance(node.op, mkl_bn.AbstractBatchNormalizationGrad):
return
if node.inputs[0].type.ndim != 4:
return
try:
x, gz, scale, shift, = node.inputs
x_u2i = U2IBatchNormalization(eps=node.op.eps)(x)
bn_out = mkl_bn.BatchNormalization(eps=node.op.eps,
bias=node.op.bias,
term=node.op.term)(x_u2i, scale, shift)
gz_u2i = I2UGrad()(bn_out, gz)
bn_GradOut = mkl_bn.BatchNormalizationGrad(eps=node.op.eps,
bias=node.op.bias,
term=node.op.term)(x_u2i, gz_u2i, scale, shift)
gx_i2u = U2IGrad()(x, bn_GradOut[0])
rval = [gx_i2u, bn_GradOut[1], bn_GradOut[2]]
return rval
except Exception as e:
msg = ('Failed to apply local opt to Op %s. '
'Exception message: %s\n') % (node.op, str(e))
_logger.warning(msg)
return
def test_bn_value(self):
X = T.ftensor4('x')
Scale = T.vector('scale')
Shift = T.vector('shift')
x_internal = U2IBatchNormalization(eps=1e-5)(X)
z_bn = mkl_bn.BatchNormalization(eps=1e-5, bias=1, term=1)(x_internal, Scale, Shift)
z_out = I2U()(z_bn)
z_sum = T.sum(z_out)
z_grad = T.grad(z_sum, [X])
fgrad = theano.function([X, Scale, Shift], z_grad, mode=with_mkl)
ival = numpy.random.rand(64, 5, 128, 128).astype(numpy.float32)
sval = numpy.random.rand(5).astype(numpy.float32)
tval = numpy.random.rand(5).astype(numpy.float32)
fgrad(ival, sval, tval)
def local_bn_mkl(node):
if not mkl_available():
return
if not isinstance(node.op, mkl_bn.AbstractBatchNormalization):
return
if node.inputs[0].type.ndim != 4:
return
try:
x, scale, shift, = node.inputs[0:3]
x_u2i = U2IBatchNormalization(eps=node.op.eps)(x)
bn_out = mkl_bn.BatchNormalization(eps=node.op.eps,
bias=node.op.bias,
term=node.op.term)(x_u2i, scale, shift)
z_i2u = I2U()(bn_out)
rval = z_i2u
return [rval]
except Exception as e:
msg = ('Failed to apply local opt to Op %s. '
'Exception message: %s\n') % (node.op, str(e))
_logger.warning(msg)
return