PyLayer¶
Paddle 通过创建 PyLayer
子类的方式实现 Python 端自定义算子,这个子类必须遵守以下规则:
子类必须包含静态的
forward
和backward
函数,它们的第一个参数必须是 PyLayerContext,如果backward
的某个返回值在forward
中对应的Tensor
是需要梯度,这个返回值必须为Tensor
。backward
除了第一个参数以外,其他参数都是forward
函数的输出Tensor
的梯度,因此,backward
输入的Tensor
的数量必须等于forward
输出Tensor
的数量。如果你需在backward
中使用forward
的输入Tensor
,你可以将这些Tensor
输入到 PyLayerContext 的save_for_backward
方法,之后在backward
中使用这些Tensor
。backward
的输出可以是Tensor
或者list/tuple(Tensor)
,这些Tensor
是forward
输出Tensor
的梯度。因此,backward
的输出Tensor
的个数等于forward
输入Tensor
的个数。
构建完自定义算子后,通过 apply
运行算子。
代码示例¶
>>> import paddle
>>> from paddle.autograd import PyLayer
>>> class cus_tanh(PyLayer):
... @staticmethod
... def forward(ctx, x):
... y = paddle.tanh(x)
... # Pass tensors to backward.
... ctx.save_for_backward(y)
... return y
...
... @staticmethod
... def backward(ctx, dy):
... # Get the tensors passed by forward.
... y, = ctx.saved_tensor()
... grad = dy * (1 - paddle.square(y))
... return grad
>>> paddle.seed(2023)
>>> data = paddle.randn([2, 3], dtype="float64")
>>> data.stop_gradient = False
>>> z = cus_tanh.apply(data)
>>> z.mean().backward()
>>> print(data.grad)
Tensor(shape=[2, 3], dtype=float64, place=Place(cpu), stop_gradient=True,
[[0.16604150, 0.05858341, 0.14051214],
[0.15677770, 0.01564609, 0.02991660]])
方法¶
forward(ctx, *args, **kwargs)¶
forward
函数必须被子类重写,它的第一个参数是 PyLayerContext 的对象,其他输入参数的类型和数量任意。
参数
*args (tuple) - 自定义算子的输入
**kwargs (dict) - 自定义算子的输入
返回
Tensor 或至少包含一个 Tensor 的 list/tuple
代码示例
>>> import paddle
>>> from paddle.autograd import PyLayer
>>> class cus_tanh(PyLayer):
... @staticmethod
... def forward(ctx, x):
... y = paddle.tanh(x)
... # Pass tensors to backward.
... ctx.save_for_backward(y)
... return y
...
... @staticmethod
... def backward(ctx, dy):
... # Get the tensors passed by forward.
... y, = ctx.saved_tensor()
... grad = dy * (1 - paddle.square(y))
... return grad
backward(ctx, *args, **kwargs)¶
backward
函数的作用是计算梯度,它必须被子类重写,其第一个参数为 PyLayerContext 的对象,其他输入参数为 forward
输出 Tensor
的梯度。它的输出 Tensor
为 forward
输入 Tensor
的梯度。
参数
*args (tuple) -
forward
输出Tensor
的梯度。**kwargs (dict) -
forward
输出Tensor
的梯度。
返回
forward
输入Tensor
的梯度。
代码示例
>>> import paddle
>>> from paddle.autograd import PyLayer
>>> class cus_tanh(PyLayer):
... @staticmethod
... def forward(ctx, x):
... y = paddle.tanh(x)
... # Pass tensors to backward.
... ctx.save_for_backward(y)
... return y
...
... @staticmethod
... def backward(ctx, dy):
... # Get the tensors passed by forward.
... y, = ctx.saved_tensor()
... grad = dy * (1 - paddle.square(y))
... return grad
apply(cls, *args, **kwargs)¶
构建完自定义算子后,通过 apply
运行算子。
参数
*args (tuple) - 自定义算子的输入
**kwargs (dict) - 自定义算子的输入
返回
Tensor 或至少包含一个 Tensor 的 list/tuple
代码示例
>>> import paddle
>>> from paddle.autograd import PyLayer
>>> class cus_tanh(PyLayer):
... @staticmethod
... def forward(ctx, x):
... y = paddle.tanh(x)
... # Pass tensors to backward.
... ctx.save_for_backward(y)
... return y
...
... @staticmethod
... def backward(ctx, dy):
... # Get the tensors passed by forward.
... y, = ctx.saved_tensor()
... grad = dy * (1 - paddle.square(y))
... return grad
>>> paddle.seed(2023)
>>> data = paddle.randn([2, 3], dtype="float64")
>>> data.stop_gradient = False
>>> z = cus_tanh.apply(data)
>>> z.mean().backward()
>>> print(data.grad)
Tensor(shape=[2, 3], dtype=float64, place=Place(cpu), stop_gradient=True,
[[0.16604150, 0.05858341, 0.14051214],
[0.15677770, 0.01564609, 0.02991660]])