C++集成ViT图像分类模型：高性能推理引擎开发指南
你是不是也遇到过这样的场景：手头有一个训练好的ViT图像分类模型，在Python里跑得挺好，但一到实际产品里，就发现Python那点性能根本不够用？延迟高、内存占用大，稍微并发多点就扛不住了。
这时候，用C++来部署模型就成了刚需。但真动手的时候，问题就来了：模型怎么转？接口怎么封？内存怎么管？多线程怎么搞？一堆坑等着你。
今天我就结合自己踩过的坑，给你讲讲怎么在C++项目里集成ViT模型，打造一个真正高性能的推理引擎。咱们不搞虚的，直接上干货，从模型转换到优化技巧，一步步带你走通。
1. 环境准备与模型选择
在开始之前，得先把基础环境搭好。C++做AI推理，现在主流的选择就那么几个框架，咱们得根据实际需求来选。
1.1 推理框架选择
目前比较成熟的C++推理框架主要有这几个：
- ONNX Runtime：微软出品，支持多种后端（CPU、CUDA、TensorRT等），社区活跃，文档齐全
- TensorRT：NVIDIA亲儿子，在GPU上优化到极致，但生态相对封闭
- OpenVINO：Intel主打，在Intel CPU上表现优秀，对边缘设备友好
- LibTorch：PyTorch的C++版本，如果你模型本来就是PyTorch的，用这个最省事
我个人的建议是：优先考虑ONNX Runtime。原因很简单，它跨平台、跨硬件，从云端到边缘都能用，而且维护得不错。咱们今天也主要用它来演示。
1.2 模型获取与转换
从搜索内容看，ViT图像分类-中文-日常物品这个模型挺实用的，覆盖1300类常见物体，日常场景够用了。咱们就用它来举例。
首先得把PyTorch模型转成ONNX格式。这是最关键的一步，转不好后面全白搭。
# convert_to_onnx.py import torch from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks # 加载模型 model_id = 'damo/cv_nextvit-small_image-classification_Dailylife-labels' image_classification = pipeline(Tasks.image_classification, model=model_id) # 获取实际的模型 model = image_classification.model # 设置为评估模式 model.eval() # 创建示例输入（注意：这个模型输入是224x224） dummy_input = torch.randn(1, 3, 224, 224) # 导出为ONNX torch.onnx.export( model, dummy_input, "vit_dailylife.onnx", input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch_size"}, "output": {0: "batch_size"} }, opset_version=13 ) print("模型转换完成！")转换的时候有几个坑要注意：
- 输入尺寸：这个模型固定输入224x224，别搞错了
- 动态轴：加上
dynamic_axes，这样推理时才能支持不同的batch size - opset版本：用13或以上，兼容性更好
1.3 开发环境搭建
C++项目我习惯用CMake来管理，这样跨平台方便。下面是个基础的CMakeLists.txt：
cmake_minimum_required(VERSION 3.14) project(ViTInference) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) # 查找ONNX Runtime find_package(onnxruntime REQUIRED) # OpenCV用于图像处理 find_package(OpenCV REQUIRED) # 添加可执行文件 add_executable(vit_inference src/main.cpp src/inference_engine.cpp) # 链接库 target_link_libraries(vit_inference PRIVATE onnxruntime::onnxruntime ${OpenCV_LIBS} ) # 包含目录 target_include_directories(vit_inference PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include ${OpenCV_INCLUDE_DIRS} )2. 核心推理引擎实现
环境搭好了，模型也转好了，现在开始写核心的推理代码。这部分是性能的关键，得仔细设计。
2.1 基础推理类设计
先设计一个基础的推理引擎类，把ONNX Runtime的接口封装起来：
// inference_engine.h #pragma once #include <onnxruntime_cxx_api.h> #include <opencv2/opencv.hpp> #include <vector> #include <string> #include <memory> class ViTInferenceEngine { public: // 构造函数：加载模型，创建会话 ViTInferenceEngine(const std::string& model_path, bool use_gpu = false); // 析构函数 ~ViTInferenceEngine(); // 单张图片推理 std::vector<float> infer(const cv::Mat& image); // 批量推理 std::vector<std::vector<float>> infer_batch( const std::vector<cv::Mat>& images); // 获取模型信息 int get_input_width() const { return input_width_; } int get_input_height() const { return input_height_; } int get_num_classes() const { return num_classes_; } private: // 预处理：把OpenCV的Mat转成模型需要的Tensor std::vector<float> preprocess(const cv::Mat& image); // 后处理：把模型输出转成可读的结果 std::vector<float> postprocess(const std::vector<float>& raw_output); // ONNX Runtime相关 Ort::Env env_; Ort::SessionOptions session_options_; std::unique_ptr<Ort::Session> session_; // 模型信息 int input_width_ = 224; int input_height_ = 224; int num_classes_ = 1300; // 根据实际模型调整 // 输入输出名称 std::vector<const char*> input_names_; std::vector<const char*> output_names_; // 内存分配器 Ort::AllocatorWithDefaultOptions allocator_; };2.2 预处理实现
预处理是性能瓶颈之一，得好好优化。ViT模型的预处理通常包括：调整大小、归一化、通道转换。
// inference_engine.cpp #include "inference_engine.h" #include <chrono> #include <iostream> ViTInferenceEngine::ViTInferenceEngine(const std::string& model_path, bool use_gpu) { // 初始化ONNX Runtime环境 env_ = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "ViTInference"); // 设置会话选项 session_options_.SetIntraOpNumThreads(1); session_options_.SetGraphOptimizationLevel( GraphOptimizationLevel::ORT_ENABLE_ALL); // 选择执行提供者 if (use_gpu) { OrtCUDAProviderOptions cuda_options; cuda_options.device_id = 0; session_options_.AppendExecutionProvider_CUDA(cuda_options); } // 创建会话 session_ = std::make_unique<Ort::Session>( env_, model_path.c_str(), session_options_); // 获取输入输出信息 Ort::AllocatorWithDefaultOptions allocator; auto input_info = session_->GetInputTypeInfo(0); auto input_tensor_info = input_info.GetTensorTypeAndShapeInfo(); auto input_shape = input_tensor_info.GetShape(); if (input_shape.size() == 4) { // NCHW格式 input_height_ = static_cast<int>(input_shape[2]); input_width_ = static_cast<int>(input_shape[3]); } // 获取输出信息 auto output_info = session_->GetOutputTypeInfo(0); auto output_tensor_info = output_info.GetTensorTypeAndShapeInfo(); auto output_shape = output_tensor_info.GetShape(); if (output_shape.size() == 2) { num_classes_ = static_cast<int>(output_shape[1]); } // 设置输入输出名称 input_names_.push_back("input"); output_names_.push_back("output"); std::cout << "模型加载成功！" << std::endl; std::cout << "输入尺寸: " << input_width_ << "x" << input_height_ << std::endl; std::cout << "类别数: " << num_classes_ << std::endl; } std::vector<float> ViTInferenceEngine::preprocess(const cv::Mat& image) { cv::Mat processed; // 1. 调整大小（保持长宽比） int target_size = std::min(input_width_, input_height_); cv::Mat resized; float scale = static_cast<float>(target_size) / std::min(image.cols, image.rows); cv::resize(image, resized, cv::Size(), scale, scale, cv::INTER_LINEAR); // 2. 中心裁剪 int start_x = (resized.cols - input_width_) / 2; int start_y = (resized.rows - input_height_) / 2; cv::Rect roi(start_x, start_y, input_width_, input_height_); cv::Mat cropped = resized(roi); // 3. 转换为float并归一化 // ViT通常使用ImageNet的均值和标准差 cropped.convertTo(processed, CV_32FC3, 1.0 / 255.0); // 减去均值，除以标准差 cv::Scalar mean(0.485, 0.456, 0.406); cv::Scalar std(0.229, 0.224, 0.225); cv::subtract(processed, mean, processed); cv::divide(processed, std, processed); // 4. 从HWC转为CHW std::vector<cv::Mat> channels(3); cv::split(processed, channels); std::vector<float> result; result.reserve(3 * input_width_ * input_height_); for (const auto& channel : channels) { result.insert(result.end(), channel.ptr<float>(), channel.ptr<float>() + channel.total()); } return result; }2.3 推理核心实现
预处理做好了，推理部分就相对简单了：
std::vector<float> ViTInferenceEngine::infer(const cv::Mat& image) { // 1. 预处理 auto start_preprocess = std::chrono::high_resolution_clock::now(); std::vector<float> input_tensor = preprocess(image); auto end_preprocess = std::chrono::high_resolution_clock::now(); // 2. 准备输入Tensor std::vector<int64_t> input_shape = {1, 3, input_height_, input_width_}; Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu( OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); Ort::Value input_tensor_ort = Ort::Value::CreateTensor<float>( memory_info, input_tensor.data(), input_tensor.size(), input_shape.data(), input_shape.size() ); // 3. 运行推理 auto start_infer = std::chrono::high_resolution_clock::now(); auto output_tensors = session_->Run( Ort::RunOptions{nullptr}, input_names_.data(), &input_tensor_ort, 1, output_names_.data(), 1 ); auto end_infer = std::chrono::high_resolution_clock::now(); // 4. 获取输出 float* output_data = output_tensors[0].GetTensorMutableData<float>(); std::vector<float> output(output_data, output_data + num_classes_); // 5. 后处理（softmax） std::vector<float> probabilities = postprocess(output); // 输出耗时信息 auto preprocess_time = std::chrono::duration_cast<std::chrono::milliseconds>( end_preprocess - start_preprocess); auto infer_time = std::chrono::duration_cast<std::chrono::milliseconds>( end_infer - start_infer); std::cout << "预处理耗时: " << preprocess_time.count() << "ms" << std::endl; std::cout << "推理耗时: " << infer_time.count() << "ms" << std::endl; return probabilities; } std::vector<float> ViTInferenceEngine::postprocess( const std::vector<float>& raw_output) { // 简单的softmax实现 std::vector<float> probabilities(num_classes_); // 防止数值溢出，减去最大值 float max_val = *std::max_element(raw_output.begin(), raw_output.end()); float sum = 0.0f; for (int i = 0; i < num_classes_; ++i) { probabilities[i] = std::exp(raw_output[i] - max_val); sum += probabilities[i]; } // 归一化 for (int i = 0; i < num_classes_; ++i) { probabilities[i] /= sum; } return probabilities; }3. 性能优化技巧
基础功能实现了，但离"高性能"还有距离。下面这些优化技巧，能让你的推理引擎快上加快。
3.1 内存池优化
ONNX Runtime默认会为每次推理分配新内存，频繁分配释放很耗时间。我们可以用内存池来优化：
class MemoryPool { public: MemoryPool(size_t max_batch_size, int input_size) { // 预分配输入输出内存 input_pool_.reserve(max_batch_size); output_pool_.reserve(max_batch_size); for (size_t i = 0; i < max_batch_size; ++i) { input_pool_.push_back(std::vector<float>(input_size)); output_pool_.push_back(std::vector<float>(num_classes_)); } } std::vector<float>& get_input_buffer() { std::lock_guard<std::mutex> lock(mutex_); if (free_inputs_.empty()) { // 动态扩展（尽量避免） input_pool_.emplace_back(input_size_); return input_pool_.back(); } auto& buffer = input_pool_[free_inputs_.back()]; free_inputs_.pop_back(); return buffer; } void return_input_buffer(size_t index) { std::lock_guard<std::mutex> lock(mutex_); free_inputs_.push_back(index); } private: std::vector<std::vector<float>> input_pool_; std::vector<std::vector<float>> output_pool_; std::vector<size_t> free_inputs_; std::mutex mutex_; size_t input_size_; int num_classes_; };3.2 批量推理优化
单张推理效率低，批量推理能充分利用硬件并行能力：
std::vector<std::vector<float>> ViTInferenceEngine::infer_batch( const std::vector<cv::Mat>& images) { size_t batch_size = images.size(); std::vector<std::vector<float>> batch_inputs(batch_size); // 并行预处理 #pragma omp parallel for for (size_t i = 0; i < batch_size; ++i) { batch_inputs[i] = preprocess(images[i]); } // 合并batch std::vector<float> merged_input; merged_input.reserve(batch_size * 3 * input_height_ * input_width_); for (const auto& input : batch_inputs) { merged_input.insert(merged_input.end(), input.begin(), input.end()); } // 准备batch输入 std::vector<int64_t> input_shape = { static_cast<int64_t>(batch_size), 3, input_height_, input_width_ }; Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu( OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); Ort::Value input_tensor_ort = Ort::Value::CreateTensor<float>( memory_info, merged_input.data(), merged_input.size(), input_shape.data(), input_shape.size() ); // 批量推理 auto output_tensors = session_->Run( Ort::RunOptions{nullptr}, input_names_.data(), &input_tensor_ort, 1, output_names_.data(), 1 ); // 处理批量输出 float* batch_output = output_tensors[0].GetTensorMutableData<float>(); std::vector<std::vector<float>> results(batch_size); for (size_t i = 0; i < batch_size; ++i) { results[i].resize(num_classes_); std::copy(batch_output + i * num_classes_, batch_output + (i + 1) * num_classes_, results[i].begin()); results[i] = postprocess(results[i]); } return results; }3.3 异步推理实现
对于实时应用，异步推理能大幅提升吞吐量：
class AsyncInferenceEngine { public: AsyncInferenceEngine(const std::string& model_path, size_t worker_count = 2) : engine_(model_path), stop_(false) { // 创建工作线程 for (size_t i = 0; i < worker_count; ++i) { workers_.emplace_back([this]() { worker_loop(); }); } } ~AsyncInferenceEngine() { stop_ = true; cv_.notify_all(); for (auto& worker : workers_) { if (worker.joinable()) { worker.join(); } } } // 提交推理任务 std::future<std::vector<float>> submit(const cv::Mat& image) { auto promise = std::make_shared<std::promise<std::vector<float>>>(); auto future = promise->get_future(); { std::lock_guard<std::mutex> lock(queue_mutex_); task_queue_.push({image, promise}); } cv_.notify_one(); return future; } private: struct InferenceTask { cv::Mat image; std::shared_ptr<std::promise<std::vector<float>>> promise; }; void worker_loop() { while (!stop_) { InferenceTask task; { std::unique_lock<std::mutex> lock(queue_mutex_); cv_.wait(lock, [this]() { return !task_queue_.empty() || stop_; }); if (stop_) break; task = std::move(task_queue_.front()); task_queue_.pop(); } try { auto result = engine_.infer(task.image); task.promise->set_value(result); } catch (...) { task.promise->set_exception(std::current_exception()); } } } ViTInferenceEngine engine_; std::queue<InferenceTask> task_queue_; std::mutex queue_mutex_; std::condition_variable cv_; std::vector<std::thread> workers_; std::atomic<bool> stop_; };4. 实际应用示例
理论讲完了，来看看怎么在实际项目里用。咱们写个完整的示例程序：
// main.cpp #include "inference_engine.h" #include <iostream> #include <fstream> #include <algorithm> // 加载类别标签（根据实际模型调整） std::vector<std::string> load_labels(const std::string& label_path) { std::vector<std::string> labels; std::ifstream file(label_path); if (!file.is_open()) { std::cerr << "无法打开标签文件: " << label_path << std::endl; return labels; } std::string line; while (std::getline(file, line)) { labels.push_back(line); } return labels; } int main(int argc, char* argv[]) { if (argc < 3) { std::cout << "用法: " << argv[0] << " <模型路径> <图片路径>" << std::endl; return 1; } std::string model_path = argv[1]; std::string image_path = argv[2]; try { // 1. 创建推理引擎 std::cout << "正在加载模型..." << std::endl; ViTInferenceEngine engine(model_path, false); // 使用CPU // 2. 加载图片 cv::Mat image = cv::imread(image_path); if (image.empty()) { std::cerr << "无法加载图片: " << image_path << std::endl; return 1; } std::cout << "图片尺寸: " << image.cols << "x" << image.rows << std::endl; // 3. 执行推理 auto start = std::chrono::high_resolution_clock::now(); auto probabilities = engine.infer(image); auto end = std::chrono::high_resolution_clock::now(); auto total_time = std::chrono::duration_cast<std::chrono::milliseconds>( end - start); std::cout << "总耗时: " << total_time.count() << "ms" << std::endl; // 4. 解析结果 // 找到top-5预测结果 std::vector<int> indices(probabilities.size()); std::iota(indices.begin(), indices.end(), 0); std::partial_sort(indices.begin(), indices.begin() + 5, indices.end(), [&probabilities](int a, int b) { return probabilities[a] > probabilities[b]; }); std::cout << "\nTop-5预测结果:" << std::endl; for (int i = 0; i < 5; ++i) { int idx = indices[i]; std::cout << i + 1 << ". 类别" << idx << ": " << probabilities[idx] * 100 << "%" << std::endl; } // 5. 批量推理示例 std::cout << "\n=== 批量推理示例 ===" << std::endl; std::vector<cv::Mat> batch_images = {image, image, image}; // 重复3次 auto batch_results = engine.infer_batch(batch_images); std::cout << "批量推理完成，处理了 " << batch_results.size() << " 张图片" << std::endl; } catch (const std::exception& e) { std::cerr << "错误: " << e.what() << std::endl; return 1; } return 0; }编译运行：
mkdir build && cd build cmake .. make ./vit_inference vit_dailylife.onnx test_image.jpg5. 总结
走完这一整套流程，你应该对C++集成ViT模型有了比较清晰的认识。从模型转换到接口封装，从基础推理到性能优化，每个环节都有需要注意的地方。
实际用下来，ONNX Runtime的表现确实不错，跨平台部署很方便。预处理部分是最容易出性能瓶颈的地方，一定要做好优化。批量推理和异步处理对于高并发场景是必须的，能大幅提升吞吐量。
内存管理是C++项目的永恒话题，预分配内存池、合理使用智能指针，能避免很多内存泄漏的问题。多线程环境下更要注意线程安全，该加锁的地方不能省。
如果你刚开始接触这块，建议先从单张图片推理做起，把流程跑通，然后再逐步加入批量处理、异步推理等高级功能。遇到问题多查文档，ONNX Runtime的社区挺活跃的，大部分问题都能找到解决方案。
最后提醒一点，不同的应用场景对性能的要求不同。如果是实时视频分析，延迟是关键；如果是后台批量处理，吞吐量更重要。根据实际需求来调整优化策略，别过度优化。
获取更多AI镜像
想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。
