深度学习模型打包

#python模型打包 生成mnist.pt文件
import torch
from torch import jit
from MLP.net import Net

# 模型的部署
if __name__ == '__main__':
    model = Net()
    model.load_state_dict(torch.load("./checkpoint/17.pt"))
    input_Data = torch.randn(1,784)
    trace_model = jit.trace(model,input_Data)
    trace_model.save("mnist.pt")

#include 
#include 
#include 
#include 
#include 
#include 

int main(){
    try {
        # 读取图片
        cv::Mat image = cv::imread("0.jpg", cv::IMREAD_GRAYSCALE);
        # 转数据类型
        torch::Tensor tensor_image = torch::from_blob(image.data, {1, image.rows * image.cols}, torch::kByte).toType(torch::kFloat);
        # 归一化
        tensor_image /= 255.;
        auto module = torch::jit::load("mnist.pt");
        std::vector inputs;
        inputs.push_back(tensor_image);
        auto rst = module.forward(inputs).toTensor();
//        std::cout << rst << std::endl;
        std::cout << torch::argmax(rst, 1) << std::endl;
    }catch(const c10::Error &e) {
        std::cerr << e.what();
        return -1;
    }
    return 0;
}

对应的cmakeList.txt
cmake_minimum_required(VERSION 3.15)
project(DeepLearning)

set(CMAKE_PREFIX_PATH D:/lib/libtorch)
set(OpenCV_DIR D:/lib/opencv/build/x64/vc15/lib)

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")


add_executable(DeepLearning  main.cpp)
target_link_libraries(DeepLearning ${TORCH_LIBRARIES} ${OpenCV_LIBS})
set_property(TARGET DeepLearning PROPERTY CXX_STANDARD 11)

if (MSVC)
    file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
    add_custom_command(TARGET DeepLearning
            POST_BUILD
            COMMAND ${CMAKE_COMMAND} -E copy_if_different
            ${TORCH_DLLS}
            $)

    file(GLOB OPENCV_DLLS "D:/lib/opencv/build/x64/vc15/bin/*.dll")
    add_custom_command(TARGET DeepLearning
            POST_BUILD
            COMMAND ${CMAKE_COMMAND} -E copy_if_different
            ${OPENCV_DLLS}
            $)
endif (MSVC)