Skip to content

Commit

Permalink
Refactor tensorflow servable to allow loader sharing with autofill
Browse files Browse the repository at this point in the history
  • Loading branch information
David Goodwin committed Nov 14, 2018
1 parent 653dfd8 commit 3179aa0
Show file tree
Hide file tree
Showing 14 changed files with 706 additions and 458 deletions.
3 changes: 1 addition & 2 deletions src/core/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,8 @@ cc_library(
":logging",
":model_config",
":model_config_proto",
"//src/servables/tensorflow:autofill",
"//src/servables/tensorrt:autofill",
"@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
"@org_tensorflow//tensorflow/c:c_api",
"@org_tensorflow//tensorflow/core:lib",
],
)
Expand Down
292 changes: 1 addition & 291 deletions src/core/autofill.cc
Original file line number Diff line number Diff line change
Expand Up @@ -29,56 +29,13 @@
#include "src/core/constants.h"
#include "src/core/logging.h"
#include "src/core/model_config.h"
#include "src/servables/tensorflow/autofill.h"
#include "src/servables/tensorrt/autofill.h"
#include "tensorflow/c/c_api.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/cc/saved_model/tag_constants.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/env.h"

namespace nvidia { namespace inferenceserver {

namespace {

DataType
ConvertDataType(tensorflow::DataType dtype)
{
switch (dtype) {
case tensorflow::DT_INVALID:
return DataType::TYPE_INVALID;
case tensorflow::DT_BOOL:
return DataType::TYPE_BOOL;
case tensorflow::DT_UINT8:
return DataType::TYPE_UINT8;
case tensorflow::DT_UINT16:
return DataType::TYPE_UINT16;
case tensorflow::DT_UINT32:
return DataType::TYPE_UINT32;
case tensorflow::DT_UINT64:
return DataType::TYPE_UINT64;
case tensorflow::DT_INT8:
return DataType::TYPE_INT8;
case tensorflow::DT_INT16:
return DataType::TYPE_INT16;
case tensorflow::DT_INT32:
return DataType::TYPE_INT32;
case tensorflow::DT_INT64:
return DataType::TYPE_INT64;
case tensorflow::DT_HALF:
return DataType::TYPE_FP16;
case tensorflow::DT_FLOAT:
return DataType::TYPE_FP32;
case tensorflow::DT_DOUBLE:
return DataType::TYPE_FP64;
default:
break;
}

return DataType::TYPE_INVALID;
}

} // namespace

//
// AutoFillNull
//
Expand Down Expand Up @@ -136,253 +93,6 @@ AutoFillSimple::Fix(ModelConfig* config)
return tensorflow::Status::OK();
}

//
// AutoFillSavedModel
//
class AutoFillSavedModel : public AutoFill {
public:
static tensorflow::Status Create(
const std::string& model_name, const std::string& model_path,
std::unique_ptr<AutoFillSavedModel>* autofill);
tensorflow::Status Fix(ModelConfig* config);

private:
AutoFillSavedModel(
const std::string& model_name, const std::string& savedmodel_dirname,
const tensorflow::SignatureDef& sig)
: AutoFill(model_name), savedmodel_dirname_(savedmodel_dirname), sig_(sig)
{
}

const std::string savedmodel_dirname_;
const tensorflow::SignatureDef sig_;
};

tensorflow::Status
AutoFillSavedModel::Create(
const std::string& model_name, const std::string& model_path,
std::unique_ptr<AutoFillSavedModel>* autofill)
{
std::set<std::string> version_dirs;
TF_RETURN_IF_ERROR(GetSubdirs(model_path, &version_dirs));

// There must be at least one version directory that we can inspect
// to attempt to determine the platform. For now we only handle the
// case where there is one version directory.
if (version_dirs.size() != 1) {
return tensorflow::errors::Internal(
"unable to autofill for '", model_name, "' due to multiple versions");
}

const auto version_path =
tensorflow::io::JoinPath(model_path, *(version_dirs.begin()));

// There must be a single savedmodel directory within the version
// directory...
std::set<std::string> savedmodel_dirs;
TF_RETURN_IF_ERROR(GetSubdirs(version_path, &savedmodel_dirs));
if (savedmodel_dirs.size() != 1) {
return tensorflow::errors::Internal(
"unable to autofill for '", model_name,
"', unable to find savedmodel directory");
}

const std::string savedmodel_dir = *(savedmodel_dirs.begin());
const auto savedmodel_path =
tensorflow::io::JoinPath(version_path, savedmodel_dir);
std::unique_ptr<tensorflow::SavedModelBundle> bundle(
new tensorflow::SavedModelBundle);

std::unordered_set<std::string> saved_model_tags;
saved_model_tags.insert(tensorflow::kSavedModelTagServe);

tensorflow::SessionOptions session_options;
tensorflow::RunOptions run_options;
TF_RETURN_IF_ERROR(tensorflow::LoadSavedModel(
session_options, run_options, savedmodel_path, saved_model_tags,
bundle.get()));

// Only autofill if there is a "serve" tag...
bool found_serve_tag = false;
for (const auto& tag : bundle->meta_graph_def.meta_info_def().tags()) {
if (tag == tensorflow::kSavedModelTagServe) {
found_serve_tag = true;
break;
}
}
if (!found_serve_tag) {
return tensorflow::errors::Internal(
"unable to autofill for '", model_name, "', unable to find '",
tensorflow::kSavedModelTagServe, "' tag");
}

// Must have a "serving_default" signature as that is what will be
// used to determine the inputs and outputs.
static const std::string DEFAULT_SERVING_SIGNATURE_DEF_KEY("serving_default");
const auto& sig_itr = bundle->meta_graph_def.signature_def().find(
DEFAULT_SERVING_SIGNATURE_DEF_KEY);
if (sig_itr == bundle->meta_graph_def.signature_def().end()) {
return tensorflow::errors::InvalidArgument(
"unable to autofill for '", model_name, "', require '",
DEFAULT_SERVING_SIGNATURE_DEF_KEY, "' signature");
}

// Save the name of the savedmodel directory since it may not be the
// default value. Save the SignatureDef as that is what we need to
// autofill inputs and outputs.
const tensorflow::SignatureDef& sig = sig_itr->second;

autofill->reset(new AutoFillSavedModel(model_name, savedmodel_dir, sig));
return tensorflow::Status::OK();
}

tensorflow::Status
AutoFillSavedModel::Fix(ModelConfig* config)
{
config->set_platform(kTensorFlowSavedModelPlatform);

if (config->name().empty()) {
config->set_name(model_name_);
}

if (config->default_model_filename().empty()) {
config->set_default_model_filename(savedmodel_dirname_);
}

// Assume model doesn't support batching unless we see a batch
// dimension in the input or output.
bool supports_batch = false;

// Inputs
if (config->input().size() == 0) {
for (const auto& sin : sig_.inputs()) {
ModelInput* config_input = config->add_input();
config_input->set_name(sin.first);
config_input->set_data_type(ConvertDataType(sin.second.dtype()));


// The first model dimension can be -1 to serve as a placeholder
// for batch. The batch dim doesn't appear in the configuration
// 'dims'.
const tensorflow::TensorShapeProto& shape = sin.second.tensor_shape();
const bool has_batch_dim =
(shape.dim().size() >= 1) && (shape.dim(0).size() == -1);

for (int i = (has_batch_dim ? 1 : 0); i < shape.dim().size(); ++i) {
config_input->mutable_dims()->Add(shape.dim(i).size());
}

supports_batch |= has_batch_dim;
}
}

// Outputs
if (config->output().size() == 0) {
for (const auto& sout : sig_.outputs()) {
ModelOutput* config_output = config->add_output();
config_output->set_name(sout.first);
config_output->set_data_type(ConvertDataType(sout.second.dtype()));


// The first model dimension can be -1 to serve as a placeholder
// for batch. The batch dim doesn't appear in the configuration
// 'dims'.
const tensorflow::TensorShapeProto& shape = sout.second.tensor_shape();
const bool has_batch_dim =
(shape.dim().size() >= 1) && (shape.dim(0).size() == -1);

for (int i = (has_batch_dim ? 1 : 0); i < shape.dim().size(); ++i) {
config_output->mutable_dims()->Add(shape.dim(i).size());
}

supports_batch |= has_batch_dim;
}
}

// Set max-batch-size to 1 if the model supports it and it is not
// already set.
if (supports_batch && (config->max_batch_size() == 0)) {
config->set_max_batch_size(1);
}

return tensorflow::Status::OK();
}

//
// AutoFillGraphDef
//
class AutoFillGraphDef : public AutoFill {
public:
static tensorflow::Status Create(
const std::string& model_name, const std::string& model_path,
std::unique_ptr<AutoFillGraphDef>* autofill);
tensorflow::Status Fix(ModelConfig* config);

private:
AutoFillGraphDef(const std::string& model_name) : AutoFill(model_name) {}
};

tensorflow::Status
AutoFillGraphDef::Create(
const std::string& model_name, const std::string& model_path,
std::unique_ptr<AutoFillGraphDef>* autofill)
{
std::set<std::string> version_dirs;
TF_RETURN_IF_ERROR(GetSubdirs(model_path, &version_dirs));

// There must be at least one version directory that we can inspect
// to attempt to determine the platform. For now we only handle the
// case where there is one version directory.
if (version_dirs.size() != 1) {
return tensorflow::errors::Internal(
"unable to autofill for '", model_name, "' due to multiple versions");
}

const auto version_path =
tensorflow::io::JoinPath(model_path, *(version_dirs.begin()));

// There must be a single graphdef file within the version
// directory...
std::set<std::string> graphdef_files;
TF_RETURN_IF_ERROR(GetFiles(version_path, &graphdef_files));
if (graphdef_files.size() != 1) {
return tensorflow::errors::Internal(
"unable to autofill for '", model_name,
"', unable to find graphdef file");
}

const std::string graphdef_file = *(graphdef_files.begin());
const auto graphdef_path =
tensorflow::io::JoinPath(version_path, graphdef_file);

// If find a file named with the default graphdef name then assume
// it is a graphdef. We could be smarter here and try to parse to
// see if it really is a graphdef. We could also guess thae
// placeholders are inputs... but we have no way to know what the
// outputs are.
if (graphdef_file != kTensorFlowGraphDefFilename) {
return tensorflow::errors::Internal(
"unable to autofill for '", model_name,
"', unable to find graphdef file named '", kTensorFlowGraphDefFilename,
"'");
}

autofill->reset(new AutoFillGraphDef(model_name));
return tensorflow::Status::OK();
}

tensorflow::Status
AutoFillGraphDef::Fix(ModelConfig* config)
{
config->set_platform(kTensorFlowGraphDefPlatform);

if (config->name().empty()) {
config->set_name(model_name_);
}

return tensorflow::Status::OK();
}

//
// AutoFillNetDef
//
Expand Down
42 changes: 41 additions & 1 deletion src/servables/tensorflow/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,49 @@ serving_proto_library(
],
)

cc_library(
name = "autofill",
srcs = ["autofill.cc"],
hdrs = ["autofill.h"],
deps = [
":loader",
":tf_utils",
"//src/core:autofill_header",
"//src/core:logging",
"@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
"@org_tensorflow//tensorflow/c:c_api",
"@org_tensorflow//tensorflow/core:lib",
],
)

cc_library(
name = "tf_utils",
srcs = ["tf_utils.cc"],
hdrs = ["tf_utils.h"],
deps = [
"//src/core:model_config",
"//src/core:model_config_proto",
"@org_tensorflow//tensorflow/core:lib",
],
)

cc_library(
name = "loader",
srcs = ["loader.cc"],
hdrs = ["loader.h"],
deps = [
"@org_tensorflow//tensorflow/c:c_api",
"@org_tensorflow//tensorflow/core:lib",
"@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
],
)

cc_library(
name = "base_bundle",
srcs = ["base_bundle.cc"],
hdrs = ["base_bundle.h"],
deps = [
":tf_utils",
"//src/core:constants",
"//src/core:infer",
"//src/core:label_provider",
Expand All @@ -73,6 +111,7 @@ cc_library(
hdrs = ["graphdef_bundle.h"],
deps = [
":base_bundle",
":tf_utils",
"//src/core:model_config",
],
)
Expand Down Expand Up @@ -106,8 +145,9 @@ cc_library(
hdrs = ["savedmodel_bundle.h"],
deps = [
":base_bundle",
":loader",
":tf_utils",
"//src/core:model_config",
"@org_tensorflow//tensorflow/cc/saved_model:tag_constants",
],
)

Expand Down
Loading

0 comments on commit 3179aa0

Please sign in to comment.