You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
506 lines
18 KiB
506 lines
18 KiB
//===- TFUtils.cpp - tensorflow evaluation utilities ----------------------===//
|
|
//
|
|
// The LLVM Compiler Infrastructure
|
|
//
|
|
// This file is distributed under the University of Illinois Open Source
|
|
// License. See LICENSE.TXT for details.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements utilities for interfacing with tensorflow C APIs.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
#include "llvm/Config/config.h"
|
|
#if defined(LLVM_HAVE_TF_API)
|
|
|
|
#include "llvm/ADT/Twine.h"
|
|
#include "llvm/Analysis/Utils/TFUtils.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include "llvm/Support/JSON.h"
|
|
#include "llvm/Support/ManagedStatic.h"
|
|
#include "llvm/Support/MemoryBuffer.h"
|
|
#include "llvm/Support/Path.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
|
|
#include "tensorflow/c/c_api.h"
|
|
#include "tensorflow/c/c_api_experimental.h"
|
|
|
|
#include <cassert>
|
|
#include <numeric>
|
|
|
|
using namespace llvm;
|
|
|
|
namespace {
|
|
|
|
using TFGraphPtr = std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)>;
|
|
using TFSessionOptionsPtr =
|
|
std::unique_ptr<TF_SessionOptions, decltype(&TF_DeleteSessionOptions)>;
|
|
using TFStatusPtr = std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)>;
|
|
|
|
struct TFInitializer {
|
|
TFInitializer() {
|
|
assert(!IsInitialized && "TFInitialized should be called only once");
|
|
int Argc = 1;
|
|
const char *Name = "";
|
|
const char **NamePtr = &Name;
|
|
TF_InitMain(Name, &Argc, const_cast<char ***>(&NamePtr));
|
|
IsInitialized = true;
|
|
}
|
|
bool IsInitialized = false;
|
|
};
|
|
|
|
llvm::ManagedStatic<TFInitializer> TFLibInitializer;
|
|
|
|
bool ensureInitTF() { return TFLibInitializer->IsInitialized; }
|
|
|
|
TFGraphPtr createTFGraph() {
|
|
return TFGraphPtr(TF_NewGraph(), &TF_DeleteGraph);
|
|
}
|
|
|
|
TFStatusPtr createTFStatus() {
|
|
return TFStatusPtr(TF_NewStatus(), &TF_DeleteStatus);
|
|
}
|
|
|
|
TFSessionOptionsPtr createTFSessionOptions() {
|
|
return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions);
|
|
}
|
|
|
|
/// Write the values of one tensor as a list.
|
|
template <typename T>
|
|
void writeTensorValues(raw_ostream &OutFile, const char *TensorData,
|
|
size_t ElemCount) {
|
|
OutFile << "[";
|
|
const T *TypedData = reinterpret_cast<const T *>(TensorData);
|
|
for (size_t I = 0; I < ElemCount; ++I) {
|
|
if (I > 0)
|
|
OutFile << ", ";
|
|
OutFile << TypedData[I];
|
|
}
|
|
OutFile << "]";
|
|
}
|
|
|
|
/// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
|
|
/// The tensors are assumed to be stored contiguously, in row-major format,
|
|
/// in the TensorData buffer. Each tensor has the shape given by Spec. The
|
|
/// feature name in the output is either the provided LoggingName, if
|
|
/// specified, otherwise it's the name of the tensor (as given by Spec).
|
|
void writeRawTensorsAsFeatureLists(raw_ostream &OutFile,
|
|
const LoggedFeatureSpec &LoggedSpec,
|
|
const char *TensorData, size_t TensorCount,
|
|
bool FinalReward = false) {
|
|
const char *FieldName = "<invalid>";
|
|
std::function<void(const char *)> ValueWriter;
|
|
const auto &Spec = LoggedSpec.Spec;
|
|
// The 'Feature' protobuf only has 3 possible fields: float_list,
|
|
// int64_list, or bytes_list, so we capture int32 values as int64. We don't
|
|
// support any other types.
|
|
if (Spec.isElementType<int64_t>()) {
|
|
FieldName = "int64_list";
|
|
ValueWriter = [&](const char *Data) {
|
|
writeTensorValues<int64_t>(OutFile, Data, Spec.getElementCount());
|
|
};
|
|
} else if (Spec.isElementType<int32_t>()) {
|
|
FieldName = "int64_list";
|
|
ValueWriter = [&](const char *Data) {
|
|
writeTensorValues<int32_t>(OutFile, Data, Spec.getElementCount());
|
|
};
|
|
|
|
} else if (Spec.isElementType<float>()) {
|
|
FieldName = "float_list";
|
|
ValueWriter = [&](const char *Data) {
|
|
writeTensorValues<float>(OutFile, Data, Spec.getElementCount());
|
|
};
|
|
|
|
} else {
|
|
llvm_unreachable("Unsupported tensor type.");
|
|
}
|
|
|
|
OutFile << " feature_list: {\n";
|
|
OutFile << " key: "
|
|
<< "\""
|
|
<< (LoggedSpec.LoggingName ? *LoggedSpec.LoggingName : Spec.name())
|
|
<< "\" ";
|
|
OutFile << "value: {\n";
|
|
size_t TensorByteSize = Spec.getElementCount() * Spec.getElementByteSize();
|
|
|
|
auto WriteFeatureProto = [&](const char *P) {
|
|
OutFile << " feature: { " << FieldName << ": { value: ";
|
|
ValueWriter(P);
|
|
OutFile << " } }\n";
|
|
};
|
|
|
|
const char *CurrentTensor = TensorData;
|
|
static int64_t Zero = 0;
|
|
// Write all but the last value. If this is the final reward, don't increment
|
|
// the CurrentTensor, and just write 0.
|
|
for (size_t I = 0; I < TensorCount - 1; ++I) {
|
|
if (FinalReward)
|
|
WriteFeatureProto(reinterpret_cast<const char *>(&Zero));
|
|
else {
|
|
WriteFeatureProto(CurrentTensor);
|
|
CurrentTensor += TensorByteSize;
|
|
}
|
|
}
|
|
|
|
WriteFeatureProto(CurrentTensor);
|
|
|
|
OutFile << " }\n";
|
|
OutFile << " }\n";
|
|
}
|
|
} // namespace
|
|
|
|
namespace llvm {
|
|
class EvaluationResultImpl {
|
|
public:
|
|
EvaluationResultImpl(size_t OutputSize)
|
|
: OutputSize(OutputSize), Output(OutputSize){};
|
|
|
|
~EvaluationResultImpl() {
|
|
for (auto *P : Output)
|
|
if (P)
|
|
TF_DeleteTensor(P);
|
|
}
|
|
|
|
EvaluationResultImpl(const EvaluationResultImpl &) = delete;
|
|
EvaluationResultImpl(EvaluationResultImpl &&Other) = delete;
|
|
std::vector<TF_Tensor *> &getOutput() { return Output; }
|
|
|
|
private:
|
|
const size_t OutputSize;
|
|
std::vector<TF_Tensor *> Output;
|
|
};
|
|
|
|
size_t TensorSpec::getElementByteSize() const {
|
|
return TF_DataTypeSize(static_cast<TF_DataType>(TypeIndex));
|
|
}
|
|
|
|
TensorSpec::TensorSpec(const std::string &Name, int Port, int TypeIndex,
|
|
const std::vector<int64_t> &Shape)
|
|
: Name(Name), Port(Port), TypeIndex(TypeIndex), Shape(Shape),
|
|
ElementCount(std::accumulate(Shape.begin(), Shape.end(), 1,
|
|
std::multiplies<int64_t>())) {}
|
|
|
|
Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
|
|
const json::Value &Value) {
|
|
auto EmitError = [&](const llvm::Twine &Message) -> Optional<TensorSpec> {
|
|
std::string S;
|
|
llvm::raw_string_ostream OS(S);
|
|
OS << Value;
|
|
Ctx.emitError("Unable to parse JSON Value as spec (" + Message + "): " + S);
|
|
return None;
|
|
};
|
|
// FIXME: accept a Path as a parameter, and use it for error reporting.
|
|
json::Path::Root Root("tensor_spec");
|
|
json::ObjectMapper Mapper(Value, Root);
|
|
if (!Mapper)
|
|
return EmitError("Value is not a dict");
|
|
|
|
std::string TensorName;
|
|
int TensorPort = -1;
|
|
std::string TensorType;
|
|
std::vector<int64_t> TensorShape;
|
|
|
|
if (!Mapper.map<std::string>("name", TensorName))
|
|
return EmitError("'name' property not present or not a string");
|
|
if (!Mapper.map<std::string>("type", TensorType))
|
|
return EmitError("'type' property not present or not a string");
|
|
if (!Mapper.map<int>("port", TensorPort))
|
|
return EmitError("'port' property not present or not an int");
|
|
if (!Mapper.map<std::vector<int64_t>>("shape", TensorShape))
|
|
return EmitError("'shape' property not present or not an int array");
|
|
|
|
#define PARSE_TYPE(T, E) \
|
|
if (TensorType == #T) \
|
|
return TensorSpec::createSpec<T>(TensorName, TensorShape, TensorPort);
|
|
TFUTILS_SUPPORTED_TYPES(PARSE_TYPE)
|
|
#undef PARSE_TYPE
|
|
return None;
|
|
}
|
|
|
|
Optional<std::vector<LoggedFeatureSpec>>
|
|
loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
|
|
StringRef ModelPath, StringRef SpecFileOverride) {
|
|
SmallVector<char, 128> OutputSpecsPath;
|
|
StringRef FileName = SpecFileOverride;
|
|
if (FileName.empty()) {
|
|
llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
|
|
FileName = {OutputSpecsPath.data(), OutputSpecsPath.size()};
|
|
}
|
|
|
|
auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
|
|
if (!BufferOrError) {
|
|
Ctx.emitError("Error opening output specs file: " + FileName + " : " +
|
|
BufferOrError.getError().message());
|
|
return None;
|
|
}
|
|
auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
|
|
if (!ParsedJSONValues) {
|
|
Ctx.emitError("Could not parse specs file: " + FileName);
|
|
return None;
|
|
}
|
|
auto ValuesArray = ParsedJSONValues->getAsArray();
|
|
if (!ValuesArray) {
|
|
Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
|
|
"logging_name:<name>} dictionaries");
|
|
return None;
|
|
}
|
|
std::vector<LoggedFeatureSpec> Ret;
|
|
for (const auto &Value : *ValuesArray)
|
|
if (const auto *Obj = Value.getAsObject())
|
|
if (const auto *SpecPart = Obj->get("tensor_spec"))
|
|
if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
|
|
if (auto LoggingName = Obj->getString("logging_name")) {
|
|
if (!TensorSpec->isElementType<int64_t>() &&
|
|
!TensorSpec->isElementType<int32_t>() &&
|
|
!TensorSpec->isElementType<float>()) {
|
|
Ctx.emitError(
|
|
"Only int64, int32, and float tensors are supported. "
|
|
"Found unsupported type for tensor named " +
|
|
TensorSpec->name());
|
|
return None;
|
|
}
|
|
Ret.push_back({*TensorSpec, LoggingName->str()});
|
|
}
|
|
|
|
if (ValuesArray->size() != Ret.size()) {
|
|
Ctx.emitError(
|
|
"Unable to parse output spec. It should be a json file containing an "
|
|
"array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
|
|
"with a json object describing a TensorSpec; and a 'logging_name' key, "
|
|
"which is a string to use as name when logging this tensor in the "
|
|
"training log.");
|
|
return None;
|
|
}
|
|
if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) {
|
|
Ctx.emitError("The first output spec must describe the decision tensor, "
|
|
"and must have the logging_name " +
|
|
StringRef(ExpectedDecisionName));
|
|
return None;
|
|
}
|
|
return Ret;
|
|
}
|
|
|
|
class TFModelEvaluatorImpl {
|
|
public:
|
|
TFModelEvaluatorImpl(StringRef SavedModelPath,
|
|
const std::vector<TensorSpec> &InputSpecs,
|
|
function_ref<TensorSpec(size_t)> GetOutputSpecs,
|
|
size_t OutputSpecsSize, const char *Tags);
|
|
|
|
bool isValid() const { return IsValid; }
|
|
size_t OutputSize() const { return OutputFeed.size(); }
|
|
|
|
void evaluate(TF_Tensor **Output, TF_Status *Status) {
|
|
TF_SessionRun(Session, nullptr, InputFeed.data(), Input.data(),
|
|
Input.size(), OutputFeed.data(), Output, OutputFeed.size(),
|
|
nullptr, 0, nullptr, Status);
|
|
}
|
|
|
|
void initInput(size_t Index, TF_DataType Type,
|
|
const std::vector<int64_t> &Dimensions);
|
|
const std::vector<TF_Tensor *> &getInput() const { return Input; }
|
|
|
|
~TFModelEvaluatorImpl();
|
|
|
|
private:
|
|
/// The objects necessary for carrying out an evaluation of the SavedModel.
|
|
/// They are expensive to set up, and we maintain them accross all the
|
|
/// evaluations of the model.
|
|
TF_Session *Session = nullptr;
|
|
TFGraphPtr Graph;
|
|
TFSessionOptionsPtr Options;
|
|
|
|
/// The specification of the input nodes.
|
|
std::vector<TF_Output> InputFeed;
|
|
|
|
/// The input tensors. They must match by index of the corresponding InputFeed
|
|
/// value. We set up the tensors once and just mutate theirs scalars before
|
|
/// each evaluation. The input tensors keep their value after an evaluation.
|
|
std::vector<TF_Tensor *> Input;
|
|
|
|
/// The specification of the output nodes. When evaluating, the tensors in the
|
|
/// output tensor vector must match by index the corresponding element in the
|
|
/// OutputFeed.
|
|
std::vector<TF_Output> OutputFeed;
|
|
|
|
void invalidate() { IsValid = false; }
|
|
|
|
bool IsValid = true;
|
|
|
|
/// Reusable utility for ensuring we can bind the requested Name to a node in
|
|
/// the SavedModel Graph.
|
|
bool checkReportAndInvalidate(const TF_Output &Output,
|
|
const TensorSpec &OutputSpec);
|
|
};
|
|
} // namespace llvm
|
|
|
|
TFModelEvaluatorImpl::TFModelEvaluatorImpl(
|
|
StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
|
|
function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
|
|
const char *Tags = "serve")
|
|
: Graph(createTFGraph()), Options(createTFSessionOptions()),
|
|
InputFeed(InputSpecs.size()), Input(InputSpecs.size()),
|
|
OutputFeed(OutputSpecsSize) {
|
|
if (!ensureInitTF()) {
|
|
errs() << "Tensorflow should have been initialized";
|
|
return;
|
|
}
|
|
auto Status = createTFStatus();
|
|
|
|
Session = TF_LoadSessionFromSavedModel(Options.get(), nullptr,
|
|
SavedModelPath.str().c_str(), &Tags, 1,
|
|
Graph.get(), nullptr, Status.get());
|
|
if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
|
|
errs() << TF_Message(Status.get());
|
|
invalidate();
|
|
}
|
|
for (size_t I = 0; I < InputSpecs.size(); ++I) {
|
|
auto &InputSpec = InputSpecs[I];
|
|
InputFeed[I] = {
|
|
TF_GraphOperationByName(Graph.get(), (InputSpec.name()).c_str()),
|
|
InputSpec.port()};
|
|
if (!checkReportAndInvalidate(InputFeed[I], InputSpec))
|
|
return;
|
|
initInput(I, static_cast<TF_DataType>(InputSpec.typeIndex()),
|
|
InputSpec.shape());
|
|
}
|
|
for (size_t I = 0; I < OutputSpecsSize; ++I) {
|
|
auto OutputSpec = GetOutputSpecs(I);
|
|
OutputFeed[I] = {
|
|
TF_GraphOperationByName(Graph.get(), (OutputSpec.name()).c_str()),
|
|
OutputSpec.port()};
|
|
if (!checkReportAndInvalidate(OutputFeed[I], OutputSpec))
|
|
return;
|
|
}
|
|
}
|
|
|
|
TFModelEvaluator::TFModelEvaluator(
|
|
StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
|
|
function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
|
|
const char *Tags)
|
|
: Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, GetOutputSpecs,
|
|
OutputSpecsSize, Tags)) {
|
|
if (!Impl->isValid())
|
|
Impl.reset();
|
|
}
|
|
|
|
TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
|
|
const std::vector<TensorSpec> &InputSpecs,
|
|
const std::vector<TensorSpec> &OutputSpecs,
|
|
const char *Tags)
|
|
: TFModelEvaluator(
|
|
SavedModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I]; },
|
|
OutputSpecs.size(), Tags) {}
|
|
|
|
TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {
|
|
for (auto *T : Input) {
|
|
TF_DeleteTensor(T);
|
|
}
|
|
if (Session == nullptr)
|
|
return;
|
|
auto Status = createTFStatus();
|
|
TF_DeleteSession(Session, Status.get());
|
|
Session = nullptr;
|
|
if (TF_GetCode(Status.get()) != TF_Code::TF_OK)
|
|
errs() << "Could not delete TF session";
|
|
}
|
|
|
|
bool TFModelEvaluatorImpl::checkReportAndInvalidate(
|
|
const TF_Output &Output, const TensorSpec &OutputSpec) {
|
|
if (Output.oper)
|
|
return true;
|
|
errs() << "Could not find TF_Output named: " + OutputSpec.name();
|
|
IsValid = false;
|
|
return IsValid;
|
|
}
|
|
|
|
Optional<TFModelEvaluator::EvaluationResult> TFModelEvaluator::evaluate() {
|
|
if (!isValid())
|
|
return None;
|
|
std::unique_ptr<EvaluationResultImpl> Ret =
|
|
std::make_unique<EvaluationResultImpl>(Impl->OutputSize());
|
|
auto Status = createTFStatus();
|
|
Impl->evaluate(Ret->getOutput().data(), Status.get());
|
|
if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
|
|
errs() << TF_Message(Status.get());
|
|
Impl.reset();
|
|
return None;
|
|
}
|
|
return EvaluationResult(std::move(Ret));
|
|
}
|
|
|
|
void TFModelEvaluatorImpl::initInput(size_t Index, TF_DataType Type,
|
|
const std::vector<int64_t> &Dimensions) {
|
|
int64_t TotalSize = TF_DataTypeSize(Type);
|
|
for (auto &D : Dimensions)
|
|
TotalSize *= D;
|
|
|
|
Input[Index] =
|
|
TF_AllocateTensor(Type, Dimensions.data(), Dimensions.size(), TotalSize);
|
|
std::memset(TF_TensorData(Input[Index]), 0, TotalSize);
|
|
}
|
|
|
|
void *TFModelEvaluator::getUntypedInput(size_t Index) {
|
|
return TF_TensorData(Impl->getInput()[Index]);
|
|
}
|
|
|
|
TFModelEvaluator::EvaluationResult::EvaluationResult(
|
|
std::unique_ptr<EvaluationResultImpl> Impl)
|
|
: Impl(std::move(Impl)) {}
|
|
|
|
TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
|
|
: Impl(std::move(Other.Impl)) {}
|
|
|
|
TFModelEvaluator::EvaluationResult &
|
|
TFModelEvaluator::EvaluationResult::operator=(EvaluationResult &&Other) {
|
|
Impl = std::move(Other.Impl);
|
|
return *this;
|
|
}
|
|
|
|
void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
|
|
return TF_TensorData(Impl->getOutput()[Index]);
|
|
}
|
|
|
|
const void *
|
|
TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) const {
|
|
return TF_TensorData(Impl->getOutput()[Index]);
|
|
}
|
|
|
|
#define TFUTILS_GETDATATYPE_IMPL(T, E) \
|
|
template <> int TensorSpec::getDataType<T>() { return E; }
|
|
|
|
TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL)
|
|
|
|
#undef TFUTILS_GETDATATYPE_IMPL
|
|
|
|
TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
|
|
TFModelEvaluator::~TFModelEvaluator() {}
|
|
|
|
void Logger::print(raw_ostream &OS) {
|
|
if (RawLogData.empty())
|
|
return;
|
|
if (RawLogData[0].empty())
|
|
return;
|
|
size_t Tensor0Size = FeatureSpecs[0].Spec.getElementCount() *
|
|
FeatureSpecs[0].Spec.getElementByteSize();
|
|
size_t NumberOfRecords = RawLogData[0].size() / Tensor0Size;
|
|
if (NumberOfRecords == 0)
|
|
return;
|
|
size_t RewardSize =
|
|
RewardSpec.getElementCount() * RewardSpec.getElementByteSize();
|
|
size_t NumberOfRewards = RawLogData.back().size() / RewardSize;
|
|
|
|
OS << "feature_lists: {\n";
|
|
for (size_t I = 0; I < FeatureSpecs.size(); ++I)
|
|
writeRawTensorsAsFeatureLists(OS, FeatureSpecs[I], RawLogData[I].data(),
|
|
NumberOfRecords);
|
|
|
|
if (IncludeReward)
|
|
writeRawTensorsAsFeatureLists(OS, {RewardSpec, None},
|
|
RawLogData.back().data(), NumberOfRecords,
|
|
NumberOfRewards == 1);
|
|
|
|
OS << "}\n";
|
|
}
|
|
#endif // defined(LLVM_HAVE_TF_API)
|