Model.cxx

Go to the documentation of this file.

#include "TensorFlowEvaluator/LSTME/model/Model.h"

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

namespace LSTME
{

const std::vector<InputConfigKeys> Model::sliceInputs({
    { "input_slice", "vars_slice" }
});

const std::vector<InputConfigKeys> Model::prongInputs({
    { "input_png2d", "vars_png2d" },
    { "input_png3d", "vars_png3d" }
});

const std::vector<std::string>     Model::outputs({
    "target_primary", "target_total"
});


Model::Model(const std::string &savedir)
    : simpleModel(savedir, sliceInputs, prongInputs, outputs)
{ }

LSTMEnergy Model::predict(const VarDict &varDict)
{
    std::vector<tensorflow::Tensor> outputs = simpleModel.predict(varDict);

    const float primaryE = outputs[0].tensor<float,2>()(0, 0);
    const float totalE   = outputs[1].tensor<float,2>()(0, 0);

    return LSTMEnergy{ primaryE, totalE };
}

}