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

Question

我正在编写一个推理代码，以C++加载转换后的pytorch模型（来自imagenet的标记模型）。我使用了 c++ pytorch 前端 API。我的代码在 CPU 上正常工作，但在 GPU 上不起作用。问题是，当我想打印最终结果时，出现分段错误（核心转储）错误。我必须将"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.txtn";
        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: %.2fsn", (double)(tend - tStart)/CLOCKS_PER_SEC);
    return 0;
}

Answer 1

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

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