Perform inference MXNet model converted to TVM

tvmInference.cpp

#include <fstream>
#include <iostream>
#include <memory>

#include "opencv2/opencv.hpp"
#include "tvm/runtime/module.h"
#include "tvm/runtime/registry.h"
#include "tvm/runtime/packed_func.h"


constexpr int FEATURE_SIZE = 512;
constexpr int BATCH_SIZE = 1;
constexpr int CHANNELS = 3;
constexpr int IMG_HEIGHT = 112;
constexpr int IMG_WIDTH = 112;


class TF_TVM {
public:
    TF_TVM(const std::string& paramsPath, const std::string& graphPath, const std::string& libPath);
    cv::Mat getTemplate(const std::string& imgPath);
private:
    std::unique_ptr<tvm::runtime::Module> m_handle = nullptr;
};

inline float getSimilarity(const cv::Mat &v1,const cv::Mat &v2){
    return static_cast<float>(v1.dot(v2));
}

int main() {
    const std::string paramsPath = "../models/tvm.params";
    const std::string graphPath = "../models/tvm.json";
    const std::string libPath = "../lib/tvm_lib.so";

    TF_TVM tfTvm(paramsPath, graphPath, libPath);

    // Two images of the same person
    const std::string person1imgAPath = "../images/chip1.jpg";
    const std::string person1imgBPath = "../images/chip2.jpg";

    // One image of a different person
    const std::string person2imgPath = "../images/chip17.jpg";

    const auto person1ATempl = tfTvm.getTemplate(person1imgAPath);
    const auto person1BTempl = tfTvm.getTemplate(person1imgBPath);
    const auto person2Templ = tfTvm.getTemplate(person2imgPath);

    // Get score for two images of the same person
    const auto scoreSimilar = getSimilarity(person1ATempl, person1BTempl);
    std::cout << "Similarity score: " << scoreSimilar << "\n";

    // Get score for two images of different people
    const auto scoreDissimilar = getSimilarity(person1ATempl, person2Templ);
    std::cout << "Similarity score: " << scoreDissimilar << "\n";
    
    return 0;
}

TF_TVM::TF_TVM(const std::string& paramsPath, const std::string& graphPath, const std::string& libPath){
    tvm::runtime::Module mod_syslib = tvm::runtime::Module::LoadFromFile(libPath);
    std::ifstream json_in(graphPath);
    std::string json_data((std::istreambuf_iterator<char>(json_in)), std::istreambuf_iterator<char>());
    json_in.close();

    int device_type = kDLCPU;
    int device_id = 0;

    // get global function module for graph runtime
    tvm::runtime::Module mod = (*tvm::runtime::Registry::Get("tvm.graph_runtime.create"))(json_data, mod_syslib, device_type, device_id);
    m_handle = std::make_unique<tvm::runtime::Module>(mod);

    //load param
    std::ifstream params_in(paramsPath, std::ios::binary);
    std::string params_data((std::istreambuf_iterator<char>(params_in)), std::istreambuf_iterator<char>());
    params_in.close();

    TVMByteArray params_arr;
    params_arr.data = params_data.c_str();
    params_arr.size = params_data.length();

    tvm::runtime::PackedFunc load_params = mod.GetFunction("load_params");
    load_params(params_arr);
}

cv::Mat TF_TVM::getTemplate(const std::string& imgPath){
    auto inputImageAligned = cv::imread(imgPath);

    cv::Mat tensor = cv::dnn::blobFromImage(inputImageAligned, 1.0, cv::Size(IMG_WIDTH, IMG_HEIGHT), cv::Scalar(0,0,0), true);
    DLTensor* input;

    constexpr int dtype_code = kDLFloat;
    constexpr int dtype_bits = 32;
    constexpr int dtype_lanes = 1;
    constexpr int device_type = kDLCPU;
    constexpr int device_id = 0;
    constexpr int in_ndim = 4;

    const int64_t in_shape[in_ndim] = {BATCH_SIZE, CHANNELS, IMG_WIDTH, IMG_HEIGHT};
    TVMArrayAlloc(in_shape, in_ndim, dtype_code, dtype_bits, dtype_lanes, device_type, device_id, &input);
    TVMArrayCopyFromBytes(input, tensor.data, IMG_WIDTH * CHANNELS * IMG_HEIGHT * sizeof(float));

    tvm::runtime::PackedFunc set_input = m_handle->GetFunction("set_input");
    set_input("data", input);
    tvm::runtime::PackedFunc run = m_handle->GetFunction("run");
    run();
    
    tvm::runtime::PackedFunc get_output = m_handle->GetFunction("get_output");
    tvm::runtime::NDArray res = get_output(0);

    cv::Mat output(FEATURE_SIZE, 1, CV_32F);
    memcpy(output.data,res->data, FEATURE_SIZE * sizeof(float));

    cv::Mat tmpMat;
    // normlize
    cv::multiply(output, output, tmpMat);
    float l2 =  cv::sqrt(cv::sum(tmpMat).val[0]);
    output = output / l2;

    TVMArrayFree(input);
    return output;
}