diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..9e5fa0c --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/pytorch"] + path = third_party/pytorch + url = https://github.com/Jonahmoon/pytorch.git diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8d79189 --- /dev/null +++ b/__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}) diff --git a/csrc/aten/PrivateUse1/myadd.cpp b/csrc/aten/PrivateUse1/myadd.cpp new file mode 100644 index 0000000..d3e357c --- /dev/null +++ b/csrc/aten/PrivateUse1/myadd.cpp @@ -0,0 +1,51 @@ +#include +#include +#include +#include +#include + +/* +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(); + const float* other_ptr = other.data_ptr(); + float* result_ptr = result.data_ptr(); + 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); +} +*/ diff --git a/csrc/jit/CMakeLists.txt b/csrc/jit/CMakeLists.txt new file mode 100644 index 0000000..5df1148 --- /dev/null +++ b/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) diff --git a/csrc/jit/README.md b/csrc/jit/README.md new file mode 100644 index 0000000..b257cc3 --- /dev/null +++ b/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 +``` diff --git a/csrc/jit/op.cpp b/csrc/jit/op.cpp new file mode 100644 index 0000000..c7711fa --- /dev/null +++ b/csrc/jit/op.cpp @@ -0,0 +1,39 @@ +#include +#include +#include +#include +#include + +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()); + // 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()); + // 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(), /*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); +} diff --git a/csrc/lltm.cpp b/csrc/lltm.cpp new file mode 100644 index 0000000..e95539f --- /dev/null +++ b/csrc/lltm.cpp @@ -0,0 +1,92 @@ +#include +#include +#include + +torch::Tensor my_add(torch::Tensor a, torch::Tensor b) { + auto c = a + b; + return c; +} + +std::vector 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 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"); +} diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..ee86d31 --- /dev/null +++ b/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) + } +) diff --git a/test/test_lltm.py b/test/test_lltm.py new file mode 100644 index 0000000..3b275f4 --- /dev/null +++ b/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)) diff --git a/test/test_warp.py b/test/test_warp.py new file mode 100644 index 0000000..f0c21b7 --- /dev/null +++ b/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) diff --git a/third_party/pytorch b/third_party/pytorch new file mode 160000 index 0000000..4cdfced --- /dev/null +++ b/third_party/pytorch @@ -0,0 +1 @@ +Subproject commit 4cdfceddd2fd5ddaa51b556f85d956822f835d9a