본문 바로가기

Deep Learning

[tensorflow] how to load and use CNN in C++

반응형

** tensorflow에서 학습한 network load해서 C++에서 사용하기 **


#include <fstream>

#include <utility>

#include <vector>

#include <stdio.h>

#include <stdlib.h>

#include <sys/shm.h>


#include "tensorflow/cc/ops/const_op.h"

#include "tensorflow/cc/ops/image_ops.h"

#include "tensorflow/cc/ops/standard_ops.h"

#include "tensorflow/core/framework/graph.pb.h"

#include "tensorflow/core/framework/tensor.h"

#include "tensorflow/core/graph/default_device.h"

#include "tensorflow/core/graph/graph_def_builder.h"

#include "tensorflow/core/lib/core/errors.h"

#include "tensorflow/core/lib/core/stringpiece.h"

#include "tensorflow/core/lib/core/threadpool.h"

#include "tensorflow/core/lib/io/path.h"

#include "tensorflow/core/lib/strings/stringprintf.h"

#include "tensorflow/core/platform/env.h"

#include "tensorflow/core/platform/init_main.h"

#include "tensorflow/core/platform/logging.h"

#include "tensorflow/core/platform/types.h"

#include "tensorflow/core/public/session.h"

#include "tensorflow/core/util/command_line_flags.h"


using tensorflow::Flag;

using tensorflow::Tensor;

using tensorflow::Status;

using tensorflow::string;

using tensorflow::int32;

using namespace tensorflow;


/*

something about 'shared->data'

*/



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




/*

something about 'shared->data'

*/



    ///////////////////////////// Initialize a tensorflow session

    Session* session;

    Status status = NewSession(SessionOptions(), &session);

    if (!status.ok()) {

        std::cout << status.ToString() << "\n";

        return 1;

    }


    // Read the protobuf graph we exported

    GraphDef graph_def;

    status = ReadBinaryProto(Env::Default(), "models/graph.pb", &graph_def);

    if (!status.ok()) {

        std::cout << status.ToString() << "\n";

        return 1;

    }


    // Add the graph to the session

    status = session->Create(graph_def);

    if (!status.ok()) {

        std::cout << status.ToString() << "\n";

        return 1;

    }



    ///////////////////////////// convert a RGB image input to a tensor

    const int batch_size = 1;

    const int height = image_height;

    const int width = image_width;

    const int channels = 3;


    // CNN input and output definition

    Tensor tensor(DataType::DT_FLOAT, TensorShape({batch_size, height, width, channels}));

    std::vector<Tensor> outputs;


    // get underlying Eigen tensor

    auto tensor_map = tensor.tensor<float, 4>();



    // write the input image to the pre-defined tensor 

// uint8_t shared->data : an RGB image of size (height x width x channels)

    for (int h = 0; h < height; h++) {

for (int w = 0; w < width; w++) {

              tensor_map(0, h, w, 0) = (float)(shared->data[(h*width+w)*3+0]);

              tensor_map(0, h, w, 1) = (float)(shared->data[(h*width+w)*3+1]);

              tensor_map(0, h, w, 2) = (float)(shared->data[(h*width+w)*3+2]);

        }

    }


    // Run CNN

    status = session->Run({{ "input_node", tensor }}, {"output_node"}, {}, &outputs);

    if (!status.ok()) {

        std::cout << status.ToString() << "\n";

        return 1;

    }


    // Get estimated value (of size (1,1))

    auto output_c = outputs[0].scalar<float>()();


    // Free any resources used by the session

    session->Close();


    return 0;

}