如何在 c++ pytorch 前端 API 中将变量从 GPU 移动到 CPU？

Question

我正在编写推理代码以在 C++ 中加载转换后的 pytorch 模型（来自 imagenet 的标记模型）。我使用了 c++ pytorch 前端 API。我的代码在 CPU 上正常工作，但在 GPU 上不工作。 问题是当我想打印最终结果时，我得到了 Segmentation fault (core dumped) 错误。 我必须将“top_scores_a”和“top_idx_a”变量传输到 CPU，但我不知道该怎么做。

我在 GPU 上加载模型和输入图像。 错误发生在以下部分：

for (int i = 0; i < 2; ++i)
    {
        // int idx = top_idxs_a[i];
        std::cout << "top-" << i+1 << " label: ";
        // std::cout << labels[idx] << ", score: " << top_scores_a[i] << std::endl;
    }

完整的代码在这里：

#include "torch/script.h"
#include <torch/script.h>
#include <torch/torch.h>
#include <ATen/Tensor.h>
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <time.h> 

#include <iostream>
#include <memory>
#include <cuda.h>
#include <cuda_runtime_api.h>

using namespace std;



// __global__
int main(int argc, const char* argv[]) {

    //// asign gpu
    torch::Device device(torch::kCPU);
    clock_t tStart = clock();

    //// check cuda visibility
    if (torch::cuda::is_available()) 
    {
        std::cout << "CUDA is available! Run on GPU." << std::endl;
        device = torch::kCUDA;

    }

    if (argc != 4) {
        cout << "ptcpp path/to/scripts/model.pt path/to/image.jpg path/to/label.txt\n";
        return -1;
    }

    cout << "Will load from " << argv[1] << endl;
    shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);
    module->to(device); // on gpu

    if (module == nullptr) {
        cerr << "model load error from " << argv[1] << endl;
    }
    cout << "Model load ok.\n";

    // load image and transform
    cv::Mat image;
    image = cv::imread(argv[2], 1);

    cv::Mat image_rgb;
    cv::cvtColor(image, image_rgb, CV_BGR2RGB);  

    cv::Mat image_resized;
    cv::resize(image_rgb, image_resized, cv::Size(224, 224));

    cv::Mat image_resized_float;
    image_resized.convertTo(image_resized_float, CV_32F, 1.0/255);

    auto img_tensor = torch::CPU(torch::kFloat32).tensorFromBlob(image_resized_float.data, {1, 224, 224, 3}).to(device); // work correctly

    cout << "img tensor loaded..\n";
    img_tensor = img_tensor.permute({0, 3, 1, 2});
    img_tensor[0][0] = img_tensor[0][0].sub(0.485).div(0.229);
    img_tensor[0][1] = img_tensor[0][1].sub(0.456).div(0.224);
    img_tensor[0][2] = img_tensor[0][2].sub(0.406).div(0.225);

    auto img_var = torch::autograd::make_variable(img_tensor, false);

    vector<torch::jit::IValue> inputs;
    inputs.push_back(img_var);
    torch::Tensor out_tensor = module->forward(inputs).toTensor();


    // load labels
    vector<string> labels;
    ifstream ins;
    ins.open(argv[3]);
    string line;
    while (getline(ins, line)) 
    {
        labels.push_back(line);
    }


    std::tuple<torch::Tensor,torch::Tensor> result = out_tensor.sort(-1, true); //-1
    torch::Tensor top_scores = std::get<0>(result)[0];
    torch::Tensor top_idxs = std::get<1>(result)[0].toType(torch::kInt32);

    auto top_scores_a = top_scores.accessor<float,1>();
    auto top_idxs_a = top_idxs.accessor<int,1>();


    for (int i = 0; i < 2; ++i)
    {
        int idx = top_idxs_a[i];
        std::cout << "top-" << i+1 << " label: ";
        std::cout << labels[idx] << ", score: " << top_scores_a[i] << std::endl;
    }


    float tend = clock();
    printf("Time taken: %.2fs\n", (double)(tend - tStart)/CLOCKS_PER_SEC);

    return 0;
}

Answer 1

要将数据从 CPU 移动到 GPU，反之亦然，您必须分配所谓的托管内存。在这里查看一些示例代码https://devblogs.nvidia.com/even-easier-introduction-cuda

如果您的 cuda 版本不支持 cudaMallocManaged，那么您必须使用 cudaMalloc + cudaMemcpy 序列。