本地wordpress平台长沙优化网站厂家
环境配置
建议在linux上配置对应环境
可以在autoDL上租一个服务器来搭建,带有pytorch的环境
https://www.autodl.com/home
我自己的win电脑上安装了pytorch,但是配置时会报错,于是到ubuntu上配置
电脑上装有pytorch的就不需要再下载libtorch了,pytorch就带有libtorch
首先通过下面的代码找出pytorch自带的环境
import torch
torch.utils.cmake_prefix_path
整体架构图
编写CMakeLists.txt文件
cmake_minimum_required(VERSION 3.6 FATAL_ERROR)
project(test-libtorch)
# 设置libtorch的位置
set(CMAKE_PREFIX_PATH "/root/miniconda3/lib/python3.8/site-packages/torch/share/cmake")
message(${CMAKE_PREFIX_PATH})
find_package(Torch REQUIRED)
#op.cpp exe
add_executable(test-libtorch op.cpp)# #link libtorch .a .so
target_link_libraries(test-libtorch "${TORCH_LIBRARIES}")
#
set_property(TARGET test-libtorch PROPERTY CXX_STANDARD 14)
编写op.cpp
#include<torch/torch.h>
#include <iostream>int main() {torch::Tensor tensor = torch::rand({ 2, 3 });std::cout << tensor << std::endl;return 0;
}
编译
先创建build文件夹,再在build里面编译
mkdir build
cd build
cmake ../
make
./test-libtorch
配置opencv
编写第一个自定义算子
算子的定义与实现
由于我们重写算子一般都是为了性能,所以多半使用C++
op.cpp
#include <torch/torch.h>
// 实现
torch::Tensor my_add(torch::Tensor t1, torch::Tensor t2)
{assert(t1.size(0) == t2.size(0));assert(t1.size(1) == t2.size(1));torch::Tensor t=t1+t2;return t;
}
// C++函数绑定到python中
TORCH_LIBRARY(my_ops, m)
{m.def("my_add", my_add);
}
CMakeLists.txt
cmake_minimum_required(VERSION 3.6 FATAL_ERROR)
project(test-libtorch)
# 设置libtorch的位置
set(CMAKE_PREFIX_PATH "/root/miniconda3/lib/python3.8/site-packages/torch/share/cmake")
message(${CMAKE_PREFIX_PATH})
find_package(Torch REQUIRED)
# find_package(Opencv REQUIRED)
add_library(test-libtorch SHARED op.cpp)
# add_executable(test-libtorch op.cpp)# #link libtorch .a .so
target_link_libraries(test-libtorch "${TORCH_LIBRARIES}")
#
set_property(TARGET test-libtorch PROPERTY CXX_STANDARD 14)
test.py
import torchlib_path = r"/root/test/build/libtest-libtorch.so"
torch.ops.load_library(lib_path)def test_add():a = torch.rand([10, 10, 3])b = torch.rand([10, 10, 3])c = torch.ops.my_ops.my_add(a, b)d = a + bassert torch.allclose(c, d)
test_add()
运行步骤
mkdir build
cd build
cmake ../
make
cd ../
python test.py
将op.cpp替换
op.cpp
#include <torch/torch.h>
// 实现
torch::Tensor my_add(torch::Tensor t1, torch::Tensor t2)
{assert(t1.size(0) == t2.size(0));assert(t1.size(1) == t2.size(1));torch::Tensor t=t1+t2;return t;
}
// 直接用Pybind绑定了C++函数
PYBIND11_MODULE(my_ops, m)
{m.def("my_add", my_add);
}
编写setup.py
from setuptools import setup
from torch.utils import cpp_extensionsetup(name='my_add',ext_modules=[cpp_extension.CppExtension('my_ops', ['op.cpp'])],cmdclass={'build_ext': cpp_extension.BuildExtension})
test.py
import torch
import my_opsdef test_add():a = torch.rand([10, 10, 3])b = torch.rand([10, 10, 3])c = my_ops.my_add(a, b)d = a + bassert torch.allclose(c, d)
test_add()
运行步骤
python setup.py install
python test.py