add submodule pytorch

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/pytorch"]
+	path = third_party/pytorch
+	url = https://github.com/Jonahmoon/pytorch.git

__init__.py

@@ -0,0 +1,6 @@
+from setuptools import setup, Extension
+from torch.utils import cpp_extension
+
+setup(name='add_cpp',
+      ext_modules=[cpp_extension.CppExtension('add_cpp', ['add.cpp'])],
+      cmdclass={'build_ext': cpp_extension.BuildExtension})

csrc/aten/PrivateUse1/myadd.cpp

@@ -0,0 +1,51 @@
+#include <iostream>
+#include <vector>
+#include <opencv2/opencv.hpp>
+#include <torch/extension.h>
+#include <torch/script.h>
+
+/*
+Generally, speaking, the structure of your registrations will look like this:
+
+1. A single TORCH_LIBRARY that lists every custom operator in your namespace in a centralized place.
+
+2. A TORCH_LIBRARY_IMPL per dispatch key that registers implementations for that key (e.g., CPU or CUDA). If you like, you can further subdivide TORCH_LIBRARY_IMPL blocks into a block per operator. This is convenient if you have a separate file per operator implementation, but don’t want to expose the operators in a header; you can just put the registration in the cpp file that defines your operator.
+*/
+
+
+
+// we need to actually provide some implementations of this operator. For concreteness, here is a really simple implementation of addition on CPU
+
+Tensor myadd_cpu(const Tensor& self_, const Tensor& other_) {
+  TORCH_CHECK(self_.sizes() == other_.sizes());
+  TORCH_INTERNAL_ASSERT(self_.device().type() == DeviceType::CPU);
+  TORCH_INTERNAL_ASSERT(other_.device().type() == DeviceType::CPU);
+  Tensor self = self_.contiguous();
+  Tensor other = other_.contiguous();
+  Tensor result = torch::empty(self.sizes(), self.options());
+  const float* self_ptr = self.data_ptr<float>();
+  const float* other_ptr = other.data_ptr<float>();
+  float* result_ptr = result.data_ptr<float>();
+  for (int64_t i = 0; i < result.numel(); i++) {
+    result_ptr[i] = self_ptr[i] + other_ptr[i];
+  }
+  return result;
+}
+
+// we just provide a schema string specifying the type signature of the operator that all of our other kernels will abide by
+
+TORCH_LIBRARY(myops, m) {
+  m.def("myadd(Tensor self, Tensor other) -> Tensor");
+}
+
+// the simple way of registering it (def("myadd", myadd_cpu)) would register the kernel to run in all cases, even if the tensor is not a CPU tensor! (Internally, we refer to these as “catch-all” kernels, since they catch all cases.) To ensure that myadd_cpu is only run for CPU tensors, we can use the TORCH_LIBRARY_IMPL macro
+
+TORCH_LIBRARY_IMPL(myops, CPU, m) {
+  m.impl("myadd", myadd_cpu);
+}
+
+/*
+TORCH_LIBRARY_IMPL(myops, CUDA, m) {
+  m.impl("myadd", myadd_cuda);
+}
+*/

csrc/jit/CMakeLists.txt

@@ -0,0 +1,16 @@
+cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
+project(warp_perspective)
+
+include_directories("/usr/include/python3.9")
+
+find_package(Torch REQUIRED)
+find_package(OpenCV REQUIRED)
+
+# Define our library target
+add_library(warp_perspective SHARED op.cpp)
+# Enable C++14
+target_compile_features(warp_perspective PRIVATE cxx_std_14)
+# Link against LibTorch
+target_link_libraries(warp_perspective "${TORCH_LIBRARIES}")
+# Link against OpenCV
+target_link_libraries(warp_perspective opencv_core opencv_imgproc)

csrc/jit/README.md

@@ -0,0 +1,6 @@
+```bash
+mkdir build
+cd build
+cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
+make -j
+```

csrc/jit/op.cpp

@@ -0,0 +1,39 @@
+#include <iostream>
+#include <vector>
+#include <opencv2/opencv.hpp>
+#include <torch/extension.h>
+#include <torch/script.h>
+
+torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) {
+  // BEGIN image_mat
+  cv::Mat image_mat(/*rows=*/image.size(0),
+                    /*cols=*/image.size(1),
+                    /*type=*/CV_32FC1,
+                    /*data=*/image.data_ptr<float>());
+  // END image_mat
+
+  // BEGIN warp_mat
+  cv::Mat warp_mat(/*rows=*/warp.size(0),
+                   /*cols=*/warp.size(1),
+                   /*type=*/CV_32FC1,
+                   /*data=*/warp.data_ptr<float>());
+  // END warp_mat
+
+  // BEGIN output_mat
+  cv::Mat output_mat;
+  cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{8, 8});
+  // END output_mat
+
+  // BEGIN output_tensor
+  torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8});
+  return output.clone();
+  // END output_tensor
+}
+
+/*
+somewhere at the top level of our op.cpp file. The TORCH_LIBRARY macro creates a function that will be called when your program starts. The name of your library (my_ops) is given as the first argument (it should not be in quotes). The second argument (m) defines a variable of type torch::Library which is the main interface to register your operators. The method Library::def actually creates an operator named warp_perspective, exposing it to both Python and TorchScript. You can define as many operators as you like by making multiple calls to def.
+*/
+
+TORCH_LIBRARY(my_ops, m) {
+  m.def("warp_perspective", warp_perspective);
+}

csrc/lltm.cpp

@@ -0,0 +1,92 @@
+#include <iostream>
+#include <vector>
+#include <torch/extension.h>
+
+torch::Tensor my_add(torch::Tensor a, torch::Tensor b) {
+  auto c = a + b;
+  return c;
+}
+
+std::vector<torch::Tensor> lltm_forward(
+    torch::Tensor input,
+    torch::Tensor weights,
+    torch::Tensor bias,
+    torch::Tensor old_h,
+    torch::Tensor old_cell) {
+  auto X = torch::cat({old_h, input}, /*dim=*/1);
+
+  auto gate_weights = torch::addmm(bias, X, weights.transpose(0, 1));
+  auto gates = gate_weights.chunk(3, /*dim=*/1);
+
+  auto input_gate = torch::sigmoid(gates[0]);
+  auto output_gate = torch::sigmoid(gates[1]);
+  auto candidate_cell = torch::elu(gates[2], /*alpha=*/1.0);
+
+  auto new_cell = old_cell + candidate_cell * input_gate;
+  auto new_h = torch::tanh(new_cell) * output_gate;
+
+  return {new_h,
+          new_cell,
+          input_gate,
+          output_gate,
+          candidate_cell,
+          X,
+          gate_weights};
+}
+
+torch::Tensor d_tanh(torch::Tensor z) {
+  return 1 - z.tanh().pow(2);
+}
+
+torch::Tensor d_elu(torch::Tensor z, torch::Scalar alpha = 1.0) {
+  auto e = z.exp();
+  auto mask = (alpha * (e - 1)) < 0;
+  return (z > 0).type_as(z) + mask.type_as(z) * (alpha * e);
+}
+
+torch::Tensor d_sigmoid(torch::Tensor z) {
+  auto s = torch::sigmoid(z);
+  return (1 - s) * s;
+}
+
+std::vector<torch::Tensor> lltm_backward(
+    torch::Tensor grad_h,
+    torch::Tensor grad_cell,
+    torch::Tensor new_cell,
+    torch::Tensor input_gate,
+    torch::Tensor output_gate,
+    torch::Tensor candidate_cell,
+    torch::Tensor X,
+    torch::Tensor gate_weights,
+    torch::Tensor weights) {
+  auto d_output_gate = torch::tanh(new_cell) * grad_h;
+  auto d_tanh_new_cell = output_gate * grad_h;
+  auto d_new_cell = d_tanh(new_cell) * d_tanh_new_cell + grad_cell;
+
+  auto d_old_cell = d_new_cell;
+  auto d_candidate_cell = input_gate * d_new_cell;
+  auto d_input_gate = candidate_cell * d_new_cell;
+
+  auto gates = gate_weights.chunk(3, /*dim=*/1);
+  d_input_gate *= d_sigmoid(gates[0]);
+  d_output_gate *= d_sigmoid(gates[1]);
+  d_candidate_cell *= d_elu(gates[2]);
+
+  auto d_gates =
+      torch::cat({d_input_gate, d_output_gate, d_candidate_cell}, /*dim=*/1);
+
+  auto d_weights = d_gates.t().mm(X);
+  auto d_bias = d_gates.sum(/*dim=*/0, /*keepdim=*/true);
+
+  auto d_X = d_gates.mm(weights);
+  const auto state_size = grad_h.size(1);
+  auto d_old_h = d_X.slice(/*dim=*/1, 0, state_size);
+  auto d_input = d_X.slice(/*dim=*/1, state_size);
+
+  return {d_old_h, d_input, d_weights, d_bias, d_old_cell};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &lltm_forward, "LLTM forward");
+  m.def("backward", &lltm_backward, "LLTM backward");
+}

setup.py

@@ -0,0 +1,14 @@
+from setuptools import setup
+from torch.utils import cpp_extension
+
+setup(
+    name='lltm_cpp',						# 编译后的链接库名称
+    ext_modules=[
+        cpp_extension.CppExtension(
+            'lltm_cpp', ['csrc/lltm.cpp']		       # 待编译文件，及编译函数
+        )
+    ],
+    cmdclass={						       # 执行编译命令设置
+        'build_ext': cpp_extension.BuildExtension.with_options(no_python_abi_suffix=True)
+    }
+)

test/test_lltm.py

@@ -0,0 +1,66 @@
+import math
+import time
+import torch
+import lltm_cpp as lltm
+
+# help(lltm.forward)
+
+
+class LLTMFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weights, bias, old_h, old_cell):
+        outputs = lltm.forward(input, weights, bias, old_h, old_cell)
+        new_h, new_cell = outputs[:2]
+        variables = outputs[1:] + [weights]
+        ctx.save_for_backward(*variables)
+
+        return new_h, new_cell
+
+    @staticmethod
+    def backward(ctx, grad_h, grad_cell):
+        outputs = lltm.backward(
+            grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_tensors)
+        d_old_h, d_input, d_weights, d_bias, d_old_cell = outputs
+        return d_input, d_weights, d_bias, d_old_h, d_old_cell
+
+
+class LLTM(torch.nn.Module):
+    def __init__(self, input_features, state_size):
+        super(LLTM, self).__init__()
+        self.input_features = input_features
+        self.state_size = state_size
+        self.weights = torch.nn.Parameter(
+            torch.empty(3 * state_size, input_features + state_size))
+        self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.state_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, +stdv)
+
+    def forward(self, input, state):
+        return LLTMFunction.apply(input, self.weights, self.bias, *state)
+
+batch_size = 16
+input_features = 32
+state_size = 128
+
+X = torch.randn(batch_size, input_features)
+h = torch.randn(batch_size, state_size)
+C = torch.randn(batch_size, state_size)
+
+rnn = LLTM(input_features, state_size)
+
+forward = 0
+backward = 0
+for _ in range(1000):
+    start = time.time()
+    new_h, new_C = rnn(X, (h, C))
+    forward += time.time() - start
+
+    start = time.time()
+    (new_h.sum() + new_C.sum()).backward()
+    backward += time.time() - start
+
+print('Forward: {:.3f} s | Backward {:.3f} s'.format(forward, backward))

test/test_warp.py

@@ -0,0 +1,13 @@
+import torch
+
+torch.ops.load_library("../csrc/jit/build/libwarp_perspective.so")
+print(torch.ops.my_ops.warp_perspective)
+print(torch.ops.my_ops.warp_perspective(torch.randn(32, 32), torch.rand(3, 3)))
+
+def compute(x, y, z):
+    x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
+    return x.matmul(y) + torch.relu(z)
+
+inputs = [torch.randn(4, 8), torch.randn(8, 5), torch.randn(8, 5)]
+trace = torch.jit.trace(compute, inputs)
+print(trace.graph)