Support static graph code-gen for temporal_shift (#52686)

paddle/fluid/operators/temporal_shift_op.cc

-/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-   http://www.apache.org/licenses/LICENSE-2.0
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License. */
-
-#include "paddle/fluid/operators/temporal_shift_op.h"
-
-#include <memory>
-#include <string>
-#include <vector>
-
-#include "paddle/fluid/framework/infershape_utils.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/phi/core/infermeta_utils.h"
-#include "paddle/phi/infermeta/unary.h"
-
-namespace paddle {
-namespace operators {
-
-class TemporalShiftOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    return phi::KernelKey(OperatorWithKernel::IndicateVarDataType(ctx, "X"),
-                          ctx.GetPlace());
-  }
-};
-
-class TemporalShiftOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput("X",
-             "The input tensor of temporal shift operator. "
-             "This is a 4-D tensor with shape of [N*T, C, H, W] "
-             "or [N*T, H, W, C]. "
-             "While N is the batch size, T is the temporal segment "
-             "number, C is the channel number, H is the height of "
-             "features and W is the width of features. "
-             "The data type is float16, float32 and float64");
-    AddOutput("Out",
-              "The output tensor of temporal shift operator. "
-              "This is a 4-D tensor in the same shape with Input(X).");
-
-    AddAttr<int>("seg_num",
-                 "The temporal segment number, this should be a positive "
-                 "integer.");
-    AddAttr<float>(
-        "shift_ratio",
-        "The shift ratio of the channels, the first :attr:`shift_ratio` part "
-        "of channels will be shifted by -1 along the temporal dimension, "
-        "and the second :attr:`shift_ratio` part of channels will be shifted "
-        "by 1 along the temporal dimension. :attr:`shift_ratio` should be in "
-        "range [0, 0.5]. Default 0.25.")
-        .SetDefault(0.25);
-    AddAttr<std::string>(
-        "data_format",
-        "(string, default NCHW) Only used in "
-        "an optional string from: \"NHWC\", \"NCHW\". "
-        "Specify that the data format of the input and output data is "
-        "channel_first or channel_last.")
-        .SetDefault("NCHW");
-
-    AddComment(R"DOC(
-          This operator calculates the temporal shifting features for Input(X).
-
-          Input(X) should be in shape of [N*T, C, H, W] or [N*T, H, W, C], while
-          N is the batch size, T is the temporal segment number specified by
-          :attr:`seg_num`, C is the channel number, H and W is the height and
-          width of features.
-
-          Temporal Shifting is calculated as follows when data format is NCHW:
-
-          Step 1: Reshape Input(X) to [N, T, C, H, W].
-
-          Step 2: Pad 0 to reshaping result in the 2nd(T) dimension with
-          padding width as 1 on each side, padding result will be in shape
-          of [N, T+2, C, H, W].
-
-          Step 3: Assume :attr:`shift_ratio` is :math:`1/4`, slice padding
-          result as follows:
-
-          $$
-          slice1 = x[:, :T, :C/4, :, :]
-          $$
-          $$
-          slice2 = x[:, 2:T+2, C/4:C/2, :, :]
-          $$
-          $$
-          slice3 = x[:, 1:T+1, C/2:, :, :]
-          $$
-
-          Step 4: Concatenate three slices along the 3rd(C) dimension and
-          reshape result to [N*T, C, H, W].
-
-          For details of temporal shifting, please refer to paper:
-          `Temporal Shift Module <http://arxiv.org/abs/1811.08383>`_ .
-
-         )DOC");
-  }
-};
-
-class TemporalShiftOpGrad : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
- protected:
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    if (ctx->HasOutput(framework::GradVarName("X"))) {
-      ctx->SetOutputDim(framework::GradVarName("X"),
-                        ctx->GetInputDim(framework::GradVarName("Out")));
-    }
-  }
-
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    return phi::KernelKey(OperatorWithKernel::IndicateVarDataType(
-                              ctx, framework::GradVarName("Out")),
-                          ctx.GetPlace());
-  }
-};
-
-template <typename T>
-class TemporalShiftGradOpMaker : public framework::SingleGradOpMaker<T> {
- public:
-  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
-
- protected:
-  void Apply(GradOpPtr<T> op) const override {
-    op->SetType("temporal_shift_grad");
-    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
-    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
-    op->SetAttrMap(this->Attrs());
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-DECLARE_INFER_SHAPE_FUNCTOR(temporal_shift,
-                            TemporalShiftInferShapeFunctor,
-                            PD_INFER_META(phi::TemporalShiftInferMeta));
-REGISTER_OPERATOR(temporal_shift,
-                  ops::TemporalShiftOp,
-                  ops::TemporalShiftOpMaker,
-                  ops::TemporalShiftGradOpMaker<paddle::framework::OpDesc>,
-                  ops::TemporalShiftGradOpMaker<paddle::imperative::OpBase>,
-                  TemporalShiftInferShapeFunctor);
-REGISTER_OPERATOR(temporal_shift_grad, ops::TemporalShiftOpGrad);
-REGISTER_OP_CPU_KERNEL(temporal_shift,
-                       ops::TemporalShiftKernel<float>,
-                       ops::TemporalShiftKernel<double>);
-REGISTER_OP_CPU_KERNEL(temporal_shift_grad,
-                       ops::TemporalShiftGradKernel<float>,
-                       ops::TemporalShiftGradKernel<double>);

paddle/phi/api/yaml/backward.yaml

@@ -1863,6 +1863,17 @@
   inplace : (grad_x_grad_forward -> grad_out_forward_grad)
   optional : grad_out_new_grad, grad_out_grad_grad
 
+- backward_op : temporal_shift_grad
+  forward : temporal_shift(Tensor x, int seg_num, float shift_ratio = 0.25f, str data_format = "NCHW") -> Tensor(out)
+  args : (Tensor out_grad, int seg_num, float shift_ratio, str data_format)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [out_grad]
+  kernel :
+    func : temporal_shift_grad
+    data_type : out_grad
+
 - backward_op : thresholded_relu_grad
   forward : thresholded_relu (Tensor x, float threshold) -> Tensor(out)
   args : (Tensor x, Tensor out_grad, float threshold)

paddle/phi/api/yaml/legacy_backward.yaml

@@ -1031,16 +1031,6 @@
     data_type : out_grad
   optional : reserve_space
 
-- backward_op : temporal_shift_grad
-  forward : temporal_shift(Tensor x, int seg_num, float shift_ratio, str data_format_str) -> Tensor(out)
-  args : (Tensor out_grad, int seg_num, float shift_ratio, str data_format_str)
-  output : Tensor(x_grad)
-  infer_meta :
-    func : UnchangedInferMeta
-    param : [out_grad]
-  kernel :
-    func : temporal_shift_grad
-
 - backward_op : tile_double_grad
   forward : tile_grad (Tensor x, Tensor grad_out, IntArray repeat_times) -> Tensor(grad_x)
   args : (Tensor grad_x_grad, IntArray repeat_times)

paddle/phi/api/yaml/legacy_ops.yaml

@@ -1229,15 +1229,6 @@
   backward : sync_batch_norm_grad
   inplace : (mean -> mean_out), (variance -> variance_out)
 
-- op : temporal_shift
-  args : (Tensor x, int seg_num, float shift_ratio, str data_format_str)
-  output : Tensor
-  infer_meta :
-    func : TemporalShiftInferMeta
-  kernel :
-    func : temporal_shift
-  backward : temporal_shift_grad
-
 - op : tile
   args : (Tensor x, IntArray repeat_times = {})
   output : Tensor(out)

paddle/phi/api/yaml/op_compat.yaml

@@ -2347,3 +2347,10 @@
     x : X
   outputs :
     out : Out
+
+- op: temporal_shift
+  backward: temporal_shift_grad
+  inputs :
+    x : X
+  outputs :
+    out : Out

paddle/phi/api/yaml/ops.yaml

@@ -1885,6 +1885,16 @@
     func : tanh_shrink
   backward : tanh_shrink_grad
 
+- op : temporal_shift
+  args : (Tensor x, int seg_num, float shift_ratio = 0.25f, str data_format = "NCHW")
+  output : Tensor(out)
+  infer_meta :
+    func : TemporalShiftInferMeta
+  kernel :
+    func : temporal_shift
+    data_type : x
+  backward : temporal_shift_grad
+
 - op : thresholded_relu
   args : (Tensor x, float threshold = 1.0)
   output : Tensor

paddle/phi/ops/compat/temporal_shift_sig.cc

-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#include "paddle/phi/core/compat/op_utils.h"
-
-namespace phi {
-
-KernelSignature TemporalShiftOpArgumentMapping(
-    const ArgumentMappingContext& ctx) {
-  return KernelSignature("temporal_shift",
-                         {"X"},
-                         {"seg_num", "shift_ratio", "data_format"},
-                         {"Out"});
-}
-
-KernelSignature TemporalShiftGradOpArgumentMapping(
-    const ArgumentMappingContext& ctx) {
-  return KernelSignature("temporal_shift_grad",
-                         {"Out@GRAD"},
-                         {"seg_num", "shift_ratio", "data_format"},
-                         {"X@GRAD"});
-}
-
-}  // namespace phi
-
-PD_REGISTER_ARG_MAPPING_FN(temporal_shift, phi::TemporalShiftOpArgumentMapping);
-PD_REGISTER_ARG_MAPPING_FN(temporal_shift_grad,
-                           phi::TemporalShiftGradOpArgumentMapping);

python/paddle/nn/functional/extension.py

@@ -28,7 +28,6 @@ from ...fluid.framework import in_dygraph_mode
 from ...fluid.layer_helper import LayerHelper
 from ...framework import convert_np_dtype_to_dtype_, core
 from ...tensor.creation import assign
-from ...tensor.layer_function_generator import templatedoc
 
 __all__ = []
 
@@ -338,13 +337,44 @@ def gather_tree(ids, parents):
         return out
 
 
-@templatedoc()
 def temporal_shift(x, seg_num, shift_ratio=0.25, name=None, data_format="NCHW"):
     """
 
     **Temporal Shift Operator**
 
-    ${comment}
+    Calculate the temporal shifting features for Input(X).
+
+    Input(X) should be in shape of [N*T, C, H, W] or [N*T, H, W, C], while
+    N is the batch size, T is the temporal segment number specified by
+    :attr:`seg_num`, C is the channel number, H and W is the height and
+    width of features.
+
+    Temporal Shifting is calculated as follows when data format is NCHW:
+
+    Step 1: Reshape Input(X) to [N, T, C, H, W].
+
+    Step 2: Pad 0 to reshaping result in the 2nd(T) dimension with
+    padding width as 1 on each side, padding result will be in shape
+    of [N, T+2, C, H, W].
+
+    Step 3: Assume :attr:`shift_ratio` is :math:`1/4`, slice padding
+    result as follows:
+
+    $$
+    slice1 = x[:, :T, :C/4, :, :]
+    $$
+    $$
+    slice2 = x[:, 2:T+2, C/4:C/2, :, :]
+    $$
+    $$
+    slice3 = x[:, 1:T+1, C/2:, :, :]
+    $$
+
+    Step 4: Concatenate three slices along the 3rd(C) dimension and
+    reshape result to [N*T, C, H, W].
+
+    For details of temporal shifting, please refer to paper:
+    `Temporal Shift Module <http://arxiv.org/abs/1811.08383>`_ .
 
     Args:
         x(Tensor): ${x_comment}