Не удается получить аналогичный результат прогнозирования DL в API Pytorch C ++ по сравнению с Python - PullRequest
Новая
       24

Не удается получить аналогичный результат прогнозирования DL в API Pytorch C ++ по сравнению с Python

1 голос
/ 05 апреля 2019

Я обучил модели глубокого обучения с использованием архитектуры unet, чтобы сегментировать ядра в python и pytorch.Я хотел бы загрузить эту предварительно обученную модель и сделать прогноз в C ++.По этой причине я получил файл трассировки (с расширением pt).Затем я запустил этот код: 

#include <iostream>

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>


using namespace cv;


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


    Mat image;
    image = imread("C:/Users/Sercan/PycharmProjects/samplepyTorch/test_2.png", CV_LOAD_IMAGE_COLOR);

    std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("C:/Users/Sercan/PycharmProjects/samplepyTorch/epistroma_unet_best_model_trace.pt");
    module->to(torch::kCUDA);


    std::vector<int64_t> sizes = { 1, 3, image.rows, image.cols };
    torch::TensorOptions options(torch::ScalarType::Byte);
    torch::Tensor tensor_image = torch::from_blob(image.data, torch::IntList(sizes), options);
    tensor_image = tensor_image.toType(torch::kFloat);
    auto result = module->forward({ tensor_image.to(at::kCUDA) }).toTensor();


    result = result.squeeze().cpu();
    result = at::sigmoid(result);

    cv::Mat img_out(image.rows, image.cols, CV_32F, result.data<float>());

    cv::imwrite("img_out.png", img_out);

}

Вывод изображения (Первое изображение: тестовое изображение, Второе изображение: результат прогнозирования Python, Третье изображение: результат прогнозирования C ++):

enter image description here

Как вы видите, вывод предсказания C ++ не похож на вывод предсказания Python.Не могли бы вы предложить решение, чтобы решить эту проблему?

...