|
| 1 | + |
| 2 | +#为 Paddle 新增 apply API 设计文档 |
| 3 | + |
| 4 | +| API名称 | tensor.apply | |
| 5 | +|---|----------------------------------| |
| 6 | +| 提交作者 | yangguohao | |
| 7 | +| 提交时间 | 2023-11-26 | |
| 8 | +| 版本号 | V1.0 | |
| 9 | +| 依赖飞桨版本 | develop | |
| 10 | +| 文件名 | 20231126_api_design_for_apply.md | |
| 11 | + |
| 12 | +# 一、概述 |
| 13 | + |
| 14 | +## 1、相关背景 |
| 15 | + |
| 16 | +为了增强Paddle框架的功能并与其他框架保持竞争力,我们决定为Tensor添加一个新的API,它允许用户应用Python函数到Tensor的每个元素。 |
| 17 | + |
| 18 | +## 2、功能目标 |
| 19 | + |
| 20 | +- 为Tensor添加 `apply(callable)` 方法,返回新的Tensor,存放计算结果。 |
| 21 | +- 为Tensor添加 `apply_(callable)` 方法,inplace修改输入Tensor。 |
| 22 | + |
| 23 | +## 3、意义 |
| 24 | + |
| 25 | +该功能将提供更大的灵活性,使开发人员能够轻松地应用自定义函数到Tensor的每个元素,从而实现更复杂的操作。 |
| 26 | + |
| 27 | +# 二、飞桨现状 |
| 28 | + |
| 29 | +Paddle当前不支持此功能。 |
| 30 | + |
| 31 | +# 三、业内方案调研 |
| 32 | + |
| 33 | +TensorFlow、PyTorch 和 NumPy 都提供了允许用户对数组或张量元素应用函数的功能。 |
| 34 | + |
| 35 | +- **TensorFlow**: TensorFlow 提供了 `tf.map_fn` 方法,使用户能够对张量的元素应用函数, 函数的范围不限于 python 函数还包括 tf 函数,同时支持 backpropagation。 |
| 36 | +``` |
| 37 | +tf.map_fn( |
| 38 | + fn, |
| 39 | + elems, |
| 40 | + dtype=None, |
| 41 | + parallel_iterations=None, |
| 42 | + back_prop=True, |
| 43 | + swap_memory=False, |
| 44 | + infer_shape=True, |
| 45 | + name=None, |
| 46 | + fn_output_signature=None |
| 47 | +) |
| 48 | +``` |
| 49 | + |
| 50 | +- **PyTorch**: PyTorch 通过 `torch.Tensor.apply_` 方法,允许用户在 inplace 应用函数到张量的元素, 无法应用 torch 函数。但需要注意,此方法仅限于CPU张量,并且性能较差,运行的速度较慢。 |
| 51 | +``` |
| 52 | +Tensor.apply_(callable) |
| 53 | +``` |
| 54 | + |
| 55 | +# 四、对比分析 |
| 56 | + |
| 57 | +- **TensorFlow**: `tf.map_fn` 方法提供了强大的功能,但其接口相对复杂,不太适合初学者。 |
| 58 | + |
| 59 | +- **PyTorch**: 虽然 `torch.Tensor.apply_` 方法提供了所需的功能,但由于其局限性,它在新代码中不被推荐使用。 |
| 60 | + |
| 61 | +Paddle 的 `apply` 和 `apply_` 方法的设计目标以 pytorch 的实现功能类似。 |
| 62 | + |
| 63 | +# 五、设计思路与实现方案 |
| 64 | + |
| 65 | +## 命名与参数设计 |
| 66 | + |
| 67 | +tensor.apply(callable) |
| 68 | +tensor.apply_(callable) |
| 69 | + |
| 70 | +参数: |
| 71 | +- `callable`: 用户提供的Python函数,将应用于Tensor的每个元素。 |
| 72 | + |
| 73 | +## 底层OP设计 |
| 74 | + |
| 75 | +无需实现底层 OP 及 Kernel |
| 76 | + |
| 77 | +## API实现方案 |
| 78 | +采取与 pytorch 相似的设计,以下对应的实现其本质都是在 C++ 端调用 python 函数,可以看作是 python 函数 f 作用在张量 x 上,即 python 中 f(x) 的调用。因此该方案对 GPU 上的 tensor 也同样支持。 |
| 79 | + |
| 80 | +### 动态图 |
| 81 | +在 tensor_patch_methods 内增加 apply 和 apply_ 的两个 api,对于求梯度的 tensor 则无法使用 apply 并报错 |
| 82 | +``` |
| 83 | +@framework.dygraph_only |
| 84 | +def apply_(self, func): |
| 85 | + """ |
| 86 | + Inplace apply the python function to the tensor. |
| 87 | + Returns: |
| 88 | + None |
| 89 | + Examples: |
| 90 | + .. code-block:: python |
| 91 | + >>> import paddle |
| 92 | + >>> x = paddle.to_tensor(5.) |
| 93 | + >>> f = lambda x: 3*x+2 |
| 94 | + >>> x.apply_(f) |
| 95 | + >>> print(x) # 17 |
| 96 | + """ |
| 97 | + if self.stop_gradient is True: |
| 98 | + raise RuntimeError( |
| 99 | + "Cannot apply function on a tensor that stop gradient." |
| 100 | + ) |
| 101 | + self._apply_(func) |
| 102 | +
|
| 103 | +def apply(self, func): |
| 104 | + """ |
| 105 | + Apply the python function to the tensor. |
| 106 | + Returns: |
| 107 | + None |
| 108 | + Examples: |
| 109 | + .. code-block:: python |
| 110 | + >>> import paddle |
| 111 | + >>> x = paddle.to_tensor(5.) |
| 112 | + >>> f = lambda x: 3*x+2 |
| 113 | + >>> y = x.apply(f) |
| 114 | + >>> print(y) # 17 |
| 115 | + >>> print(x) # 5 |
| 116 | + """ |
| 117 | + if self.stop_gradient is True: |
| 118 | + raise RuntimeError( |
| 119 | + "Cannot apply function on a tensor that stop gradient." |
| 120 | + ) |
| 121 | + return self._apply(func) |
| 122 | +``` |
| 123 | + |
| 124 | +C++ 端则是在 eager_method.cc 中增加相应的逻辑,可以复用 PyTensorHook 来对 tensor 作用 callable python function。 |
| 125 | + |
| 126 | +``` |
| 127 | +static PyObject* tensor_apply(TensorObject* self, |
| 128 | + PyObject* args, |
| 129 | + PyObject* kwargs) { |
| 130 | + EAGER_TRY |
| 131 | + PyObject* apply_func = PyTuple_GET_ITEM(args, 0); |
| 132 | + PyTensorHook func = PyTensorHook(apply_func); |
| 133 | + paddle::Tensor out = func(self->tensor); |
| 134 | + return ToPyObject(out); |
| 135 | + EAGER_CATCH_AND_THROW_RETURN_NULL |
| 136 | +} |
| 137 | +
|
| 138 | +static PyObject* tensor_apply_(TensorObject* self, |
| 139 | + PyObject* args, |
| 140 | + PyObject* kwargs) { |
| 141 | + EAGER_TRY |
| 142 | + PyObject* apply_func = PyTuple_GET_ITEM(args, 0); |
| 143 | + PyTensorHook func = PyTensorHook(apply_func); |
| 144 | + paddle::Tensor out = func(self->tensor); |
| 145 | + self->tensor.set_impl(out.impl()); |
| 146 | + RETURN_PY_NONE |
| 147 | + EAGER_CATCH_AND_THROW_RETURN_NULL |
| 148 | +} |
| 149 | +``` |
| 150 | + |
| 151 | +### 静态图 |
| 152 | + |
| 153 | +Legacy IR 在 python/paddle/base/framework.py 的 Variable 下添加 apply 函数 |
| 154 | +``` |
| 155 | +def apply(self, func): |
| 156 | + if not self.stop_gradient: |
| 157 | + raise RuntimeError( |
| 158 | + "Cannot apply function on a tensor that required gradient." |
| 159 | + ) |
| 160 | + try: |
| 161 | + return func(self) |
| 162 | + except: |
| 163 | + raise ValueError(f"The PyFunc {func.__name__} could not be applied") |
| 164 | +``` |
| 165 | + |
| 166 | +PIR 则在 paddle/fluid/pybind/pir.cc 中实现 apply 并 bind 到 Value 上 |
| 167 | + |
| 168 | +``` |
| 169 | +pir::OpResult apply(Value self, py::object func) { |
| 170 | + py::gil_scoped_acquire gil; |
| 171 | + PyObject *py_func = func.release().ptr(); |
| 172 | + Py_INCREF(py_func); |
| 173 | + PyObject *res = nullptr; |
| 174 | + try { |
| 175 | + py::object obj = py::cast(self); |
| 176 | + PyObject *tmp_self = obj.release().ptr(); |
| 177 | + Py_INCREF(tmp_self); |
| 178 | + res = PyObject_CallFunctionObjArgs(py_func, tmp_self, nullptr); |
| 179 | + Py_DECREF(tmp_self); |
| 180 | + } catch (std::exception &e) { |
| 181 | + PADDLE_THROW(phi::errors::Unavailable( |
| 182 | + "Hook function of Tensor raises an exception: %s.", e.what())); |
| 183 | + } catch (...) { |
| 184 | + PADDLE_THROW(phi::errors::Fatal( |
| 185 | + "Hook function of Tensor raises an unknown exception.")); |
| 186 | + } |
| 187 | + if (res == Py_None) { |
| 188 | + return self.dyn_cast<OpResult>(); |
| 189 | + } |
| 190 | + auto out = CastPyArg2Value(res, "apply", 0); |
| 191 | + Py_DECREF(py_func); |
| 192 | + Py_DECREF(res); |
| 193 | + return out.dyn_cast<OpResult>(); |
| 194 | +} |
| 195 | +``` |
| 196 | + |
| 197 | +# 六、测试和验收的考量 |
| 198 | + |
| 199 | +- 单测代码,apply api 测试 |
| 200 | +- inplace 测试 |
| 201 | +- 对 stop_gradient=False 的 tensor 报错测试 |
| 202 | +- 动转静下测试 |
| 203 | + |
| 204 | +# 七、可行性分析和排期规划 |
| 205 | + |
| 206 | +已大部分完成代码开发工作,工期上可以满足在当前版本周期内开发完成。 |
| 207 | + |
| 208 | +# 八、影响面 |
| 209 | + |
| 210 | +为独立新增API,对其他模块没有影响。 |
| 211 | + |
| 212 | +# 名词解释 |
| 213 | +无 |
0 commit comments