def Addn(data, target=utils.CCE):
"""
Compute sum of all elements in tensor.
Args:
data (tvm.tensor.Tensor): Tensor of of type float16, float32.
Returns:
tvm.tensor.Tensor, compute result, get all elements' sum.
Supported Platforms:
'Ascend', 'GPU', 'CPU'
"""
utils.check_supported_target(target)
# check types
dtype = data[0].dtype
if target == utils.CCE:
utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT)
res = data[0]
for i in range(1, len(data)):
utils.elemwise_dtype_check(res.dtype, data[i].dtype)
utils.elemwise_shape_check(res.shape, data[i].shape)
res = akg.topi.elemwise_sum(data)
return res
def matrix_diag_part(input_diagonal, input_help):
"""
Calculate the batched diagonal part of a batched tensor.
Note:
input_help is a tensor with a diagonal element of 1 and other positions of 0,
the last two dimensions can be unequal.
Args:
input_diagonal (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8. The last two dimensions
can be unequal.
input_help (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8, and with a diagonal element of 1
and other positions of 0.
Returns:
tvm.tensor.Tensor, has the same type as input_diagonal, the shape dims is equal to dims(input_diagonal) - 1.
"""
dtype_input_diagonal = input_diagonal.dtype
dtype_input_help = input_help.dtype
utils.elemwise_shape_check(input_help.shape, input_diagonal.shape)
if len(input_help.shape) < 2:
raise ValueError("Input tensors of rank>=2 are supported!")
utils.ops_dtype_check([dtype_input_diagonal, dtype_input_help], [
utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT8,
utils.DtypeForDavinci.INT32, utils.DtypeForDavinci.UINT8
])
res = matrix_diag_part_compute(input_diagonal, input_help)
return res
def atan_grad(head, input_x):
"""
Compute gradient of input_x in atan.
.. math::
dx = \\frac{1}{1 + x^2} \\cdot dy
Args:
head (tvm.tensor.Tensor): Gradient tensor of forward's output with the
same shape and dtype as input_x.
input_x (tvm.tensor.Tensor): Forward's input tensor support float16
and float32.
Returns:
A tvm.tensor.Tensor as gradient of forward's input.
Supported Platforms:
'Ascend'
"""
utils.elemwise_shape_check(head.shape, input_x.shape)
utils.elemwise_dtype_check(head.dtype, input_x.dtype,
utils.DtypeForDavinci.ALL_FLOAT)
dtype = input_x.dtype
tensor_one = dc.one_const(dtype)
def _compute(*i):
return tensor_one / (tensor_one + input_x(*i) * input_x(*i)) * head(*i)
out_tensor = tvm.compute(input_x.shape, _compute, name="out")
return out_tensor
def minimum_ad(head, data_x, data_y, grad_x=True, grad_y=True):
"""
Calculating the reversed outputs of the operator minimum by using automatic differentiate.
Args:
head (tvm.tensor.Tensor): Input tensor of float32, float16 and int32.
data_x (tvm.tensor.Tensor): Input tensor of float32, float16 and int32.
data_y (tvm.tensor.Tensor): Input tensor of float32, float16 and int32.
grad_x (bool): Default is True, whether to differentiate x.
grad_y (bool): Default is True, whether to differentiate y.
Returns:
tvm.tensor.Tensor, has the same type and shape as grads, if grad_x and grad_y all equal to True, need return
a list like: [jacs[0], jacs[1]].
"""
utils.elemwise_shape_check(data_x.shape, data_y.shape)
utils.elemwise_shape_check(head.shape, data_x.shape)
utils.elemwise_dtype_check(
data_x.dtype, head.dtype,
[utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT32])
utils.elemwise_dtype_check(
data_x.dtype, data_y.dtype,
[utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT32])
if not grad_x and not grad_y:
raise ValueError("At least one of grad_x and grad_y is True.")
op = minimum(data_x, data_y)
jacs = list(akg.differentiate(op, [data_x, data_y], head))
if grad_x and grad_y:
return jacs[0], jacs[1]
if grad_x:
return jacs[0]
return jacs[1]
def fake_quant_with_min_max_vars_per_channel_gradient(input_gradients,
input_data,
input_min,
input_max,
num_bits=8,
narrow_range=False):
"""
Computes gradients of Fake-quantize on the 'input_data' tensor,
output_backprops = input_gradients*(if input_data>=nudged_min and <=nudged_max 1 else 0)
Args:
input_gradients (tvm.tensor.Tensor): input gradients from previously operation
input_data (tvm.tensor.Tensor): input of fake-quantize, only supports "float32"
input_min (tvm.tensor.Tensor): input_min shape equals to input_max shape
The last dimension shoud be same for shapes of min, max and shape_inputs
only support fp32
input_max (tvm.tensor.Tensor): only support fp32
num_bits (int): Defaults to 8. bitwidth of the quantization,between 2 and 16
narrow_range (bool):
True, quantized into the quantization range [1, 2^num_bits - 1]
False,quantized into the quantization range [0, 2^num_bits - 1]
Returns:
tvm.tensor.Tensor
"""
input_gradients_shape = get_shape(input_gradients)
input_data_shape = get_shape(input_data)
input_min_shape = get_shape(input_min)
input_max_shape = get_shape(input_max)
utils.check_shape(input_gradients_shape)
utils.check_shape(input_data_shape)
utils.check_shape(input_min_shape)
utils.check_shape(input_max_shape)
utils.elemwise_shape_check(input_gradients.shape, input_data.shape)
utils.elemwise_shape_check(input_min_shape, input_max_shape)
if input_min_shape[0] != input_data_shape[-1]:
raise RuntimeError(
"The shapes of min,max and shape_inputs last one dimension shoud be same"
)
utils.ops_dtype_check(input_gradients.dtype, utils.DtypeForDavinci.FLOAT32)
utils.ops_dtype_check(input_data.dtype, utils.DtypeForDavinci.FLOAT32)
utils.ops_dtype_check(input_min.dtype, utils.DtypeForDavinci.FLOAT32)
utils.ops_dtype_check(input_max.dtype, utils.DtypeForDavinci.FLOAT32)
if num_bits > 16 or num_bits < 2:
raise RuntimeError("numbits should be range[2,16]")
input_min_broadcast = topi.broadcast_to(input_min, input_data_shape)
input_max_broadcast = topi.broadcast_to(input_max, input_data_shape)
res = fake_quant_with_min_max_vars_per_channel_gradient_compute(
input_gradients, input_data, input_min_broadcast, input_max_broadcast,
num_bits, narrow_range)
return res
def apply_ftrl(var, accum, linear, grad, lr, l1, l2, lr_power, target=utils.CCE):
"""
Ftrl-proximal optimization algorithm.
Note:
accum_new = accum + grad * grad
linear_new = linear + grad - (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
x = clip(linear_new, -l1, l1) - linear_new
y = accum_new^(-lr_power) / lr + 2 * l2
var_new = x / y
Args:
var (tvm.tensor.Tensor): The tensor to be updated. Should be float16 or float32.
accum (tvm.tensor.Tensor): A tensor of same shape and type as var. Eatch entry in it must be
greater or equal to zero.
linear (tvm.tensor.Tensor): A tensor of same shape and type as var.
grad (tvm.tensor.Tensor): A tensor of same shape and type as var.
lr (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
l1 (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
l2 (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
lr_power (tvm.tensor.Tensor): A scalar tensor of the same type as `var`. Value of it
must be less or equal to zero.
Returns:
tvm.tensor.Tensor, updated var.
tvm.tensor.Tensor, updated accum.
tvm.tensor.Tensor, updated linear.
"""
# As vlog instruction on mini product has a percision problem and mini product used to infer
# rather than train
if product_is_mini():
raise RuntimeError("The apply_ftrl operator does not support the mini product")
# check_shape
utils.check_shape(var)
shape = get_shape(var)
for tensor in (accum, linear, grad):
utils.elemwise_shape_check(shape, tensor.shape)
sclar_shape = (1,)
for sclar in (lr, l1, l2, lr_power):
utils.elemwise_shape_check(sclar.shape, sclar_shape)
# check dtype
dtype = var.dtype
utils.ops_dtype_check(dtype, [utils.DtypeForDavinci.FLOAT16, utils.DtypeForDavinci.FLOAT32])
for tensor in (var, accum, linear, grad, lr, l1, l2, lr_power):
utils.elemwise_dtype_check(tensor.dtype, dtype)
var_new, accum_new, linear_new = apply_ftrl_impl(var, accum, linear, grad, lr, l1, l2, None,
lr_power, with_l2_shrinkage=False)
# update by inplace
(var_new, accum_new, linear_new), binds_info = \
TensorUtils.inplace_set_tensors((var, accum, linear), (var_new, accum_new, linear_new))
attrs = {utils.BINDS: binds_info}
return var_new, accum_new, linear_new, attrs
def _apply_rms_prop_check(var, ms, mom, grad, lr, momentum, rho, epsilon):
"""Check inputs"""
utils.check_shape(var)
for i in (ms, mom, grad, lr, momentum, rho):
utils.elemwise_dtype_check(var.dtype, i.dtype)
for i in (ms, mom, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (lr, rho, momentum):
if tuple(get_shape(i)) != (1, ):
raise RuntimeError(
"lr, rho and momentum only support scalar tensor.")
if epsilon <= 0:
raise ValueError("epsilon should greater than zero.")
def apply_power_sign(var,
m,
grad,
lr,
logbase,
sign_decay,
beta,
target=utils.CCE):
"""
Update 'var' according to the PowerSign update
m_out = beta * m + (1 - beta) * grad
var_out = var - lr_t * (exp(logbase * sign_decay * Sign(grad) * Sign(m_out)) * grad)
Args:
var (tvm.tensor.Tensor): A tensor of type float16 or float32
m (tvm.tensor.Tensor): A tensor of same shape and type as var.
grad (tvm.tensor.Tensor): A tensor of same shape and type as var.
lr (tvm.tensor.Tensor): A scalar tensor of of same type as var.
logbase (tvm.tensor.Tensor): A scalar tensor of of same type as var.
sign_decay (tvm.tensor.Tensor): A scalar tensor of of same type as var.
beta (tvm.tensor.Tensor): A scalar tensor of of same type as var.
Returns:
tvm.tensor.Tensor, updated var.
tvm.tensor.Tensor, updated m.
"""
# check dtypes
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (m, grad, lr, logbase, sign_decay, beta):
utils.elemwise_dtype_check(var.dtype, i.dtype)
# check shapes
for i in (m, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (lr, logbase, sign_decay, beta):
if tuple(get_shape(i)) != (1, ):
raise RuntimeError(
"lr, logbase, sign_decay and beta only support scalar tensor.")
# compute
out_var, out_m = _apply_power_sign_compute(var, m, grad, lr, logbase,
sign_decay, beta)
# reuse var, m
out_var, binds_info = TensorUtils.inplace_set(var, out_var, "var_buf")
out_m, binds_info2 = TensorUtils.inplace_set(m, out_m, "m_buf")
binds_info.update(binds_info2)
attrs = {utils.BINDS: binds_info}
return out_var, out_m, attrs
def fake_quant_with_min_max_args_gradient(input_gradients,
input_data,
min=-6,
max=6,
num_bits=8,
narrow_range=False):
"""
Computes gradients of Fake-quantize on the 'input_data' tensor,
output_backprops = input_gradients*(if input_data>=nudged_min and <=nudged_max 1 else 0)
Args:
input_gradients (tvm.tensor.Tensor): input gradients from previously operation
input_data (tvm.tensor.Tensor): input of fake-quantize, only supports "float32"
min ([float, int]): scalar, defaults to -6
max ([float, int]): scalar, defaults to 6. [min; max] define the
clamping range for the input_data data
num_bits ([float, int]): Defaults to 8. num_bits is the bitwidth
of the quantization,between 2 and 16
narrow_range ([bool]):
True, quantized into the quantization range [1; 2^num_bits - 1]
False,quantized into the quantization range [0; 2^num_bits - 1]
Returns:
tvm.tensor.Tensor
"""
shape = get_shape(input_data)
utils.check_shape(shape)
utils.elemwise_shape_check(input_gradients.shape, input_data.shape)
utils.ops_dtype_check(input_data.dtype, utils.DtypeForDavinci.FLOAT32)
utils.ops_dtype_check(input_gradients.dtype, utils.DtypeForDavinci.FLOAT32)
nudged_min, nudged_max, scale = nudge_min_max(min, max, num_bits,
narrow_range)
zero_tensor = tvm.compute(input_data.shape,
lambda *i: tvm.const(0, dtype="float32"),
name="zero_tensor")
nudged_max_tensor = topi.add(zero_tensor, nudged_max)
nudged_min_tensor = topi.add(zero_tensor, nudged_min)
# where((input_data<=nudged_max)&(x>=nudged_min),1,0),Convert the input to 0 and 1 tensor
between_nudged_min_max = _cmpare_value(input_data, nudged_min_tensor,
nudged_max_tensor)
res = topi.multiply(input_gradients, between_nudged_min_max)
return res
def apply_proximal_adagrad(var, accum, lr, l1, l2, grad, target=utils.CCE):
"""
The FOBOS optimization algorithm with Adagrad learning rate.
Note:
accum_new = accum + grad * grad
ada_lr = lr * rsqrt(accum_new)
prox_var = var - ada_lr * grad
if l1 > 0:
var_new = Sign(prox_var)/(1+ada_lr*l2) * max{|prox_var|-ada_lr*l1,0}
else:
var_new = prox_var/(1+ada_lr*l2)
Args:
var (tvm.tensor.Tensor): The tensor to be updated. Should be float16 or float32.
accum (tvm.tensor.Tensor): A tensor of same shape and type as var. Eatch entry in it must be
greater or equal to zero.
lr (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
l1 (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
l2 (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
grad (tvm.tensor.Tensor): A tensor of same shape and type as var.
Returns:
tvm.tensor.Tensor, updated var.
tvm.tensor.Tensor, updated accum.
"""
# check_shape
utils.check_shape(var)
shape = get_shape(var)
for tensor in (accum, grad):
utils.elemwise_shape_check(shape, tensor.shape)
sclar_shape = (1, )
for sclar in (lr, l1, l2):
utils.elemwise_shape_check(sclar.shape, sclar_shape)
# check dtype
dtype = var.dtype
utils.ops_dtype_check(
dtype, [utils.DtypeForDavinci.FLOAT16, utils.DtypeForDavinci.FLOAT32])
for tensor in (var, accum, lr, l1, l2, grad):
utils.elemwise_dtype_check(tensor.dtype, dtype)
var_new, accum_new = _apply_proximal_adagrad_compute(
var, accum, lr, l1, l2, grad)
(var_new, accum_new), binds_info = TensorUtils.inplace_set_tensors(
[var, accum], [var_new, accum_new])
attrs = {utils.BINDS: binds_info}
return var_new, accum_new, attrs
def eltwise(data, mode=1, coeff=()):
"""
Compute elementwise modes, such as 0:PRODUCT, 1:SUM and 2:MAX.
Args:
data (list of tvm.tensor.Tensor): a list of tensor, tensor support fp16 and fp32.
mode (int): 0:product, 1:sum, 2:max.
coeff (tuple): tensor name of data should be equal with coeff size, only
used by sum, support int and float.
Returns:
tvm.tensor.Tensor.
"""
dtype = data[0].dtype
utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT)
utils.check_shape(data[0].shape)
shape_data = get_shape(data[0])
if not mode in [0, 1, 2]:
raise RuntimeError("mode only support 0, 1, or 2")
if not len(data) == len(coeff) and len(coeff) != 0:
raise RuntimeError(
"coeff should be [] or its length be same as data")
tensor_num = len(data)
#tensor num must be [1, 120]
if tensor_num < 1 or tensor_num > 120:
raise RuntimeError("tensor_num need in range [1,120].")
if mode == 1 and len(coeff) == 0:
return addn.addn(data)
if len(coeff) != 0:
if type(coeff[0]) != int and type(coeff[0]) != float:
raise RuntimeError("ele of coeff must be a number.")
for i in range(1, len(data)):
utils.elemwise_dtype_check(data[0].dtype, data[i].dtype)
utils.elemwise_shape_check(data[0].shape, data[i].shape)
if mode == 1 and len(coeff) > 0:
return _addn(data, coeff)
if mode == 0:
return _product(data)
if mode == 2:
return _max(data)
def apply_centered_rms_prop(var, mg, ms, mom, grad, lr, momentum, rho, epsilon, target=utils.CCE):
"""
Update `var` according to the centered RMSProp algorithm.
out_mean_grad = decay * mg + (1-decay) * grad
out_mean_square = decay * ms + (1-decay) * grad * grad
out_mom = momentum * mom + lr * grad / sqrt(out_mean_square - out_mean_grad^2 + epsilon)
out_var = var - out_mom
Args:
var (tvm.tensor.Tensor): Input data of type float16 or float32.
mg (tvm.tensor.Tensor): A tensor of the same type and shape as `var`.
ms (tvm.tensor.Tensor): A tensor of the same type and shape as `var`.
mom (tvm.tensor.Tensor): A tensor of the same type and shape as `var`.
grad (tvm.tensor.Tensor): A tensor of the same type and shape as `var`.
lr (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
momentum (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
rho (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
epsilon (float): A scalar tensor of the same type as `var`.
Returns:
tvm.tensor.Tensor, updated var.
tvm.tensor.Tensor, updated mean_grad.
tvm.tensor.Tensor, updated mean_square.
tvm.tensor.Tensor, updated mom.
"""
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (mg, ms, mom, lr, rho, momentum, grad):
utils.elemwise_dtype_check(var.dtype, i.dtype)
for i in (mg, ms, mom, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (lr, rho, momentum):
if tuple(get_shape(i)) != (1,):
raise RuntimeError("lr, rho and momentum only support scalar tensor.")
if epsilon <= 0:
raise ValueError("epsilon should be greater than 0.")
out_var, out_mg, out_ms, out_mom = _apply_centered_rms_prop_compute(
var, mg, ms, mom, grad, lr, momentum, rho, epsilon)
out_var, binds_info = TensorUtils.inplace_set(var, out_var, "var_buf")
out_mg, binds_info2 = TensorUtils.inplace_set(mg, out_mg, "mg_buf")
out_ms, binds_info3 = TensorUtils.inplace_set(ms, out_ms, "ms_buf")
out_mom, binds_info4 = TensorUtils.inplace_set(mom, out_mom, "mom_buf")
binds_info.update(binds_info2)
binds_info.update(binds_info3)
binds_info.update(binds_info4)
attrs = {utils.BINDS: binds_info}
return out_var, out_mg, out_ms, out_mom, attrs
def _check_inputs(var, grad_accum, grad_squared_accum, grad, lr, l1, l2,
global_step):
"""Check op inputs"""
# check dtype
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (grad_accum, grad_squared_accum, grad, lr, l1, l2):
utils.elemwise_dtype_check(var.dtype, i.dtype)
utils.ops_dtype_check(global_step.dtype, utils.DtypeForDavinci.INT32)
# check shape
for i in (grad_accum, grad_squared_accum, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (lr, l1, l2, global_step):
if tuple(get_shape(i)) != (1, ):
raise RuntimeError(
"lr, l1, l2 and global_step only support scalar tensor.")
def inv_grad(input_y, input_dy):
"""
Calculate data's Reciprocal grad,dx = -1 * input_dy * input_y * input_y.
Args:
input_y (tvm.tensor.Tensor): Tensor of type float16, float32, int8, int32.
input_dy (tvm.tensor.Tensor): Tensor of type float16, float32, int8, int32.
Returns:
tvm.tensor.Tensor, has the same type and shape as input_y.
"""
# Check shapes and dtypes.
utils.elemwise_shape_check(input_y.shape, input_dy.shape)
utils.elemwise_dtype_check(input_y.dtype, input_dy.dtype, supported_type=["float16", "float32", "int8", "int32"])
res = inv_grad_compute(input_y, input_dy)
return res
def matrix_set_diag(input_matrix, input_diagonal, input_help):
"""
Return a batched matrix tensor with new batched diagonal values.
Args:
input_matrix (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8. The last two dimensions
can be unequal.
input_diagonal (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8.The last shape need equal
to min(input_matrix[-1], input_matrix[-2]).
input_help (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8,and with a diagonal element of 1
and other positions of 0.
Returns:
tvm.tensor.Tensor, has the same type and shape as input_matrix.
"""
shape_input = get_shape(input_matrix)
shape_diag = get_shape(input_diagonal)
shape_help = get_shape(input_help)
dtype = input_matrix.dtype
utils.check_shape(shape_input)
utils.check_shape(shape_diag)
utils.check_shape(shape_help)
# Check help_matrix.
if (len(shape_input) < 2) or (len(shape_help) < 2):
raise RuntimeError("Only the rank of input tensors >= 2 are supported!")
utils.elemwise_shape_check(shape_input, shape_help)
# Check support dtype.
utils.ops_dtype_check(dtype, [utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT8,
utils.DtypeForDavinci.INT32, utils.DtypeForDavinci.UINT8])
# Adjust diag's shape according to input shape.
# Extend the shape_diag dimension for broadcast.
# if input_shape is [2,4,7,9] and shape_diag is [2,4,7] then new_shape is [2,4,7,1]
# if input_shape is [2,4,9,7] and shape_diag is [2,4,7], then new_shape is [2,4,1,7]
if shape_input[-2] <= shape_input[-1]:
shape_b_newshape = list(shape_diag) + [1]
# The penultimate dimension of the shape_diag is extended for broadcast.
else:
shape_b_newshape = list(shape_diag)
shape_b_newshape.insert(-1, 1)
input_diagonal = topi.reshape(input_diagonal, shape_b_newshape)
res = matrix_set_diag_compute(input_matrix, input_diagonal, input_help)
return res
def _check_inputs(var, accum, accum_update, grad, lr, rho, epsilon):
"""Check op inputs"""
# check dtype
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (accum, accum_update, grad, lr, rho):
utils.elemwise_dtype_check(var.dtype, i.dtype)
# check shape
for i in (accum, accum_update, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (lr, rho):
if tuple(get_shape(i)) != (1,):
raise RuntimeError("lr and rho only support scalar tensor.")
# check value
if epsilon <= 0:
raise ValueError("epsilon should be greater than zero.")
def apply_add_sign(var,
m,
grad,
lr,
alpha,
sign_decay,
beta,
target=utils.CCE):
"""
Update 'var' according to the AddSign update.
m_out = m * beta + grad * (1 - beta)
var_out = var - lr * (alpha + sign_decay * Sign(grad) *Sign(m)) * grad
Args:
var (tvm.tensor.Tensor): A tensor of type float16 or float32
m (tvm.tensor.Tensor): A tensor of type float16 or float32
grad (tvm.tensor.Tensor): A tensor of type float16 or float32
lr (tvm.tensor.Tensor): A scalar tensor of type float16 or float32
alpha (tvm.tensor.Tensor): A scalar tensor of type float16 or float32
sign_decay (tvm.tensor.Tensor): A scalar tensor of type float16 or float32
beta (tvm.tensor.Tensor): A scalar tensor of type float16 or float32
Returns:
tvm.tensor.Tensor, updated var.
tvm.tensor.Tensor, updated m.
"""
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (m, lr, alpha, sign_decay, beta, grad):
utils.elemwise_dtype_check(var.dtype, i.dtype)
for i in (m, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (lr, alpha, sign_decay, beta):
if tuple(get_shape(i)) != (1, ):
raise RuntimeError(
"lr, alpha, sign_decay and beta only support scalar.")
out_var, out_m = _apply_add_sign_compute(var, m, grad, lr, alpha,
sign_decay, beta)
out_var, binds_info = TensorUtils.inplace_set(var, out_var, "var_buf")
out_m, binds_info2 = TensorUtils.inplace_set(m, out_m, "m_buf")
binds_info.update(binds_info2)
attrs = {utils.BINDS: binds_info}
return out_var, out_m, attrs
def fake_quant_with_min_max_vars_per_channel(input_data,
input_min,
input_max,
num_bits=8,
narrow_range=False):
"""
Generate fake_quantize the input_data for every channel.
Note:
For input_data last dim must be equal to d. And need to satisfy: input_min <= 0 <= input_max.
Args:
input_data (tvm.tensor.Tensor): Tensor of type float32, shape must be equal to [b, d] or [b, h, w, d] or [d].
input_min (tvm.tensor.Tensor): Tensor of type float32, shape must be equal to [d].
input_max (tvm.tensor.Tensor): Tensor of type float32, shape must be equal to [d].
num_bits (int): The quantization bits, must be int, defaults to 8.
narror_range (Union[bool, None]): if True, quant_min equal to 1, else 0, defaults to False.
Returns:
tvm.tensor.Tensor of same type and shape as input_data.
"""
# get shape and check
shape_inputs = get_shape(input_data)
shape_min = get_shape(input_min)
shape_max = get_shape(input_max)
utils.elemwise_shape_check(shape_min, shape_max)
utils.auto_broadcast_check(shape_min, shape_inputs)
if shape_min[0] != shape_inputs[-1]:
raise RuntimeError(
"The shapes of min,max and shape_inputs last one dimension should be same!"
)
# check dtype
utils.ops_dtype_check(input_data.dtype, utils.DtypeForDavinci.FLOAT32)
utils.elemwise_dtype_check(input_min.dtype, input_max.dtype,
utils.DtypeForDavinci.FLOAT32)
# check num_bits range
if num_bits > 16 or num_bits < 2:
raise ValueError("numbits should be in range [2, 16]!")
# get output by fake_quant_with_min_max_vars_per_channel_compute function
res = fake_quant_with_min_max_vars_per_channel_compute(
input_data, input_min, input_max, num_bits, narrow_range)
return res
def _check_inputs(var, m, v, grad, lr, beta1, beta1_power, beta2, epsilon):
"""Check op inputs"""
# check dtype
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (m, v, grad, beta1_power, lr, beta1, beta2):
utils.elemwise_dtype_check(var.dtype, i.dtype)
# check shape
for i in (m, v, grad):
utils.elemwise_shape_check(var.shape, i.shape)
for i in (beta1_power, lr, beta1, beta2):
if tuple(get_shape(i)) != (1, ):
raise RuntimeError(
"beta1_power, lr, beta1 and beta2 only support scalar tensor.")
# check value
if epsilon <= 0:
raise ValueError("epsilon should be greater than zero.")
def l1_loss(inputs, outputs, reduction='none', target="cce"):
"""
Computes l1 loss.
Args:
inputs(akg.tvm.Tensor): Supported data type is float16, float32.
outputs(akg.tvm.Tensor): With same type as inputs.
reduction(str): Default is 'none', could be 'sum' or 'mean', if 'mean', loss result will be divided
by the size of inputs.
Returns:
akg.tvm.Tensor of same type as input tensors.
"""
inputs_dtype = inputs.dtype
target_dtype = outputs.dtype
# check inputs data types
utils.ops_dtype_check([inputs_dtype, target_dtype], utils.DtypeForDavinci.ALL_FLOAT)
target_shape = [x.value for x in outputs.shape]
inputs_shape = [x.value for x in inputs.shape]
utils.elemwise_shape_check(target_shape, inputs_shape)
inputs_dtype_old = inputs_dtype
if product_is_mini() and inputs_dtype == 'float32':
inputs = akg.topi.cast(inputs, "float16")
outputs = akg.topi.cast(outputs, "float16")
inputs_dtype = "float16"
diff = akg.topi.subtract(inputs, outputs)
loss = akg.topi.abs(diff)
if reduction == 'sum':
loss = akg.topi.sum(loss)
if reduction == 'mean':
loss = akg.topi.sum(loss)
deno = 1.0
for num in inputs.shape:
deno = deno * num
deno = akg.topi.cast(deno, dtype=inputs_dtype)
loss = akg.topi.divide(loss, deno)
if product_is_mini() and inputs_dtype_old == 'float32':
loss = akg.topi.cast(loss, inputs_dtype_old)
return loss
def acosh_grad(y, dy):
"""
Gradient for acosh.
Note:
dx = dy * 1/sinh(y)
Args:
y (tvm.tensor.Tensor): tensor of type float16, float32.
dy (tvm.tensor.Tensor): same type and shape as y.
Returns:
tvm.tensor.Tensor, same type and shape as y.
Supported Platforms:
'Ascend'
"""
# mini product just used to infer
if product_is_mini():
raise RuntimeError(
"The mini product does not support the acosh_grad operator")
dtype = y.dtype
utils.ops_dtype_check(y.dtype, utils.DtypeForDavinci.ALL_FLOAT)
utils.elemwise_dtype_check(dtype, dy.dtype)
utils.check_shape(y.shape)
utils.elemwise_shape_check(y.shape, dy.shape)
if dtype == "float16":
y = topi.cast(y, "float32")
dy = topi.cast(dy, "float32")
# If we use sinh(y) = (exp(y) - exp(-y))/2 directly, there will be some precision problems
# For example, as dx = dy/sinh(y), if we use exp directly, when exp(y) and exp(-y) are close,
# the small precision error of exp calculation will be greatly enlarged in the final result
sinh_y = _sinh_taylor(y)
dx = topi.divide(dy, sinh_y)
if dx.dtype != dtype:
dx = topi.cast(dx, dtype)
attrs = {"enable_auto_inline": False}
return dx, attrs
def apply_proximal_gradient_descent(var, alpha, l1, l2, delta, target=utils.CCE):
"""
The FOBOS algorithm with fixed learning rate.
Note:
prox_var = var - alpha * delta
if l1 > 0:
var_new = Sign(prox_var)/(1+alpha*l2) * max{|prox_var|-alpha*l1,0}
else:
var_new = prox_var/(1+alpha*l2)
Args:
var (tvm.tensor.Tensor): The tensor to be updated. Should be float16 or float32.
alpha (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
l1 (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
l2 (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
delta (tvm.tensor.Tensor): A tensor of same shape and type as var.
Returns:
tvm.tensor.Tensor, updated var.
"""
# check_shape
utils.check_shape(var)
shape = get_shape(var)
utils.elemwise_shape_check(shape, delta.shape)
sclar_shape = (1,)
for sclar in (alpha, l1, l2):
utils.elemwise_shape_check(sclar.shape, sclar_shape)
# check dtype
dtype = var.dtype
utils.ops_dtype_check(dtype, [utils.DtypeForDavinci.FLOAT16, utils.DtypeForDavinci.FLOAT32])
for tensor in (var, alpha, l1, l2, delta):
utils.elemwise_dtype_check(tensor.dtype, dtype)
var_new = apply_proximal_gradient_descent_impl(var, alpha, l1, l2, delta)
var_new, binds_info = TensorUtils.inplace_set(var, var_new, "var_buf")
attrs = {utils.BINDS: binds_info}
return var_new, attrs
def asinh_grad(y, dy):
"""
Gradient for asinh.
Note:
dx = dy * 1/cosh(y)
Args:
y (tvm.tensor.Tensor): tensor of type float16, float32.
dy (tvm.tensor.Tensor): same type and shape as y.
Returns:
tvm.tensor.Tensor, same type and shape as y.
Supported Platforms:
'Ascend'
"""
# mini product just used to infer
if product_is_mini():
raise RuntimeError(
"The mini product does not support the asinh_grad operator")
dtype = y.dtype
utils.ops_dtype_check(y.dtype, utils.DtypeForDavinci.ALL_FLOAT)
utils.elemwise_dtype_check(dtype, dy.dtype)
utils.check_shape(y.shape)
utils.elemwise_shape_check(y.shape, dy.shape)
if dtype == "float16":
y = topi.cast(y, "float32")
dy = topi.cast(dy, "float32")
dx = topi.divide(dy, cosh(y))
if dx.dtype != dtype:
dx = topi.cast(dx, dtype)
return dx
def atan2(y, x):
"""
Compute arc tangent of y/x.
.. math::
\\arctan2(y, x) = \\arctan(\\frac{y}{x})
Args:
y (tvm.tensor.Tensor): Input tensor, only support float16, float32.
x (tvm.tensor.Tensor): Input tensor, only support float16, float32.
Returns:
A tvm.tensor.Tensor as angles in radians.
Supported Platforms:
'Ascend'
"""
utils.elemwise_shape_check(get_shape(y), get_shape(x))
utils.elemwise_dtype_check(y.dtype, x.dtype,
utils.DtypeForDavinci.ALL_FLOAT)
return _atan2_compute(y, x), {"enable_auto_inline": False}
def asin_grad(x, dy):
"""
Gradient for arcsin.
.. math::
\\frac {\\partial arcsin(x)} {\\partial x} = \\frac{1}{\\sqrt{1 - x^2}}
Args:
x (tvm.tensor.Tensor): Tensor of type float16, float32.
dy (tvm.tensor.Tensor): Tensor of same type and shape as x.
Rerurns:
tvm.tensor.Tensor of same type and shape as x.
Supported Platforms:
'Ascend'
"""
utils.ops_dtype_check(x.dtype, utils.DtypeForDavinci.ALL_FLOAT)
utils.elemwise_dtype_check(x.dtype, dy.dtype)
utils.elemwise_shape_check(x.shape, dy.shape)
return _asin_grad_compute(x, dy)
def assign_sub(data1, data2):
"""
Computes data1 - data2 elementwise.
Args:
data1 (tvm.tensor.Tensor): Tensor of type float16, float32, int32, int8, uint8.
data2 (tvm.tensor.Tensor): Tensor of same shape and type as data1.
Returns:
Subtracted result, with same shape and type as input tensors.
Supported Platforms:
'Ascend'
"""
dtype = data1.dtype
utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_TYPES)
utils.elemwise_dtype_check(data1.dtype, data2.dtype)
utils.elemwise_shape_check(data1.shape, data2.shape)
need_cast_dtype = ["int8", "uint8"]
cast_type = "float16"
if dtype in need_cast_dtype:
data1 = akg.topi.cast(data1, cast_type)
data2 = akg.topi.cast(data2, cast_type)
res = akg.topi.subtract(data1, data2)
if dtype in need_cast_dtype:
if dtype == "uint8":
cons = akg.tvm.const(256, dtype=cast_type)
res = akg.tvm.compute(res.shape,
lambda *indice: akg.tvm.expr.Select(
res(*indice) < 0,
res(*indice) + cons, res(*indice)),
name="positive_res")
res = akg.topi.cast(res, dtype)
return res
def apply_adagrad(var,
accum,
learning_rate,
grad,
update_slots=True,
target=utils.CCE):
"""
Update `var` according to the Adagrad algorithm.
.. math:
accum += grad^2
var -= learning_rate * grad / accum.sqrt()
Args:
var (tvm.tensor.Tensor): input var to be updated of type float16, float32
accum (tvm.tensor.Tensor): accumulation of the squared gradients of type float16, float32
learning_rate (tvm.tensor.Tensor): A scalar tensor of type float16, float32
grad (tvm.tensor.Tensor): input grad of type float16, float32
update_slots (Bool): If True, the accum tensor will be updated;
otherwise the option is False, the accum tensor will not be update.
Defaults to 'True'.
Returns:
tvm.tensor.Tensor, the updated var.
tvm.tensor.Tensor, the updated accum.
"""
utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.ALL_FLOAT)
for i in (accum, learning_rate, grad):
utils.elemwise_dtype_check(var.dtype, i.dtype)
for i in (accum, grad):
utils.elemwise_shape_check(var.shape, i.shape)
if tuple(get_shape(learning_rate)) != (1, ):
raise RuntimeError("learning_rate only support scalar tensor")
return _apply_adagrad_compute(var, accum, learning_rate, grad,
update_slots)
def sgd(parameters, gradient, accum, stat, learning_rate, momentum, dampening=0.0, weight_decay=0.0, nesterov=False):
"""
Update parameters, accum and stat according to the SGD algorithm.
accum = accum * momentum + grad
if use_nesterov is True:
parameters -= grad * lr + accum * momentum * lr
else:
parameters -= accum * lr
Args:
parameters (tvm.tensor.Tensor): parameters tensor of float32, float16, to be updated.
gradient (tvm.tensor.Tensor): gradient tensor of float32, float16.
accum (tvm.tensor.Tensor): accum tensor of float32, float16, to be updated.
stat (tvm.tensor.Tensor): stat tensor of float32, float16, to be updated.
momentum (tvm.tensor.Tensor): momentum tensor of float32, float16, shape must be equal to (1,).
learning_rate (tvm.tensor.Tensor): learning_rate tensor of float32, float16, shape must be equal to (1,).
dampening (float): Default value is 0.0.
weight_decay (float): Default value is 0.0.
nesterov (bool): Default is False.
Return:
accum_t (tvm.tensor.Tensor): updated accum with same type and shape as accum.
stat_t (tvm.tensor.Tensor): updated stat with same type and shape as stat.
parameters_t (tvm.tensor.Tensor): updated parameters with same type and shape as parameters.
"""
if nesterov and dampening != 0:
raise ValueError("Nesterov requires zero dampening!")
if weight_decay < 0:
raise ValueError("weight_decay must > 0.")
# shape check
utils.elemwise_shape_check(parameters.shape, gradient.shape)
utils.elemwise_shape_check(parameters.shape, accum.shape)
utils.elemwise_shape_check(parameters.shape, stat.shape)
# dtype check
utils.ops_dtype_check([parameters.dtype, gradient.dtype, accum.dtype, stat.dtype],
utils.DtypeForDavinci.ALL_FLOAT)
parameters_t, accum_t, stat_t = sgd_compute(parameters, gradient, learning_rate, accum, momentum, stat, dampening,
weight_decay, nesterov)
parameters_t, binds_info = TensorUtils.inplace_set(parameters, parameters_t, "parameters_buf")
accum_t, binds_info2 = TensorUtils.inplace_set(accum, accum_t, "accum_buf")
stat_t, binds_info3 = TensorUtils.inplace_set(stat, stat_t, "stat_buf")
binds_info.update(binds_info2)
binds_info.update(binds_info3)
attrs = {utils.BINDS: binds_info}
return parameters_t, accum_t, stat_t, attrs
def cross(x, y, target=utils.CCE):
"""
Compute cross product of x and y.
Note:
The first dim of x or y must be 3, it will be calculated as (two dims for example)
.. math::
res = x \\times y = \\left[ \\begin{matrix}
l, & \\cdots \\\\ m, & \\cdots \\\\ n, & \\cdots
\\end{matrix} \\right] \\times \\left[ \\begin{matrix}
o, & \\cdots \\\\ p, & \\cdots \\\\ q, & \\cdots
\\end{matrix} \\right] = \\left[ \\begin{matrix}
mq-np, & \\cdots \\\\ no-lq, & \\cdots \\\\ lp-mo, & \\cdots \\\\
\\end{matrix} \\right]
Args:
x (tvm.tensor.Tensor): Input tensor, only support float16, float32,
int32, int8, uint8.
y (tvm.tensor.Tensor): Input tensor, must have the same shape and dtype
as x.
Returns:
A tvm.tensor.Tensor with the same shape and dtype as x.
"""
utils.elemwise_shape_check(get_shape(y), get_shape(x))
utils.elemwise_dtype_check(
y.dtype, x.dtype,
(utils.DtypeForDavinci.ALL_FLOAT) if product_is_mini() \
else (utils.DtypeForDavinci.FLOAT16,
utils.DtypeForDavinci.FLOAT32,
utils.DtypeForDavinci.INT32,
utils.DtypeForDavinci.INT8, utils.DtypeForDavinci.UINT8))
shape = get_shape(x)
if shape[0] != 3:
raise RuntimeError(
"The first axis of input must be 3, actual input is %d" % shape[0])
inp_dtype = x.dtype
need_type_convert = inp_dtype in ("int8", "uint8")
shape = get_shape(x)
shp = shape[1:]
if need_type_convert:
x = Cast(x, "float16", target=utils.CCE)
y = Cast(y, "float16", target=utils.CCE)
a0b1 = tvm.compute(shp, lambda *i: x(0, *i) * y(1, *i), name="a0b1")
a0b2 = tvm.compute(shp, lambda *i: x(0, *i) * y(2, *i), name="a0b2")
a1b0 = tvm.compute(shp, lambda *i: x(1, *i) * y(0, *i), name="a1b0")
a1b2 = tvm.compute(shp, lambda *i: x(1, *i) * y(2, *i), name="a1b2")
a2b0 = tvm.compute(shp, lambda *i: x(2, *i) * y(0, *i), name="a2b0")
a2b1 = tvm.compute(shp, lambda *i: x(2, *i) * y(1, *i), name="a2b1")
res0 = tvm.compute(shp, lambda *i: a1b2(*i) - a2b1(*i), name="res0")
res1 = tvm.compute(shp, lambda *i: a2b0(*i) - a0b2(*i), name="res1")
res2 = tvm.compute(shp, lambda *i: a0b1(*i) - a1b0(*i), name="res2")
res = tvm.compute(
shape,
lambda *i: tvm.expr.Select(
i[0] == 0, res0(*i[1:]),
tvm.expr.Select(i[0] == 1, res1(*i[1:]), res2(*i[1:]))),
name='res')
if need_type_convert:
res = Cast(res, inp_dtype, target=utils.CCE)
return res
def apply_adam(var,
m,
v,
beta1_power,
beta2_power,
lr,
beta1,
beta2,
epsilon,
grad,
use_nesterov=False,
target=utils.CCE):
"""
Adam and Nadam optimization algorithm.
Note:
lr_t = lr*sqrt(1-beta2_power)/(1-beta1_power)
m_new = m + (1-beta1)*(grad-m)
v_new = v + (1-beta2)*(grad*grad-v)
if user_nesterov == True:
var_new = var - lr_t*(m_new*beta1 + (1-beta1)*grad) / (epsilon + sqrt(v_new))
else:
var_new = var - lr_t*m_new / (epsilon + sqrt(v_new))
Args:
var (tvm.tensor.Tensor): The tensor to be updated. Should be float16 or float32.
m (tvm.tensor.Tensor): The first moment estimate. A tensor of same shape and type as var.
v (tvm.tensor.Tensor): The second moment estimate. A tensor of same shape and type as var.
beta1_power (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
beta2_power (tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
lr (tvm.tensor.Tensor): The learning rate. A scalar tensor of the same type as `var`.
beta1(tvm.tensor.Tensor): A tensor with shape (1,) and type is same as var.
beta2(tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
epsilon(tvm.tensor.Tensor): A scalar tensor of the same type as `var`.
grad (tvm.tensor.Tensor): A tensor of same shape and type as var.
use_nesterov(bool): Default value is False. If use_nesterov is True, the Nadam algorithm be implemented,
otherwise the adam algorithm be implemented.
Returns:
tvm.tensor.Tensor, updated var.
tvm.tensor.Tensor, updated m.
tvm.tensor.Tensor, updated v.
"""
# check shape
utils.check_shape(var)
shape = get_shape(var)
for tensor in (m, v, grad):
utils.elemwise_shape_check(shape, tensor.shape)
sclar_shape = (1, )
for sclar in (beta1_power, beta2_power, lr, beta1, beta2, epsilon):
utils.elemwise_shape_check(sclar.shape, sclar_shape)
# check dtype
dtype = var.dtype
utils.ops_dtype_check(
dtype, [utils.DtypeForDavinci.FLOAT16, utils.DtypeForDavinci.FLOAT32])
for tensor in (var, m, v, beta1_power, beta2_power, lr, beta1, beta2,
epsilon, grad):
utils.elemwise_dtype_check(tensor.dtype, dtype)
var_new, m_new, v_new = _apply_adam_compute(var, m, v, beta1_power,
beta2_power, lr, beta1, beta2,
epsilon, grad, use_nesterov)
# update by inplace
(var_new, m_new, v_new), binds_info = TensorUtils.inplace_set_tensors(
[var, m, v], [var_new, m_new, v_new])
attrs = {utils.BINDS: binds_info}
return var_new, m_new, v_new, attrs
