first commit

test_tipc/supplementary/custom_op/custom_relu_op.cc

@@ -0,0 +1,106 @@
+// 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.
+
+// reference from :
+// https://github.com/PaddlePaddle/Paddle-Inference-Demo/blob/master/python/custom-operator/custom_relu_op.cc
+#include <iostream>
+#include <vector>
+
+#include "paddle/extension.h"
+
+template <typename data_t>
+void relu_cpu_forward_kernel(const data_t *x_data, data_t *out_data,
+                             int64_t x_numel) {
+  for (int i = 0; i < x_numel; ++i) {
+    out_data[i] = std::max(static_cast<data_t>(0.), x_data[i]);
+  }
+}
+
+template <typename data_t>
+void relu_cpu_backward_kernel(const data_t *grad_out_data,
+                              const data_t *out_data, data_t *grad_x_data,
+                              int64_t out_numel) {
+  for (int i = 0; i < out_numel; ++i) {
+    grad_x_data[i] =
+        grad_out_data[i] * (out_data[i] > static_cast<data_t>(0) ? 1. : 0.);
+  }
+}
+
+std::vector<paddle::Tensor> relu_cpu_forward(const paddle::Tensor &x) {
+  auto out = paddle::Tensor(paddle::PlaceType::kCPU);
+
+  out.reshape(x.shape());
+  PD_DISPATCH_FLOATING_TYPES(
+      x.type(), "relu_cpu_forward", ([&] {
+        relu_cpu_forward_kernel<data_t>(
+            x.data<data_t>(), out.mutable_data<data_t>(x.place()), x.size());
+      }));
+
+  return {out};
+}
+
+std::vector<paddle::Tensor> relu_cpu_backward(const paddle::Tensor &x,
+                                              const paddle::Tensor &out,
+                                              const paddle::Tensor &grad_out) {
+  auto grad_x = paddle::Tensor(paddle::PlaceType::kCPU);
+  grad_x.reshape(x.shape());
+
+  PD_DISPATCH_FLOATING_TYPES(out.type(), "relu_cpu_backward", ([&] {
+                               relu_cpu_backward_kernel<data_t>(
+                                   grad_out.data<data_t>(), out.data<data_t>(),
+                                   grad_x.mutable_data<data_t>(x.place()),
+                                   out.size());
+                             }));
+
+  return {grad_x};
+}
+
+std::vector<paddle::Tensor> relu_cuda_forward(const paddle::Tensor &x);
+std::vector<paddle::Tensor> relu_cuda_backward(const paddle::Tensor &x,
+                                               const paddle::Tensor &out,
+                                               const paddle::Tensor &grad_out);
+
+std::vector<paddle::Tensor> ReluForward(const paddle::Tensor &x) {
+  // TODO(chenweihang): Check Input
+  if (x.place() == paddle::PlaceType::kCPU) {
+    return relu_cpu_forward(x);
+  } else if (x.place() == paddle::PlaceType::kGPU) {
+    return relu_cuda_forward(x);
+  } else {
+    throw std::runtime_error("Not implemented.");
+  }
+}
+
+std::vector<paddle::Tensor> ReluBackward(const paddle::Tensor &x,
+                                         const paddle::Tensor &out,
+                                         const paddle::Tensor &grad_out) {
+  // TODO(chenweihang): Check Input
+  if (x.place() == paddle::PlaceType::kCPU) {
+    return relu_cpu_backward(x, out, grad_out);
+  } else if (x.place() == paddle::PlaceType::kGPU) {
+    return relu_cuda_backward(x, out, grad_out);
+  } else {
+    throw std::runtime_error("Not implemented.");
+  }
+}
+
+PD_BUILD_OP(custom_relu)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .SetKernelFn(PD_KERNEL(ReluForward));
+
+PD_BUILD_GRAD_OP(custom_relu)
+    .Inputs({"X", "Out", paddle::Grad("Out")})
+    .Outputs({paddle::Grad("X")})
+    .SetKernelFn(PD_KERNEL(ReluBackward));

test_tipc/supplementary/custom_op/custom_relu_op.cu

@@ -0,0 +1,71 @@
+// 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.
+
+// reference
+// https://github.com/PaddlePaddle/Paddle-Inference-Demo/blob/master/python/custom-operator/custom_relu_op.cu
+
+#include "paddle/extension.h"
+
+template <typename data_t>
+__global__ void relu_cuda_forward_kernel(const data_t *x, data_t *y,
+                                         const int num) {
+  int gid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int i = gid; i < num; i += blockDim.x * gridDim.x) {
+    y[i] = max(x[i], static_cast<data_t>(0.));
+  }
+}
+
+template <typename data_t>
+__global__ void relu_cuda_backward_kernel(const data_t *dy, const data_t *y,
+                                          data_t *dx, const int num) {
+  int gid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int i = gid; i < num; i += blockDim.x * gridDim.x) {
+    dx[i] = dy[i] * (y[i] > 0 ? 1. : 0.);
+  }
+}
+
+std::vector<paddle::Tensor> relu_cuda_forward(const paddle::Tensor &x) {
+  auto out = paddle::Tensor(paddle::PlaceType::kGPU);
+
+  out.reshape(x.shape());
+  int numel = x.size();
+  int block = 512;
+  int grid = (numel + block - 1) / block;
+  PD_DISPATCH_FLOATING_TYPES(
+      x.type(), "relu_cuda_forward_kernel", ([&] {
+        relu_cuda_forward_kernel<data_t><<<grid, block, 0, x.stream()>>>(
+            x.data<data_t>(), out.mutable_data<data_t>(x.place()), numel);
+      }));
+
+  return {out};
+}
+
+std::vector<paddle::Tensor> relu_cuda_backward(const paddle::Tensor &x,
+                                               const paddle::Tensor &out,
+                                               const paddle::Tensor &grad_out) {
+  auto grad_x = paddle::Tensor(paddle::PlaceType::kGPU);
+  grad_x.reshape(x.shape());
+
+  int numel = out.size();
+  int block = 512;
+  int grid = (numel + block - 1) / block;
+  PD_DISPATCH_FLOATING_TYPES(
+      out.type(), "relu_cuda_backward_kernel", ([&] {
+        relu_cuda_backward_kernel<data_t><<<grid, block, 0, x.stream()>>>(
+            grad_out.data<data_t>(), out.data<data_t>(),
+            grad_x.mutable_data<data_t>(x.place()), numel);
+      }));
+
+  return {grad_x};
+}

test_tipc/supplementary/custom_op/test.py

@@ -0,0 +1,77 @@
+import paddle
+import paddle.nn as nn
+from paddle.vision.transforms import Compose, Normalize
+from paddle.utils.cpp_extension import load
+from paddle.inference import Config
+from paddle.inference import create_predictor
+import numpy as np
+
+EPOCH_NUM = 4
+BATCH_SIZE = 64
+
+# jit compile custom op
+custom_ops = load(
+    name="custom_jit_ops", sources=["custom_relu_op.cc", "custom_relu_op.cu"]
+)
+
+
+class LeNet(nn.Layer):
+    def __init__(self):
+        super(LeNet, self).__init__()
+        self.conv1 = nn.Conv2D(
+            in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2
+        )
+        self.max_pool1 = nn.MaxPool2D(kernel_size=2, stride=2)
+        self.conv2 = nn.Conv2D(in_channels=6, out_channels=16, kernel_size=5, stride=1)
+        self.max_pool2 = nn.MaxPool2D(kernel_size=2, stride=2)
+        self.linear1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
+        self.linear2 = nn.Linear(in_features=120, out_features=84)
+        self.linear3 = nn.Linear(in_features=84, out_features=10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = custom_ops.custom_relu(x)
+        x = self.max_pool1(x)
+        x = custom_ops.custom_relu(x)
+        x = self.conv2(x)
+        x = self.max_pool2(x)
+        x = paddle.flatten(x, start_axis=1, stop_axis=-1)
+        x = self.linear1(x)
+        x = custom_ops.custom_relu(x)
+        x = self.linear2(x)
+        x = custom_ops.custom_relu(x)
+        x = self.linear3(x)
+        return x
+
+
+# set device
+paddle.set_device("gpu")
+
+# model
+net = LeNet()
+loss_fn = nn.CrossEntropyLoss()
+opt = paddle.optimizer.Adam(learning_rate=0.001, parameters=net.parameters())
+
+# data loader
+transform = Compose([Normalize(mean=[127.5], std=[127.5], data_format="CHW")])
+train_dataset = paddle.vision.datasets.MNIST(mode="train", transform=transform)
+train_loader = paddle.io.DataLoader(
+    train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=2
+)
+
+# train
+for epoch_id in range(EPOCH_NUM):
+    for batch_id, (image, label) in enumerate(train_loader()):
+        out = net(image)
+        loss = loss_fn(out, label)
+        loss.backward()
+
+        if batch_id % 300 == 0:
+            print(
+                "Epoch {} batch {}: loss = {}".format(
+                    epoch_id, batch_id, np.mean(loss.numpy())
+                )
+            )
+
+        opt.step()
+        opt.clear_grad()