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

Stack Overflow用户

提问于 2019-01-25 14:50:20

回答 1查看 2.1K关注 0票数 1

我正在编写一段推理代码，以便在C++中加载转换后的pytorch模型(来自imagenet的标记模型)。我使用的是c++ pytorch前端接口。我的代码在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.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;
}