Load a TensorFlow graph in C++

example_trainer.cc

/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <cstdio>
#include <functional>
#include <string>
#include <vector>

#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/threadpool.h"
#include "tensorflow/core/lib/strings/stringprintf.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"

using tensorflow::string;
using tensorflow::int32;

namespace tensorflow
{
    namespace example
    {
        /**
        Load the graph from checkpoint and graph definition files.
        Parameters:
        model_path              path to the graph definition file
        checkpoint_name         root checkpoint name. That is, if all
                                checkpoint files are called
                                "models/model_final.x", this should be
                                "models/model_final".
        filename_tensor_name    provided by the Python script. Probably
                                something like "save/Const". Strip any ":0"
                                off the end.
        restore_op_name         provided by the Python script. Probably
                                something like "save/restore_all".
        session                 a TensorFlow session
        scope                   a TensorFlow scope
        **/
        void load(std::string model_path, std::string checkpoint_name,
            std::string filename_tensor_name, std::string restore_op_name,
            std::unique_ptr<Session> &session, tensorflow::Scope &scope)
        {
            // First create the graph specified in the graph def file
            GraphDef graph_def;
            auto load_graph_status = ReadTextProto(tensorflow::Env::Default(), model_path,
                &graph_def);
            if (!load_graph_status.ok())
            {
                std::cout << load_graph_status.ToString() << std::endl;
            }
            auto session_status = session->Create(graph_def);
            if (!session_status.ok())
            {
                std::cout << session_status.ToString() << std::endl;
            }
            
            // Now restore variable values using the Saver ops.
            std::vector<std::pair<std::string, Tensor>> input;
            Tensor filename_tensor = Tensor(DT_STRING, TensorShape({ 1 }));
            filename_tensor.vec<string>()(0) = checkpoint_name;
            input.emplace_back(filename_tensor_name, filename_tensor);
            TF_CHECK_OK(session->Run(input, {}, { restore_op_name }, {}));
        }
    }
}

int main()
{
    tensorflow::SessionOptions options;
    std::unique_ptr<tensorflow::Session> session(NewSession(options));
    tensorflow::Scope root = tensorflow::Scope::NewRootScope();
    std::cout << "loading..." << std::flush;
    tensorflow::example::load("model/graph", "model/model_final",
        "save/Const", "save/restore_all", session, root);
    std::cout << "done." << std::endl;
    
    // Test input
    std::vector<std::pair<std::string, tensorflow::Tensor>> input;
    std::vector<tensorflow::Tensor> output;
    
    std::cout << "filling input tensor..." << std::flush;
    auto input_tensor = tensorflow::Tensor(tensorflow::DT_FLOAT,
        tensorflow::TensorShape({ 1, 72 }));
    auto flat = input_tensor.flat<float>();
    for (int i = 0; i < 72; i++)
    {
        flat(i) = 0;
    }
    flat(71) = 1;
    std::cout << "done." << std::endl;
    
    input.emplace_back("input", input_tensor);
    
    std::cout << "running session..." << std::flush;
    TF_CHECK_OK(session->Run(input, { "output" }, {}, &output));
    std::cout << "done." << std::endl;
    
    std::cout << "printing result..." << std::endl;
    std::cout << "[ ";
    for (int i = 0; i < 72; i++)
    {
        std::cout << output[0].flat<float>()(i) << " ";
    }
    std::cout << "]" << std::endl;
    
    return 0;
}