libtorch环境配置

环境配置

建议在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 torch

lib_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 + b
    assert 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_extension

setup(name='my_add',
      ext_modules=[
          cpp_extension.CppExtension('my_ops', ['op.cpp'])
      ],
      cmdclass={'build_ext': cpp_extension.BuildExtension})

test.py

import torch
import my_ops

def test_add():
    a = torch.rand([10, 10, 3])
    b = torch.rand([10, 10, 3])
    c = my_ops.my_add(a, b)
    d = a + b
    assert torch.allclose(c, d)
test_add()

运行步骤

python setup.py install
python test.py

原文地址：https://blog.csdn.net/qq_41921315/article/details/140881330

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！