1.动态图自定义Python算子
先说一下辅助用的,这个接口:
Paddle 通过 PyLayer
接口和PyLayerContext
接口支持动态图的Python端自定义OP。
自定义好的算子如何使用呢?需要特定的接口,这个接口就是PyLayer
PyLayer
接口描述如下:
class PyLayer:
@staticmethod
def forward(ctx, *args, **kwargs):
pass
@staticmethod
def backward(ctx, *args, **kwargs):
pass
@classmethod
def apply(cls, *args, **kwargs):
pass
其中,
--forward 是自定义Op的前向函数,必须被子类重写,它的第一个参数是 PyLayerContext 对象,其他输入参数的类型和数量任意。
--backward 是自定义Op的反向函数,必须被子类重写,其第一个参数为 PyLayerContext 对象,其他输入参数为forward输出Tensor的梯度。它的输出Tensor为forward输入Tensor的梯度。
--apply 是自定义Op的执行方法,构建完自定义Op后,通过apply运行Op。
PyLayerContext
接口描述如下:
class PyLayerContext:
def save_for_backward(self, *tensors):
pass
def saved_tensor(self):
pass
其中,
save_for_backward
用于暂存backward
需要的Tensor
,这个API只能被调用一次,且只能在forward
中调用。saved_tensor
获取被save_for_backward
暂存的Tensor
。
=====================================
========重点开始了===========
动态图自定义
如何编写动态图Python Op¶
以下以tanh为例,介绍如何利用 PyLayer
编写Python Op。
- 第一步:创建
PyLayer
子类并定义前向函数和反向函数
前向函数和反向函数均由Python编写,可以方便地使用Paddle相关API来实现一个自定义的OP。需要遵守以下规则:
forward
和backward
都是静态函数,它们的第一个参数是PyLayerContext
对象。backward
除了第一个参数以外,其他参数都是forward
函数的输出Tensor
的梯度,因此,backward
输入的Tensor
的数量必须等于forward
输出Tensor
的数量。如果您需在backward
中使用forward
中的Tensor
,您可以利用save_for_backward
和saved_tensor
这两个方法传递Tensor
。backward
的输出可以是Tensor
或者list/tuple(Tensor)
,这些Tensor
是forward
输入Tensor
的梯度。因此,backward
的输出Tensor
的个数等于forward输入Tensor
的个数。如果backward
的某个返回值(梯度)在forward
中对应的Tensor
的stop_gradient
属性为False
,这个返回值必须是Tensor
类型。
两步搞定,就像写函数一样,先写方法,然后调用。这是一个类似的比喻。
第一步:以正切算子为例,其他算子 比葫芦画瓢 应该都会吧。
import paddle
from paddle.autograd import PyLayer
# 通过创建`PyLayer`子类的方式实现动态图Python Op
class cus_tanh(PyLayer):
@staticmethod
def forward(ctx, x):
y = paddle.tanh(x)
# ctx 为PyLayerContext对象,可以把y从forward传递到backward。
ctx.save_for_backward(y)
return y
@staticmethod
# 因为forward只有一个输出,因此除了ctx外,backward只有一个输入。
def backward(ctx, dy):
# ctx 为PyLayerContext对象,saved_tensor获取在forward时暂存的y。
y, = ctx.saved_tensor()
# 调用Paddle API自定义反向计算
grad = dy * (1 - paddle.square(y))
# forward只有一个Tensor输入,因此,backward只有一个输出。
return grad
- 第二步:通过
apply
方法组建网络。apply
的输入为forward
中除了第一个参数(ctx
)以外的输入,apply
的输出即为forward
的输出。
data = paddle.randn([2, 3], dtype="float32")
data.stop_gradient = False
# 通过 apply运行这个Python算子
z = cus_tanh.apply(data)
z.mean().backward()
print(data.grad)
运行结果;
Tensor(shape=[2, 3], dtype=float32, place=Place(gpu:0), stop_gradient=False,
[[0.15810412, 0.16457692, 0.15256986],
[0.16439323, 0.10458803, 0.16661729]])
Process finished with exit code 0
下面是一些注意事项,特殊情况特殊分析
- 为了从
forward
到backward
传递信息,您可以在forward
中给PyLayerContext
添加临时属性,在backward
中读取这个属性。如果传递Tensor
推荐使用save_for_backward
和saved_tensor
,如果传递非Tensor
推荐使用添加临时属性的方式。
import paddle
from paddle.autograd import PyLayer
import numpy as np
class tanh(PyLayer):
@staticmethod
def forward(ctx, x1, func1, func2=paddle.square):
# 添加临时属性的方式传递func2
ctx.func = func2
y1 = func1(x1)
# 使用save_for_backward传递y1
ctx.save_for_backward(y1)
return y1
@staticmethod
def backward(ctx, dy1):
y1, = ctx.saved_tensor()
# 获取func2
re1 = dy1 * (1 - ctx.func(y1))
return re1
input1 = paddle.randn([2, 3]).astype("float64")
input2 = input1.detach().clone()
input1.stop_gradient = False
input2.stop_gradient = False
z = tanh.apply(x1=input1, func1=paddle.tanh)
- forward的输入和输出的类型任意,但是至少有一个输入和输出为
Tensor
类型。
# 错误示例
class cus_tanh(PyLayer):
@staticmethod
def forward(ctx, x1, x2):
y = x1+x2
# y.shape: 列表类型,非Tensor,输出至少包含一个Tensor
return y.shape
@staticmethod
def backward(ctx, dy):
return dy, dy
data = paddle.randn([2, 3], dtype="float32")
data.stop_gradient = False
# 由于forward输出没有Tensor引发报错
z, y_shape = cus_tanh.apply(data, data)
# 正确示例
class cus_tanh(PyLayer):
@staticmethod
def forward(ctx, x1, x2):
y = x1+x2
# y.shape: 列表类型,非Tensor
return y, y.shape
@staticmethod
def backward(ctx, dy):
# forward两个Tensor输入,因此,backward有两个输出。
return dy, dy
data = paddle.randn([2, 3], dtype="float32")
data.stop_gradient = False
z, y_shape = cus_tanh.apply(data, data)
z.mean().backward()
print(data.grad)
- 如果forward的某个输入为
Tensor
且stop_gredient = True
,则在backward
中与其对应的返回值应为None
。
class cus_tanh(PyLayer):
@staticmethod
def forward(ctx, x1, x2):
y = x1+x2
return y
@staticmethod
def backward(ctx, dy):
# x2.stop_gradient=True,其对应梯度需要返回None
return dy, None
data1 = paddle.randn([2, 3], dtype="float32")
data1.stop_gradient = False
data2 = paddle.randn([2, 3], dtype="float32")
z = cus_tanh.apply(data1, data2)
fake_loss = z.mean()
fake_loss.backward()
print(data1.grad)
- 如果forward的所有输入
Tensor
都是stop_gredient = True
的,则backward
不会被执行。
class cus_tanh(PyLayer):
@staticmethod
def forward(ctx, x1, x2):
y = x1+x2
return y
@staticmethod
def backward(ctx, dy):
return dy, None
data1 = paddle.randn([2, 3], dtype="float32")
data2 = paddle.randn([2, 3], dtype="float32")
z = cus_tanh.apply(data1, data2)
fake_loss = z.mean()
fake_loss.backward()
# 因为data1.stop_gradient = True、data2.stop_gradient = True,所以backward不会被执行。
print(data1.grad is None)
自于静态的自定义算子,我不太感兴趣
可以参考官网