README.md · 薛鹏/tf_custom_ascendc_op

`tf_custom_ascendc_op`目录包含了一个Tensorflow调用AscendC自定义算子的开发示例。
> 针对AscendC自定义算子的tf适配动作，会在后续的AscendC开发手册中描述，该工程为临时示例。

## 适用场景说明
> 同时满足以下两个条件的场景，适合使用本示例：
- 已经基于AscendC实现了CANN上的算子，希望在Tensorflow中调用该算子。
- Tensorflow中没有对应的算子，需要自行实现一个tf算子，并在tf脚本中调用。

针对上述场景，需要进行两部分的适配开发：
- 在Tensorflow中实现自定义算子，并在tf脚本中调用。
- 实现将Tensorflow算子原型转换为AscendC算子原型的插件。

## 示例说明
- 示例不涉及AscendC算子开发过程，仅展示如何在Tensorflow中调用AscendC算子。
- 示例中将CANN中已有的算子AssignAdd视为已经完成开发的自定义AscendC算子，并假设其没有对应的Tensorflow算子。
CANN中`AssignAdd`算子的原型如下：
```c++
REG_OP(AssignAdd)
.INPUT(ref, TensorType::BasicType())
.INPUT(value, TensorType::BasicType())
.OUTPUT(ref, TensorType::BasicType())
.ATTR(use_locking, Bool, false)
```
AssignAdd算子的功能是将value累加到ref上，即ref = ref + value。
为了在Tensorflow中调用`AssignAdd`算子，需要实现一个对应的Tensorflow算子`CustomAssignaddVariable`，其原型如下：
> 我们忽略了`use_locking`属性，以忽略不重要的细节。
```c++
REGISTER_OP("CustomAssignaddVariable")
    .Input("var: resource")
    .Input("value: float")
    .SetIsStateful();
```
> 可以看到我们将需要修改的变量数据类型修改为了`resource`，如果您熟悉Tensorflow1.x版本，您可能预期自定义的Tensorflow算子原型应该是：
```c++
REGISTER_OP("RefZeroOut")
    .Input("var: Ref(float)")
    .Input("value: float")
    .Output("output: Ref(float)")
    .SetIsStateful();
```
> 这是因为Tensorflow2.0版本后默认不再支持`Ref`类型的输入，变量统一使用`resource`类型表达，因此我们采用resource而非Ref类型输入，否则将无法在Tensorflow2.x版本中正常工作（当然，在小心地满足诸多限制后的某些场景下可以工作）。

了解上述背景后，我们将演示如何在Tensorflow中实现`CustomAssignaddVariable`自定义算子，以及如何实现将`CustomAssignaddVariable`算子原型转换为`AssignAdd`算子原型的插件。并最终在tf脚本中完成对CANN中AssignAdd算子的调用。

## 目录结构说明
```
├── LICENSE
├── README.md # 说明文档
├── tf_custom_op # 自定义tensorflow CustomAssignaddVariable算子示例
└── tf_parser # tf CustomAssignaddVariable算子原型转换为AscendC算子原型插件开发示例
```
需要您阅读tf_custom_op和tf_parser目录下的README.md文件，以了解更多示例细节。本文档仅提供了示例的整体说明。

## 使用示例
1. 首先，您需要按照tf_custom_op目录下的README.md文件，完成CustomAssignaddVariable算子的实现。完成后，您将得到一个`libcustom_ops.so`文件。
2. 然后，您需要按照tf_parser目录下的README.md文件，完成CustomAssignaddVariable算子原型转换为AssignAdd算子原型的插件开发。并按照对应的说明完成插件的编译和安装。

完成上述两步后，您可以在tf脚本中通过CustomAssignaddVariable的python API调用CANN中的AssignAdd算子了，如同原生的tf算子一般。
> 注意，您仍然需要完成Tensorflow到NPU上的迁移动作。

### Tensorflow1.x使用示例（[session模式NPU迁移](https://www.hiascend.com/document/detail/zh/canncommercial/700/modeldev/tfmigr1/tfmigr1_000017.html)）
```python
import os
import tensorflow as tf
tf.enable_resource_variables()
from npu_bridge.npu_init import *

custom_op_lib = tf.load_op_library('./libcustom_ops.so')

x = tf.Variable([[1., 2.], [3., 4.]])
value = tf.constant([[1., 2.], [3., 4.]])
y = custom_op_lib.custom_assignadd_variable(x.handle, value)

config = tf.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
config.graph_options.rewrite_options.memory_optimization = RewriterConfig.OFF

with tf.Session(config=config) as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(x))
    sess.run(y)
    print(sess.run(x))
```

### Tensorflow2.x使用示例（[compat session模式NPU迁移](https://www.hiascend.com/document/detail/zh/canncommercial/700/modeldev/tfmigr2/tfmigr2_000020.html)）
```python
import os
import tensorflow as tf
import npu_device
from npu_device.compat.v1.npu_init import *
npu_device.compat.enable_v1()

custom_op_lib = tf.load_op_library('./libcustom_ops.so')

x = tf.Variable([[1., 2.], [3., 4.]])
value = tf.constant([[1., 2.], [3., 4.]])
y = custom_op_lib.custom_assignadd_variable(x.handle, value)

config = tf.compat.v1.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
config.graph_options.rewrite_options.memory_optimization = RewriterConfig.OFF

with tf.compat.v1.Session(config=config) as sess:
    sess.run(tf.compat.v1.global_variables_initializer())
    print(sess.run(x))
    sess.run(y)
    print(sess.run(x))
```

薛鹏/tf_custom_ascendc_op

简介

发行版

贡献者

近期动态

薛鹏/tf_custom_ascendc_op .gitee-modal { width: 500px !important; }

简介

发行版

贡献者

近期动态

搜索帮助

薛鹏/tf_custom_ascendc_op