158 lines · cpp
1//===- ModelUnderTrainingRunner.cpp - 'development' mode runner -----------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// Implementation of a MLModelRunner for 'development' mode, i.e. evaluation10// happens off a model that's provided from the command line and is interpreted.11//12//===----------------------------------------------------------------------===//13 14#include "llvm/ADT/STLExtras.h"15#include "llvm/Config/config.h"16#if defined(LLVM_HAVE_TFLITE)17#include "llvm/Analysis/ModelUnderTrainingRunner.h"18#include "llvm/Support/MemoryBuffer.h"19#include "llvm/Support/Path.h"20#include <optional>21 22using namespace llvm;23namespace {24struct LoggedFeatureSpec {25 TensorSpec Spec;26 std::optional<std::string> LoggingName;27};28 29std::optional<std::vector<LoggedFeatureSpec>>30loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,31 StringRef ModelPath, StringRef SpecFileOverride) {32 SmallVector<char, 128> OutputSpecsPath;33 StringRef FileName = SpecFileOverride;34 if (FileName.empty()) {35 llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");36 FileName = {OutputSpecsPath.data(), OutputSpecsPath.size()};37 }38 39 auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);40 if (!BufferOrError) {41 Ctx.emitError("Error opening output specs file: " + FileName + " : " +42 BufferOrError.getError().message());43 return std::nullopt;44 }45 auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());46 if (!ParsedJSONValues) {47 Ctx.emitError("Could not parse specs file: " + FileName);48 return std::nullopt;49 }50 auto ValuesArray = ParsedJSONValues->getAsArray();51 if (!ValuesArray) {52 Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "53 "logging_name:<name>} dictionaries");54 return std::nullopt;55 }56 std::vector<LoggedFeatureSpec> Ret;57 for (const auto &Value : *ValuesArray)58 if (const auto *Obj = Value.getAsObject())59 if (const auto *SpecPart = Obj->get("tensor_spec"))60 if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))61 if (auto LoggingName = Obj->getString("logging_name")) {62 if (!TensorSpec->isElementType<int64_t>() &&63 !TensorSpec->isElementType<int32_t>() &&64 !TensorSpec->isElementType<float>()) {65 Ctx.emitError(66 "Only int64, int32, and float tensors are supported. "67 "Found unsupported type for tensor named " +68 TensorSpec->name());69 return std::nullopt;70 }71 Ret.push_back({*TensorSpec, LoggingName->str()});72 }73 74 if (ValuesArray->size() != Ret.size()) {75 Ctx.emitError(76 "Unable to parse output spec. It should be a json file containing an "77 "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "78 "with a json object describing a TensorSpec; and a 'logging_name' key, "79 "which is a string to use as name when logging this tensor in the "80 "training log.");81 return std::nullopt;82 }83 if (Ret.empty() || *Ret[0].LoggingName != ExpectedDecisionName) {84 Ctx.emitError("The first output spec must describe the decision tensor, "85 "and must have the logging_name " +86 StringRef(ExpectedDecisionName));87 return std::nullopt;88 }89 return Ret;90}91} // namespace92 93ModelUnderTrainingRunner::ModelUnderTrainingRunner(94 LLVMContext &Ctx, const std::string &ModelPath,95 const std::vector<TensorSpec> &InputSpecs,96 const std::vector<TensorSpec> &OutputSpecs,97 const std::vector<TensorSpec> &ExtraOutputsForLogging)98 : MLModelRunner(Ctx, MLModelRunner::Kind::Development, InputSpecs.size()),99 OutputSpecs(OutputSpecs), ExtraOutputsForLogging(ExtraOutputsForLogging) {100 Evaluator =101 std::make_unique<TFModelEvaluator>(ModelPath, InputSpecs, OutputSpecs);102 if (!Evaluator || !Evaluator->isValid()) {103 Ctx.emitError("Failed to create saved model evaluator");104 Evaluator.reset();105 return;106 }107 108 for (size_t I = 0, E = InputSpecs.size(); I < E; ++I) {109 setUpBufferForTensor(I, InputSpecs[I], Evaluator->getUntypedInput(I));110 }111}112 113void *ModelUnderTrainingRunner::evaluateUntyped() {114 LastEvaluationResult = Evaluator->evaluate();115 if (!LastEvaluationResult.has_value()) {116 Ctx.emitError("Error evaluating model.");117 return nullptr;118 }119 return LastEvaluationResult->getUntypedTensorValue(0);120}121 122std::unique_ptr<ModelUnderTrainingRunner>123ModelUnderTrainingRunner::createAndEnsureValid(124 LLVMContext &Ctx, const std::string &ModelPath, StringRef DecisionName,125 const std::vector<TensorSpec> &InputSpecs,126 StringRef OutputSpecsPathOverride) {127 if (auto MaybeOutputSpecs = loadOutputSpecs(Ctx, DecisionName, ModelPath,128 OutputSpecsPathOverride)) {129 std::unique_ptr<ModelUnderTrainingRunner> MUTR;130 std::vector<TensorSpec> OutputSpecs;131 std::vector<TensorSpec> ExtraOutputsForLogging;132 append_range(OutputSpecs,133 map_range(*MaybeOutputSpecs, [](const LoggedFeatureSpec &LFS) {134 return LFS.Spec;135 }));136 append_range(ExtraOutputsForLogging,137 map_range(drop_begin(*MaybeOutputSpecs),138 [](const LoggedFeatureSpec &LFS) {139 return TensorSpec(LFS.LoggingName140 ? *LFS.LoggingName141 : LFS.Spec.name(),142 LFS.Spec);143 }));144 145 MUTR.reset(new ModelUnderTrainingRunner(146 Ctx, ModelPath, InputSpecs, OutputSpecs, ExtraOutputsForLogging));147 if (MUTR && MUTR->isValid())148 return MUTR;149 150 Ctx.emitError("Could not load or create model evaluator.");151 return nullptr;152 }153 Ctx.emitError("Could not load the policy model from the provided path");154 return nullptr;155}156 157#endif // defined(LLVM_HAVE_TFLITE)158