Add paddle.lerp API to do a linear interpolation

paddle/fluid/operators/lerp_op.cc

@@ -0,0 +1,146 @@
+// Copyright (c) 2021 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/fluid/operators/lerp_op.h"
+
+namespace paddle {
+namespace operators {
+
+class LerpOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "lerp");
+    OP_INOUT_CHECK(ctx->HasInput("Y"), "Input", "Y", "lerp");
+    OP_INOUT_CHECK(ctx->HasInput("Weight"), "Input", "Weight", "lerp");
+    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "lerp");
+
+    auto x_dims = ctx->GetInputDim("X");
+    auto y_dims = ctx->GetInputDim("Y");
+    auto w_dims = ctx->GetInputDim("Weight");
+    framework::DDim out_dims;
+    out_dims = GetOutputDims(x_dims, y_dims);
+    if (w_dims.size() > 1 || w_dims[0] != 1) {
+      out_dims = GetOutputDims(out_dims, w_dims);
+    }
+
+    ctx->SetOutputDim("Out", out_dims);
+    ctx->ShareLoD("X", /*->*/ "Out");
+  }
+
+ private:
+  framework::DDim GetOutputDims(const framework::DDim& s_dims,
+                                const framework::DDim& l_dims) const {
+    if (s_dims.size() > l_dims.size()) {
+      return GetOutputDims(l_dims, s_dims);
+    }
+    std::vector<int64_t> shapes = framework::vectorize<int64_t>(l_dims);
+    for (int i = s_dims.size() - 1, j = l_dims.size() - 1; i >= 0; --i, --j) {
+      int64_t s = s_dims[i];
+      int64_t l = l_dims[j];
+      if (s != l) {
+        if (l == 1) {
+          shapes[j] = s;
+        } else if (s != 1) {
+          PADDLE_THROW(platform::errors::InvalidArgument(
+              "The shape of tensor a %s:%d must match shape of tensor b "
+              "%s:%d.",
+              s_dims.to_str(), i, l_dims.to_str(), j));
+        }
+      }
+    }
+    return framework::make_ddim(shapes);
+  }
+};
+
+class LerpOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X", "(Tensor), The input tensor of lerp op.");
+    AddInput("Y", "(Tensor), The input tensor of lerp op.");
+    AddInput("Weight", "(Tensor, optional), The input tensor of lerp op.");
+    AddOutput("Out", "(Tensor), The output tensor of lerp op.");
+    AddComment(R"DOC(
+Lerp Operator.
+
+This operator is used to do a linear interpolation of input $X$ and $Y$ with $Weight$.
+
+The equation is:
+
+$$Out = X + Weight * (Y - X)$$
+
+Both the input $X$ and $Y$ can carry the LoD (Level of Details) information,
+or not. But the output only shares the LoD information with input $X$.
+
+)DOC");
+  }
+};
+
+class LerpGradOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    if (ctx->HasOutput(framework::GradVarName("X"))) {
+      ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
+    }
+    if (ctx->HasOutput(framework::GradVarName("Y"))) {
+      ctx->SetOutputDim(framework::GradVarName("Y"), ctx->GetInputDim("Y"));
+    }
+  }
+};
+
+template <typename T>
+class LerpOpGradMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+  void Apply(GradOpPtr<T> op) const override {
+    op->SetType("lerp_grad");
+    op->SetInput("X", this->Input("X"));
+    op->SetInput("Y", this->Input("Y"));
+    op->SetInput("Weight", this->Input("Weight"));
+    op->SetInput("Out", this->Output("Out"));
+    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
+    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
+    op->SetOutput(framework::GradVarName("Y"), this->InputGrad("Y"));
+    op->SetAttrMap(this->Attrs());
+  }
+};
+
+DECLARE_INPLACE_OP_INFERER(LerpInplaceInferer, {"X", "Out"});
+
+}  // namespace operators
+}  // namespace paddle
+
+REGISTER_OPERATOR(
+    lerp, paddle::operators::LerpOp, paddle::operators::LerpOpMaker,
+    paddle::operators::LerpOpGradMaker<paddle::framework::OpDesc>,
+    paddle::operators::LerpOpGradMaker<paddle::imperative::OpBase>,
+    paddle::operators::LerpInplaceInferer);
+
+REGISTER_OPERATOR(lerp_grad, paddle::operators::LerpGradOp);
+
+REGISTER_OP_CPU_KERNEL(
+    lerp,
+    paddle::operators::LerpKernel<paddle::platform::CPUDeviceContext, float>,
+    paddle::operators::LerpKernel<paddle::platform::CPUDeviceContext, double>);
+
+REGISTER_OP_CPU_KERNEL(
+    lerp_grad,
+    paddle::operators::LerpGradKernel<paddle::platform::CPUDeviceContext,
+                                      float>,
+    paddle::operators::LerpGradKernel<paddle::platform::CPUDeviceContext,
+                                      double>);

paddle/fluid/operators/lerp_op.cu

@@ -0,0 +1,27 @@
+// Copyright (c) 2021 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/fluid/operators/lerp_op.h"
+
+REGISTER_OP_CUDA_KERNEL(
+    lerp,
+    paddle::operators::LerpKernel<paddle::platform::CUDADeviceContext, float>,
+    paddle::operators::LerpKernel<paddle::platform::CUDADeviceContext, double>);
+
+REGISTER_OP_CUDA_KERNEL(
+    lerp_grad,
+    paddle::operators::LerpGradKernel<paddle::platform::CUDADeviceContext,
+                                      float>,
+    paddle::operators::LerpGradKernel<paddle::platform::CUDADeviceContext,
+                                      double>);

paddle/fluid/operators/lerp_op.h

@@ -0,0 +1,217 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+
+#ifdef _WIN32
+#ifndef NOMINMAX
+#define NOMINMAX  // msvc max/min macro conflict with std::min/max
+#endif
+#endif
+
+namespace paddle {
+namespace operators {
+
+static framework::DDim ExtendDims2Rank(const framework::DDim& in_dims,
+                                       int rank) {
+  if (in_dims.size() == rank) {
+    return in_dims;
+  }
+  std::vector<int64_t> shapes(rank, 1);
+  for (int i = in_dims.size() - 1, j = rank - 1; i >= 0; --i, --j) {
+    shapes[j] = in_dims[i];
+  }
+  return framework::make_ddim(shapes);
+}
+
+template <size_t D>
+static void GetBroadcastDims(const framework::DDim& in_dims,
+                             const framework::DDim& out_dims,
+                             Eigen::DSizes<int, D>* bcast_dims) {
+  for (size_t i = 0; i < D; ++i) {
+    if (in_dims[i] == out_dims[i]) {
+      (*bcast_dims)[i] = 1;
+    } else {
+      (*bcast_dims)[i] = std::max(in_dims[i], out_dims[i]);
+    }
+  }
+}
+
+template <typename DeviceContext, typename T, size_t D>
+static void LerpFunction(const framework::ExecutionContext& ctx) {
+  auto x = ctx.Input<framework::Tensor>("X");
+  auto y = ctx.Input<framework::Tensor>("Y");
+  auto w = ctx.Input<framework::Tensor>("Weight");
+  auto out = ctx.Output<framework::Tensor>("Out");
+  out->mutable_data<T>(ctx.GetPlace());
+
+  auto out_dims = out->dims();
+  auto x_dims = ExtendDims2Rank(x->dims(), D);
+  auto y_dims = ExtendDims2Rank(y->dims(), D);
+  auto w_dims = ExtendDims2Rank(w->dims(), D);
+  Eigen::DSizes<int, D> x_bcast_dims;
+  Eigen::DSizes<int, D> y_bcast_dims;
+  Eigen::DSizes<int, D> w_bcast_dims;
+  GetBroadcastDims<D>(x_dims, out_dims, &x_bcast_dims);
+  GetBroadcastDims<D>(y_dims, out_dims, &y_bcast_dims);
+  GetBroadcastDims<D>(w_dims, out_dims, &w_bcast_dims);
+
+  auto eigen_x = framework::EigenTensor<T, D>::From(*x, x_dims);
+  auto eigen_y = framework::EigenTensor<T, D>::From(*y, y_dims);
+  auto eigen_w = framework::EigenTensor<T, D>::From(*w, w_dims);
+  auto eigen_out = framework::EigenTensor<T, D>::From(*out);
+
+  auto& place = *ctx.template device_context<DeviceContext>().eigen_device();
+  eigen_out.device(place) =
+      eigen_x.broadcast(x_bcast_dims) +
+      eigen_w.broadcast(w_bcast_dims) *
+          (eigen_y.broadcast(y_bcast_dims) - eigen_x.broadcast(x_bcast_dims));
+}
+
+template <typename DeviceContext, typename T, size_t D>
+static void LerpGradFunction(const framework::ExecutionContext& ctx) {
+  auto w = ctx.Input<framework::Tensor>("Weight");
+  auto dout = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
+  auto dx = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
+  auto dy = ctx.Output<framework::Tensor>(framework::GradVarName("Y"));
+
+  auto dout_dims = dout->dims();
+  auto dx_dims = ExtendDims2Rank(dx->dims(), D);
+  auto dy_dims = ExtendDims2Rank(dy->dims(), D);
+  auto w_dims = ExtendDims2Rank(w->dims(), D);
+  Eigen::DSizes<int, D> dx_bcast_dims;
+  Eigen::DSizes<int, D> dy_bcast_dims;
+  Eigen::DSizes<int, D> w_bcast_dims;
+  GetBroadcastDims<D>(dx_dims, dout_dims, &dx_bcast_dims);
+  GetBroadcastDims<D>(dy_dims, dout_dims, &dy_bcast_dims);
+  GetBroadcastDims<D>(w_dims, dout_dims, &w_bcast_dims);
+
+  auto eigen_w = framework::EigenTensor<T, D>::From(*w, w_dims);
+  auto eigen_dout = framework::EigenTensor<T, D>::From(*dout);
+
+  Eigen::DSizes<int, D * 2> dx_reshape_dims;
+  Eigen::DSizes<int, D * 2> dy_reshape_dims;
+  Eigen::DSizes<int, D> reduce_dims;
+  for (int i = 0; i < dout_dims.size(); ++i) {
+    dx_reshape_dims[2 * i] = dx_bcast_dims[i];
+    dx_reshape_dims[2 * i + 1] = dx_dims[i];
+    dy_reshape_dims[2 * i] = dy_bcast_dims[i];
+    dy_reshape_dims[2 * i + 1] = dy_dims[i];
+    reduce_dims[i] = 2 * i;
+  }
+
+  auto& place = *ctx.template device_context<DeviceContext>().eigen_device();
+
+  if (dx) {
+    dx->mutable_data<T>(ctx.GetPlace());
+    auto eigen_dx = framework::EigenTensor<T, D>::From(*dx, dx_dims);
+    auto eigen_expr = (1 - eigen_w.broadcast(w_bcast_dims)) * eigen_dout;
+    eigen_dx.device(place) = eigen_expr.reshape(dx_reshape_dims)
+                                 .sum(reduce_dims)
+                                 .reshape(eigen_dx.dimensions());
+  }
+  if (dy) {
+    dy->mutable_data<T>(ctx.GetPlace());
+    auto eigen_dy = framework::EigenTensor<T, D>::From(*dy, dy_dims);
+    auto eigen_expr = eigen_w.broadcast(w_bcast_dims) * eigen_dout;
+    eigen_dy.device(place) = eigen_expr.reshape(dy_reshape_dims)
+                                 .sum(reduce_dims)
+                                 .reshape(eigen_dy.dimensions());
+  }
+}
+
+template <typename DeviceContext, typename T>
+class LerpKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    int rank = ctx.Output<framework::Tensor>("Out")->dims().size();
+    PADDLE_ENFORCE_GE(
+        rank, 1,
+        platform::errors::InvalidArgument(
+            "The number of dimensions for LerpOp must be "
+            "greater than or equal to 1, but the value received is %d.",
+            rank));
+    PADDLE_ENFORCE_LE(
+        rank, 6, platform::errors::InvalidArgument(
+                     "The number of dimensions for LerpOp must be "
+                     "less than or equal to 6, but the value received is %d.",
+                     rank));
+    switch (rank) {
+      case 1:
+        LerpFunction<DeviceContext, T, 1>(ctx);
+        break;
+      case 2:
+        LerpFunction<DeviceContext, T, 2>(ctx);
+        break;
+      case 3:
+        LerpFunction<DeviceContext, T, 3>(ctx);
+        break;
+      case 4:
+        LerpFunction<DeviceContext, T, 4>(ctx);
+        break;
+      case 5:
+        LerpFunction<DeviceContext, T, 5>(ctx);
+        break;
+      case 6:
+        LerpFunction<DeviceContext, T, 6>(ctx);
+        break;
+    }
+  }
+};
+
+template <typename DeviceContext, typename T>
+class LerpGradKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    int rank = ctx.Input<framework::Tensor>(framework::GradVarName("Out"))
+                   ->dims()
+                   .size();
+    PADDLE_ENFORCE_GE(
+        rank, 1,
+        platform::errors::InvalidArgument(
+            "The number of dimensions for LerpGradOp must be "
+            "greater than or equal to 1, but the value received is %d.",
+            rank));
+    PADDLE_ENFORCE_LE(
+        rank, 6, platform::errors::InvalidArgument(
+                     "The number of dimensions for LerpGradOp must be "
+                     "less than or equal to 6, but the value received is %d.",
+                     rank));
+    switch (rank) {
+      case 1:
+        LerpGradFunction<DeviceContext, T, 1>(ctx);
+        break;
+      case 2:
+        LerpGradFunction<DeviceContext, T, 2>(ctx);
+        break;
+      case 3:
+        LerpGradFunction<DeviceContext, T, 3>(ctx);
+        break;
+      case 4:
+        LerpGradFunction<DeviceContext, T, 4>(ctx);
+        break;
+      case 5:
+        LerpGradFunction<DeviceContext, T, 5>(ctx);
+        break;
+      case 6:
+        LerpGradFunction<DeviceContext, T, 6>(ctx);
+        break;
+    }
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle