Abstract: This paper briefly introduces the registration and association process of AKG forward and reverse operators.
This article is shared from “AKG Positive and negative Operator Registration + Association” in Huawei cloud community, by Muzi _007.
A, environmental
Hardware: eulerosv2r8 aarch64
Mindspore: 1.1
Operator registration needs to be compiled and installed to take effect, so mindspore’s source code is already available in the default environment and can already be compiled and installed
Two, forward operator production and testing
I’m going to make an operator to square a vector
Forward: y = x**2
Reverse: y = 2*x
Let’s start with forward
2.1 Defining a forward operator
Path: mindspore akg/python/akg/ms/cce, create cus_square. Py
Define the calculation logic of vector squared by referring to the definition of calculation logic in the same directory
"""cus_square"""
from akg.tvm.hybrid import script
from akg.ops.math import mul
import akg
def CusSquare(x):
output_shape = x.shape
k = output_shape[0]
n = output_shape[1]
@script
def cus_square_compute(x):
y = output_tensor(output_shape, dtype=x.dtype)
for i in range(k):
for j in range(n):
y[i, j] = x[i, j] * x[i, j]
return y
output = cus_square_compute(x)
attrs = {
'enable_post_poly_loop_partition': False,
'enable_double_buffer': False,
'enable_feature_library': True,
'RewriteVarTensorIdx': True
}
return output, attrs
Copy the codeThen add the contents to the __init__.py file in the sibling directory
from .cus_square import CusSquare
Copy the code2.2 Registration Operator
To mindspore/ops/_op_impl/akg/ascend, create Cus_square.py and add the following code
"""CusSquare op"""
from mindspore.ops.op_info_register import op_info_register, AkgAscendRegOp, DataType as DT
op_info = AkgAscendRegOp("CusSquare") \
.fusion_type("ELEMWISE") \
.input(0, "x") \
.output(0, "output") \
.dtype_format(DT.F32_Default, DT.F32_Default) \
.get_op_info()
@op_info_register(op_info)
def _cus_square_akg():
"""CusSquare Akg register"""
return
Copy the codeThen add the following code to __init__.py of the sibling directory
from .cus_square import _cus_square_akg
Copy the code2.3 Defining operator primitives
Go to: mindspore/ops/operations, create a new _cus_ops.py and add the following code
Description operator input: x, output output
Infer_shape: Describes the shape of output data
Infer_dtype: indicates the type of output data
X1_shape: refers to the shape entered first
X1_dtype: indicates the dtype of the first input parameter
import math from .. primitive import prim_attr_register, PrimitiveWithInfer from ... common import dtype as mstype from ... _checkparam import Validator as validator from ... _checkparam import Rel class CusSquare(PrimitiveWithInfer): """CusSquare""" @prim_attr_register def __init__(self): self.init_prim_io_names(inputs=['x'], outputs=['output']) def infer_shape(self, x1_shape): return x1_shape def infer_dtype(self, x1_dtype): return x1_dtypeCopy the codeThen add the primitive information to the __init__.py file in the same directory
from ._cus_ops import CusSquare
Copy the code2.4 Add operator query information to CCSRC
In mindspore/CCSRC/backend/kernel_compiler/kernel_query cc KernelQuery function to add the following information
// cus_square
const PrimitivePtr kPrimCusSquare = std::make_shared<Primitive>("CusSquare");
if (IsPrimitiveCNode(kernel_node, kPrimCusSquare)) {
kernel_type = KernelType::AKG_KERNEL;
}
Copy the code2.5 Compiling and Installing the Framework
Go back to the Mindspore root directory
Sh -e ascend- j4 CD./build/package PIP install mindspore_ascend-1.1.2- cp37-cp37m-linux_aARCH64.whl --force-reinstallCopy the code2.6 test
import numpy as np import mindspore.nn as nn import mindspore.context as context from mindspore import Tensor from mindspore.ops import operations as P context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") class Net(nn.Cell): def __init__(self): super(Net, self).__init__() self.square = P.CusSquare() def construct(self, data): return self.square(data) def test_net(): X = np.array([[1.0, 4.0, 9.0]]).astype(np.float32) net = NET () output = net(x) print("x: ", x) print("output: ", output) if __name__ == "__main__": test_net()Copy the codeThe output
Three, the production and test of the reverse operator
3.1 Production Process
The calculation logic of the reverse operator: take the derivative of the vector element, such as y = x^2, then y ‘= 2x
The actual example is the input vector [1, 4, 9] and the output is [2, 8, 18]
The reverse operator is clearly CusSquareGrad, which is the same as the previous square calculation operator process. Only the key codes are posted here, and the process will not be described again
Calculate the logic code cus_square_grad.py
"""cus_square_grad"""
from akg.tvm.hybrid import script
import akg
def CusSquareGrad(x):
output_shape = x.shape
k = output_shape[0]
n = output_shape[1]
@script
def cus_square_compute_grad(x):
y = output_tensor(output_shape, dtype=x.dtype)
for i in range(k):
for j in range(n):
y[i, j] = x[i, j] * 2
return y
output = cus_square_compute_grad(x)
attrs = {
'enable_post_poly_loop_partition': False,
'enable_double_buffer': False,
'enable_feature_library': True,
'RewriteVarTensorIdx': True
}
return output, attrs
Copy the codeRegistered primitives
class CusSquareGrad(PrimitiveWithInfer):
"""
CusSquareGrad
"""
@prim_attr_register
def __init__(self):
self.init_prim_io_names(inputs=['x'], outputs=['output'])
def infer_shape(self, x1_shape):
return x1_shape
def infer_dtype(self, x1_dtype):
return x1_dtype
Copy the code3.2 test
import numpy as np import mindspore.nn as nn import mindspore.context as context from mindspore import Tensor from mindspore.ops import operations as P context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") class Net(nn.Cell): def __init__(self): Super (Net, self).__init__() self.square = p.cosquaregrad () # replace the grad operator def construct(self, data): return self.square(data) def test_net(): X = np.array([[1.0, 4.0, 9.0]]).astype(np.float32) net = NET () output = net(x) print("x: ", x) print("output: ", output) if __name__ == "__main__": test_net()Copy the codeThe output
Four, positive and negative operator correlation and testing
Add the following code to the source mindspore/mindspore/ops/_grad/grad_array_ops.py
@bprop_getters.register(P.CusSquare)
def get_bprop_cussquare(self):
"""Generate bprop of CusSquare"""
cus_square_grad = P.CusSquareGrad()
matmul = ops.Mul()
def bprop(x, out, dout):
gradient = cus_square_grad(x)
dx = matmul(gradient, dout)
return (dx,)
return bprop
Copy the codeThe input of the BPROP function is a forward input x, a forward output OUT, and a reverse gradient input dout
The above code is meant to specify the calculation of the inverse gradient of the operator CusSquareGrad, which is used as one of the functions
Gradient = Cus_square_grad (x) calculates the gradient of the square operator, but cannot directly return this gradient
Reverse the network to this operator, and finally return dx. Note that the inverse gradient calculation of the operator must be included in the inverse chain gradient calculation of the whole network
test
import numpy as np import mindspore.nn as nn import mindspore.context as context from mindspore import Tensor from mindspore.ops import operations as P from mindspore.ops import composite as C context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") class Net(nn.Cell): def __init__(self): super(Net, self).__init__() self.square = P.CusSquare() def construct(self, data): return self.square(data) def test_net(): X = Tensor (np) array ([[1.0, 4.0, Net = net () output = grad(net)(x) print("x: ", x) print("output: ", output) if __name__ == "__main__": test_net()Copy the codeThe output
Click to follow, the first time to learn about Huawei cloud fresh technology ~