Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-10-25 09:58:36

0001 /*
0002  * Tests for loading graphs via the converted protobuf files.
0003  * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
0004  *
0005  * Author: Marcel Rieger
0006  */
0007 
0008 #include <stdexcept>
0009 #include <cppunit/extensions/HelperMacros.h>
0010 
0011 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0012 
0013 #include "testBase.h"
0014 
0015 class testGraphLoading : public testBase {
0016   CPPUNIT_TEST_SUITE(testGraphLoading);
0017   CPPUNIT_TEST(test);
0018   CPPUNIT_TEST_SUITE_END();
0019 
0020 public:
0021   std::string pyScript() const override;
0022   void test() override;
0023 };
0024 
0025 CPPUNIT_TEST_SUITE_REGISTRATION(testGraphLoading);
0026 
0027 std::string testGraphLoading::pyScript() const { return "createconstantgraph.py"; }
0028 
0029 void testGraphLoading::test() {
0030   std::string pbFile = dataPath_ + "/constantgraph.pb";
0031 
0032   std::cout << "Testing CPU backend" << std::endl;
0033   tensorflow::Backend backend = tensorflow::Backend::cpu;
0034 
0035   // load the graph
0036   tensorflow::Options options{backend};
0037   tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
0038   CPPUNIT_ASSERT(graphDef != nullptr);
0039 
0040   // create a new session and add the graphDef
0041   tensorflow::Session* session = tensorflow::createSession(graphDef, options);
0042   CPPUNIT_ASSERT(session != nullptr);
0043 
0044   // check for exception
0045   CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, options), cms::Exception);
0046 
0047   // example evaluation
0048   tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
0049   float* d = input.flat<float>().data();
0050   for (size_t i = 0; i < 10; i++, d++) {
0051     *d = float(i);
0052   }
0053   tensorflow::Tensor scale(tensorflow::DT_FLOAT, {});
0054   scale.scalar<float>()() = 1.0;
0055 
0056   std::vector<tensorflow::Tensor> outputs;
0057   tensorflow::Status status = session->Run({{"input", input}, {"scale", scale}}, {"output"}, {}, &outputs);
0058   if (!status.ok()) {
0059     std::cout << status.ToString() << std::endl;
0060     CPPUNIT_ASSERT(false);
0061   }
0062 
0063   // check the output
0064   CPPUNIT_ASSERT(outputs.size() == 1);
0065   std::cout << outputs[0].DebugString() << std::endl;
0066   CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
0067 
0068   // run again using the convenience helper
0069   outputs.clear();
0070   tensorflow::run(session, {{"input", input}, {"scale", scale}}, {"output"}, &outputs);
0071   CPPUNIT_ASSERT(outputs.size() == 1);
0072   std::cout << outputs[0].DebugString() << std::endl;
0073   CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
0074 
0075   // check for exception
0076   CPPUNIT_ASSERT_THROW(tensorflow::run(session, {{"foo", input}}, {"output"}, &outputs), cms::Exception);
0077 
0078   // cleanup
0079   CPPUNIT_ASSERT(tensorflow::closeSession(session));
0080   CPPUNIT_ASSERT(session == nullptr);
0081   delete graphDef;
0082 }