Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improved options handling for TensorFlow sessions #41276

Merged
merged 1 commit into from
Apr 12, 2023
Merged

Improved options handling for TensorFlow sessions #41276

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 40 additions & 47 deletions PhysicsTools/TensorFlow/interface/TensorFlow.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,82 +30,75 @@ namespace tensorflow {
typedef std::pair<std::string, Tensor> NamedTensor;
typedef std::vector<NamedTensor> NamedTensorList;

struct Options {
int _nThreads;
Backend _backend;
SessionOptions _options;

Options(Backend backend) : _nThreads{1}, _backend{backend} {
setThreading(_nThreads);
setBackend(_backend);
};

Options() : _nThreads{1}, _backend{Backend::cpu} {
setThreading(_nThreads);
setBackend(_backend);
};

// updates the config of sessionOptions so that it uses nThreads
void setThreading(int nThreads = 1);

// Set the backend option cpu/cuda
// The gpu memory is set to "allow_growth" to avoid TF getting all the CUDA memory at once.
void setBackend(Backend backend = Backend::cpu);

SessionOptions& getSessionOptions() { return _options; };
int getNThreads() const { return _nThreads; };
Backend getBackend() const { return _backend; };
};

// set the tensorflow log level
void setLogging(const std::string& level = "3");

// updates the config of sessionOptions so that it uses nThreads
void setThreading(SessionOptions& sessionOptions, int nThreads = 1);

// deprecated
// updates the config of sessionOptions so that it uses nThreads, prints a deprecation warning
// since the threading configuration is done per run() call as of 2.1
void setThreading(SessionOptions& sessionOptions, int nThreads, const std::string& singleThreadPool);

// Set the backend option cpu/cuda
// The gpu memory is set to "allow_growth" to avoid TF getting all the CUDA memory at once.
void setBackend(SessionOptions& sessionOptions, Backend backend = Backend::cpu);

// loads a meta graph definition saved at exportDir using the SavedModel interface for a tag and
// predefined sessionOptions
// predefined options
// transfers ownership
MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, SessionOptions& sessionOptions);

// deprecated in favor of loadMetaGraphDef
MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, SessionOptions& sessionOptions);
MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag = kSavedModelTagServe);

// loads a meta graph definition saved at exportDir using the SavedModel interface for a tag and
// nThreads
// user provided options
// transfers ownership
MetaGraphDef* loadMetaGraphDef(const std::string& exportDir,
const std::string& tag = kSavedModelTagServe,
Backend backend = Backend::cpu,
int nThreads = 1);
MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, Options& options);

// deprecated in favor of loadMetaGraphDef
MetaGraphDef* loadMetaGraph(const std::string& exportDir,
const std::string& tag = kSavedModelTagServe,
Backend backend = Backend::cpu,
int nThreads = 1);
MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, Options& Options);

// loads a graph definition saved as a protobuf file at pbFile
// transfers ownership
GraphDef* loadGraphDef(const std::string& pbFile);

// return a new, empty session using predefined sessionOptions
// transfers ownership
Session* createSession(SessionOptions& sessionOptions);
// return a new, empty session using the predefined options
Session* createSession();

// return a new, empty session with nThreads and selected backend
// return a new, empty session using user provided options
// transfers ownership
Session* createSession(Backend backend = Backend::cpu, int nThreads = 1);
Session* createSession(Options& options);

// return a new session that will contain an already loaded meta graph whose exportDir must be
// given in order to load and initialize the variables, sessionOptions are predefined
// an error is thrown when metaGraphDef is a nullptr or when the graph has no nodes
// transfers ownership
Session* createSession(const MetaGraphDef* metaGraphDef,
const std::string& exportDir,
SessionOptions& sessionOptions);

// return a new session that will contain an already loaded meta graph whose exportDir must be given
// in order to load and initialize the variables, threading options are inferred from nThreads
// an error is thrown when metaGraphDef is a nullptr or when the graph has no nodes
// transfers ownership
Session* createSession(const MetaGraphDef* metaGraphDef,
const std::string& exportDir,
Backend backend = Backend::cpu,
int nThreads = 1);
Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, Options& options);

// return a new session that will contain an already loaded graph def, sessionOptions are predefined
// an error is thrown when graphDef is a nullptr or when the graph has no nodes
// transfers ownership
Session* createSession(const GraphDef* graphDef, SessionOptions& sessionOptions);
Session* createSession(const GraphDef* graphDef);

// return a new session that will contain an already loaded graph def, threading options are
// inferred from nThreads
// return a new session that will contain an already loaded graph def, sessionOptions are user defined
// an error is thrown when graphDef is a nullptr or when the graph has no nodes
// transfers ownership
Session* createSession(const GraphDef* graphDef, Backend backend = Backend::cpu, int nThreads = 1);
Session* createSession(const GraphDef* graphDef, Options& options);

// closes a session, calls its destructor, resets the pointer, and returns true on success
bool closeSession(Session*& session);
Expand Down
123 changes: 48 additions & 75 deletions PhysicsTools/TensorFlow/src/TensorFlow.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,14 @@

namespace tensorflow {

void setLogging(const std::string& level) { setenv("TF_CPP_MIN_LOG_LEVEL", level.c_str(), 0); }

void setThreading(SessionOptions& sessionOptions, int nThreads) {
void Options::setThreading(int nThreads) {
_nThreads = nThreads;
// set number of threads used for intra and inter operation communication
sessionOptions.config.set_intra_op_parallelism_threads(nThreads);
sessionOptions.config.set_inter_op_parallelism_threads(nThreads);
}

void setThreading(SessionOptions& sessionOptions, int nThreads, const std::string& singleThreadPool) {
edm::LogInfo("PhysicsTools/TensorFlow") << "setting the thread pool via tensorflow::setThreading() is deprecated";

setThreading(sessionOptions, nThreads);
_options.config.set_intra_op_parallelism_threads(nThreads);
_options.config.set_inter_op_parallelism_threads(nThreads);
}

void setBackend(SessionOptions& sessionOptions, Backend backend) {
void Options::setBackend(Backend backend) {
/*
* The TensorFlow backend configures the available devices using options provided in the sessionOptions proto.
* // Options from https://github.com/tensorflow/tensorflow/blob/c53dab9fbc9de4ea8b1df59041a5ffd3987328c3/tensorflow/core/protobuf/config.proto
Expand All @@ -43,17 +36,17 @@ namespace tensorflow {
edm::Service<edm::ResourceInformation> ri;
if (backend == Backend::cpu) {
// disable GPU usage
(*sessionOptions.config.mutable_device_count())["GPU"] = 0;
sessionOptions.config.mutable_gpu_options()->set_visible_device_list("");
(*_options.config.mutable_device_count())["GPU"] = 0;
_options.config.mutable_gpu_options()->set_visible_device_list("");
}
// NVidia GPU
else if (backend == Backend::cuda) {
if (not ri->nvidiaDriverVersion().empty()) {
// Take only the first GPU in the CUDA_VISIBLE_DEVICE list
(*sessionOptions.config.mutable_device_count())["GPU"] = 1;
sessionOptions.config.mutable_gpu_options()->set_visible_device_list("0");
(*_options.config.mutable_device_count())["GPU"] = 1;
_options.config.mutable_gpu_options()->set_visible_device_list("0");
// Do not allocate all the memory on the GPU at the beginning.
sessionOptions.config.mutable_gpu_options()->set_allow_growth(true);
_options.config.mutable_gpu_options()->set_allow_growth(true);
} else {
edm::Exception ex(edm::errors::UnavailableAccelerator);
ex << "Cuda backend requested, but no NVIDIA GPU available in the job";
Expand All @@ -73,26 +66,41 @@ namespace tensorflow {
// Check if a Nvidia GPU is availabl
if (not ri->nvidiaDriverVersion().empty()) {
// Take only the first GPU in the CUDA_VISIBLE_DEVICE list
(*sessionOptions.config.mutable_device_count())["GPU"] = 1;
sessionOptions.config.mutable_gpu_options()->set_visible_device_list("0");
(*_options.config.mutable_device_count())["GPU"] = 1;
_options.config.mutable_gpu_options()->set_visible_device_list("0");
// Do not allocate all the memory on the GPU at the beginning.
sessionOptions.config.mutable_gpu_options()->set_allow_growth(true);
_options.config.mutable_gpu_options()->set_allow_growth(true);
} else {
// Just CPU support
(*sessionOptions.config.mutable_device_count())["GPU"] = 0;
sessionOptions.config.mutable_gpu_options()->set_visible_device_list("");
(*_options.config.mutable_device_count())["GPU"] = 0;
_options.config.mutable_gpu_options()->set_visible_device_list("");
}
}
}

MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, SessionOptions& sessionOptions) {
void setLogging(const std::string& level) {
/*
* 0 = all messages are logged (default behavior)
* 1 = INFO messages are not printed
* 2 = INFO and WARNING messages are not printed
* 3 = INFO, WARNING, and ERROR messages are not printed
*/
setenv("TF_CPP_MIN_LOG_LEVEL", level.c_str(), 0);
}

MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag) {
Options default_options{};
return loadMetaGraphDef(exportDir, tag, default_options);
}

MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, Options& options) {
// objects to load the graph
Status status;
RunOptions runOptions;
SavedModelBundle bundle;

// load the model
status = LoadSavedModel(sessionOptions, runOptions, exportDir, {tag}, &bundle);
status = LoadSavedModel(options.getSessionOptions(), runOptions, exportDir, {tag}, &bundle);
if (!status.ok()) {
throw cms::Exception("InvalidMetaGraphDef")
<< "error while loading metaGraphDef from '" << exportDir << "': " << status.ToString();
Expand All @@ -102,27 +110,11 @@ namespace tensorflow {
return new MetaGraphDef(bundle.meta_graph_def);
}

MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, SessionOptions& sessionOptions) {
MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, Options& options) {
edm::LogInfo("PhysicsTools/TensorFlow")
<< "tensorflow::loadMetaGraph() is deprecated, use tensorflow::loadMetaGraphDef() instead";

return loadMetaGraphDef(exportDir, tag, sessionOptions);
}

MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, Backend backend, int nThreads) {
// create session options and set thread options
SessionOptions sessionOptions;
setThreading(sessionOptions, nThreads);
setBackend(sessionOptions, backend);

return loadMetaGraphDef(exportDir, tag, sessionOptions);
}

MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, Backend backend, int nThreads) {
edm::LogInfo("PhysicsTools/TensorFlow")
<< "tensorflow::loadMetaGraph() is deprecated, use tensorflow::loadMetaGraphDef() instead";

return loadMetaGraphDef(exportDir, tag, backend, nThreads);
return loadMetaGraphDef(exportDir, tag, options);
}

GraphDef* loadGraphDef(const std::string& pbFile) {
Expand All @@ -142,32 +134,26 @@ namespace tensorflow {
return graphDef;
}

Session* createSession(SessionOptions& sessionOptions) {
Session* createSession() {
Options default_options{};
return createSession(default_options);
}

Session* createSession(Options& options) {
// objects to create the session
Status status;

// create a new, empty session
Session* session = nullptr;
status = NewSession(sessionOptions, &session);
status = NewSession(options.getSessionOptions(), &session);
if (!status.ok()) {
throw cms::Exception("InvalidSession") << "error while creating session: " << status.ToString();
}

return session;
}

Session* createSession(Backend backend, int nThreads) {
// create session options and set thread options
SessionOptions sessionOptions;
setThreading(sessionOptions, nThreads);
setBackend(sessionOptions, backend);

return createSession(sessionOptions);
}

Session* createSession(const MetaGraphDef* metaGraphDef,
const std::string& exportDir,
SessionOptions& sessionOptions) {
Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, Options& options) {
// check for valid pointer
if (metaGraphDef == nullptr) {
throw cms::Exception("InvalidMetaGraphDef") << "error while creating session: metaGraphDef is nullptr";
Expand All @@ -178,7 +164,7 @@ namespace tensorflow {
throw cms::Exception("InvalidMetaGraphDef") << "error while creating session: graphDef has no nodes";
}

Session* session = createSession(sessionOptions);
Session* session = createSession(options);

// add the graph def from the meta graph
Status status;
Expand Down Expand Up @@ -214,16 +200,12 @@ namespace tensorflow {
return session;
}

Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, Backend backend, int nThreads) {
// create session options and set thread options
SessionOptions sessionOptions;
setThreading(sessionOptions, nThreads);
setBackend(sessionOptions, backend);

return createSession(metaGraphDef, exportDir, sessionOptions);
Session* createSession(const GraphDef* graphDef) {
Options default_options{};
return createSession(graphDef, default_options);
}

Session* createSession(const GraphDef* graphDef, SessionOptions& sessionOptions) {
Session* createSession(const GraphDef* graphDef, Options& options) {
// check for valid pointer
if (graphDef == nullptr) {
throw cms::Exception("InvalidGraphDef") << "error while creating session: graphDef is nullptr";
Expand All @@ -235,7 +217,7 @@ namespace tensorflow {
}

// create a new, empty session
Session* session = createSession(sessionOptions);
Session* session = createSession(options);

// add the graph def
Status status;
Expand All @@ -249,15 +231,6 @@ namespace tensorflow {
return session;
}

Session* createSession(const GraphDef* graphDef, Backend backend, int nThreads) {
// create session options and set thread options
SessionOptions sessionOptions;
setThreading(sessionOptions, nThreads);
setBackend(sessionOptions, backend);

return createSession(graphDef, sessionOptions);
}

bool closeSession(Session*& session) {
if (session == nullptr) {
return true;
Expand Down
4 changes: 2 additions & 2 deletions PhysicsTools/TensorFlow/test/testConstSession.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,12 @@ void testConstSession::test() {
tensorflow::Backend backend = tensorflow::Backend::cpu;

// load the graph
tensorflow::setLogging();
tensorflow::Options options{backend};
tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
CPPUNIT_ASSERT(graphDef != nullptr);

// create a new session and add the graphDef
const tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
const tensorflow::Session* session = tensorflow::createSession(graphDef, options);
CPPUNIT_ASSERT(session != nullptr);

// example evaluation
Expand Down
4 changes: 3 additions & 1 deletion PhysicsTools/TensorFlow/test/testConstSessionCUDA.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,11 +55,13 @@ process.add_(cms.Service('CUDAService'))
// load the graph
std::string pbFile = dataPath_ + "/constantgraph.pb";
tensorflow::setLogging();
tensorflow::Options options{backend};

tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
CPPUNIT_ASSERT(graphDef != nullptr);

// create a new session and add the graphDef
const tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
const tensorflow::Session* session = tensorflow::createSession(graphDef, options);
CPPUNIT_ASSERT(session != nullptr);

// example evaluation
Expand Down
6 changes: 3 additions & 3 deletions PhysicsTools/TensorFlow/test/testGraphLoading.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,16 +33,16 @@ void testGraphLoading::test() {
tensorflow::Backend backend = tensorflow::Backend::cpu;

// load the graph
tensorflow::setLogging();
tensorflow::Options options{backend};
tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
CPPUNIT_ASSERT(graphDef != nullptr);

// create a new session and add the graphDef
tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
tensorflow::Session* session = tensorflow::createSession(graphDef, options);
CPPUNIT_ASSERT(session != nullptr);

// check for exception
CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, backend), cms::Exception);
CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, options), cms::Exception);

// example evaluation
tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
Expand Down
Loading