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]])