Skip to content

Commit

Permalink
Improved options struct for TF session
Browse files Browse the repository at this point in the history
  • Loading branch information
@valsdav
valsdav committed Apr 4, 2023
1 parent 7530dfd commit f0fc754
Show file tree
Hide file tree
Showing 17 changed files with 144 additions and 165 deletions.
87 changes: 40 additions & 47 deletions PhysicsTools/TensorFlow/interface/TensorFlow.h
View file
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
125 changes: 50 additions & 75 deletions PhysicsTools/TensorFlow/src/TensorFlow.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,15 @@

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 +37,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 +67,42 @@ 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 +112,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 +136,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 +166,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 +202,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 +219,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 +233,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
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
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
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

0 comments on commit f0fc754

Please sign in to comment.