Как переместить переменную из GPU в CPU в API внешнего интерфейса Pytorch c ++? - PullRequest
Новая
       25

Как переместить переменную из GPU в CPU в API внешнего интерфейса Pytorch c ++?

0 голосов
/ 25 января 2019

Я пишу код логического вывода для загрузки преобразованной модели Pytorch (модель тегирования из imagenet) в C ++.Я использовал c ++ Pytorch API внешнего интерфейса.Мой код работает правильно на CPU, но не работает на GPU.Проблема в том, что когда я хочу напечатать окончательные результаты, я получаю ошибку Ошибка сегментации (ядро сброшено).Я должен передать переменные "top_scores_a" и "top_idx_a" в ЦПУ, но я не знаю, как это сделать.

Я загружаю модель и входное изображение в графический процессор.Ошибка произошла в следующей части: 

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;
}

1 Ответ

0 голосов
/ 25 января 2019

Для перемещения данных из CPU в GPU и наоборот необходимо выделить так называемую управляемую память. Посмотрите здесь пример кода https://devblogs.nvidia.com/even-easier-introduction-cuda

Если ваша версия cuda не поддерживает cudaMallocManaged, вы должны использовать последовательность cudaMalloc + cudaMemcpy.

...