Skip to content

Commit c7312bd

Browse files
authored
【Hackathon No.7】为 Paddle 新增 apply API rfcs (#761)
* add rfcs * add static apply api * fix rfcs
1 parent 289b309 commit c7312bd

File tree

1 file changed

+213
-0
lines changed

1 file changed

+213
-0
lines changed
Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
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

Comments
 (0)