diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 37440ece..ff87cdf3 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -36,7 +36,7 @@ jobs: - name: Conda install dependencies shell: bash -l {0} run: | - conda create -n rootbench -y -c conda-forge root cmake pytest pytest-benchmark pytest-csv numpy numba + conda create -n rootbench -y -c conda-forge root cmake pytest pytest-benchmark pytest-csv numpy numba onnx - name: Configure and build shell: bash -l {0} diff --git a/cmake/modules/FindONNXRuntime.cmake b/cmake/modules/FindONNXRuntime.cmake new file mode 100644 index 00000000..8464a8b0 --- /dev/null +++ b/cmake/modules/FindONNXRuntime.cmake @@ -0,0 +1,47 @@ +# Copyright (C) 1995-2019, Rene Brun and Fons Rademakers. +# All rights reserved. +# +# For the licensing terms see $ROOTSYS/LICENSE. +# For the list of contributors see $ROOTSYS/README/CREDITS. + +# Find the ONNXRuntime includes and library. +# +# This module defines +# ONNXRuntime_INCLUDE_DIR, where to locate ONNXRuntime include file +# ONNXRuntime_LIBRARIES, the libraries to link against to use ONNXRuntime +# ONNXRuntime_FOUND. If false, you cannot build anything that requires ONNXRuntime. +# ONNXRuntime_LIBRARY, where to find the libONNXRuntime library. + +set(ONNXRuntime_FOUND 0) +if(ONNXRuntime_LIBRARY AND ONNXRuntime_INCLUDE_DIR) + set(ONNXRuntime_FIND_QUIETLY TRUE) +endif() + +find_path(ONNXRuntime_INCLUDE_DIR onnxruntime_cxx_api.h + $ENV{ONNXRuntime_DIR}/include + $ENV{ONNXRuntime} $ENV{ONNXRuntime}/include + /usr/local/include + /usr/include + DOC "Specify the directory containing ONNXRuntime.h" +) + +find_library(ONNXRuntime_LIBRARY NAMES onnxruntime PATHS + $ENV{ONNXRuntime_DIR}/lib + $ENV{ONNXRuntime} $ENV{ONNXRuntime}/lib $ENV{ONNXRuntime}/.libs + /usr/local/lib + /usr/lib + /opt/ONNXRuntime/lib + DOC "Specify the ONNXRuntime library here." +) + +if(ONNXRuntime_INCLUDE_DIR AND ONNXRuntime_LIBRARY) + set(ONNXRuntime_FOUND 1 ) + if(NOT ONNXRuntime_FIND_QUIETLY) + message(STATUS "Found ONNXRuntime includes at ${ONNXRuntime_INCLUDE_DIR}") + message(STATUS "Found ONNXRuntime library at ${ONNXRuntime_LIBRARY}") + endif() +endif() + +set(ONNXRuntime_LIBRARIES ${ONNXRuntime_LIBRARY}) + +mark_as_advanced(ONNXRuntime_FOUND ONNXRuntime_LIBRARY ONNXRuntime_INCLUDE_DIR) diff --git a/root/tmva/CMakeLists.txt b/root/tmva/CMakeLists.txt index 2f1c8c67..136ab7d3 100644 --- a/root/tmva/CMakeLists.txt +++ b/root/tmva/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(tmva) +add_subdirectory(sofie) diff --git a/root/tmva/sofie/CMakeLists.txt b/root/tmva/sofie/CMakeLists.txt new file mode 100644 index 00000000..3a890803 --- /dev/null +++ b/root/tmva/sofie/CMakeLists.txt @@ -0,0 +1,160 @@ +# TMVA SOFIE inference benchmarks. +# @author Federico Sossai (fsossai), Lorenzo Moneta + +# Check if SOFIE is available. ROOT versions up to 6.40 have a dedicated +# tmva-sofie build option that is advertised as a ROOT feature, while in +# later versions SOFIE is built unconditionally with TMVA and we probe for +# its libraries directly: in ROOT-builtin builds the library targets exist, +# and in standalone builds the libraries are found in the ROOT installation. +set(RB_HAVE_SOFIE FALSE) +if(ROOT_tmva-sofie_FOUND) + set(RB_HAVE_SOFIE TRUE) +elseif(TARGET ROOTTMVASofie AND TARGET ROOTTMVASofieParser) + set(RB_HAVE_SOFIE TRUE) +else() + find_library(RB_SOFIE_LIBRARY ROOTTMVASofie HINTS ${ROOT_LIBRARY_DIR}) + find_library(RB_SOFIE_PARSER_LIBRARY ROOTTMVASofieParser HINTS ${ROOT_LIBRARY_DIR}) + if(RB_SOFIE_LIBRARY AND RB_SOFIE_PARSER_LIBRARY) + set(RB_HAVE_SOFIE TRUE) + endif() +endif() +if(NOT (ROOT_tmva_FOUND AND RB_HAVE_SOFIE)) + message(STATUS "TMVA SOFIE not found: disabling the SOFIE benchmarks") + return() +endif() + +# The code generated by SOFIE uses BLAS for the matrix operations. +find_package(BLAS) +if(NOT BLAS_FOUND) + message(STATUS "BLAS not found: disabling the SOFIE benchmarks") + return() +endif() + +# The ONNX input models are generated at build time by make_input_models.py, +# which needs Python with the onnx and numpy packages. +find_package(Python3 COMPONENTS Interpreter) +if(Python3_FOUND) + execute_process(COMMAND ${Python3_EXECUTABLE} -c "import onnx, numpy" + RESULT_VARIABLE onnx_missing OUTPUT_QUIET ERROR_QUIET) +endif() +if(NOT Python3_FOUND OR onnx_missing) + message(STATUS "Python with the onnx package not found: disabling the SOFIE benchmarks") + return() +endif() + +# The models to benchmark. For each of them, an ONNX file is generated with +# make_input_models.py, from which emitFromONNX then generates the inference +# code compiled into the benchmarks. Everything happens at build time, so +# the benchmarks always exercise the SOFIE version of the ROOT build they +# run against. +set(sofie_models + Conv3d_d32_L4_B1 + ConvTModel_G4 + ConvTrans2dModel_B1 + Conv_d100_L14_B1 + Conv_d100_L14_B32 + Conv_d100_L1_B1 + Generator_B1 + Generator_B64 + Linear_16 + Linear_32 + Linear_64 + Linear_event + SimpleNN_Alice + higgs_model_dense) + +# Command line tool that generates the inference code for an ONNX model. +add_executable(emitFromONNX EmitFromONNX.cxx) +target_link_libraries(emitFromONNX Core ROOTTMVASofie ROOTTMVASofieParser) +set_target_properties(emitFromONNX PROPERTIES POSITION_INDEPENDENT_CODE TRUE) + +set(model_dir ${CMAKE_CURRENT_BINARY_DIR}/input_models) +set(model_generator ${CMAKE_CURRENT_SOURCE_DIR}/make_input_models.py) + +set(sofie_headers "") +foreach(name ${sofie_models}) + add_custom_command( + OUTPUT ${model_dir}/${name}.onnx + COMMAND ${Python3_EXECUTABLE} ${model_generator} --outdir ${model_dir} ${name} + DEPENDS ${model_generator} + COMMENT "Generating ONNX model ${name}") + add_custom_command( + OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/${name}.hxx ${CMAKE_CURRENT_BINARY_DIR}/${name}.dat + COMMAND emitFromONNX ${model_dir}/${name}.onnx + DEPENDS emitFromONNX ${model_dir}/${name}.onnx + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + COMMENT "Generating SOFIE inference code for ${name}") + list(APPEND sofie_headers ${CMAKE_CURRENT_BINARY_DIR}/${name}.hxx) +endforeach() +add_custom_target(SofieCompileModels DEPENDS ${sofie_headers}) + +# Benchmark of the inference code emitted by SOFIE +RB_ADD_GBENCHMARK(SOFIEInference + SOFIEInference.cxx + LABEL short + DEPENDS SofieCompileModels + LIBRARIES Core MathCore ROOTTMVASofie ${BLAS_LIBRARIES}) + +# Benchmark of RSofieReader, which parses the model and JITs the generated +# code at runtime +RB_ADD_GBENCHMARK(SOFIEInference_Reader + SOFIEInference_Reader.cxx + LABEL short + DEPENDS SofieCompileModels + LIBRARIES Core Cling MathCore ROOTTMVASofie ${BLAS_LIBRARIES}) + +# Benchmark of SOFIE inference inside an RDataFrame event loop +RB_ADD_GBENCHMARK(RDF_SOFIE_Inference + RDF_SOFIE_Inference.cxx + LABEL short + DEPENDS SofieCompileModels + LIBRARIES Core Hist Imt RIO Tree TreePlayer ROOTDataFrame ROOTVecOps ROOTTMVASofie ${BLAS_LIBRARIES}) + +# Compile the benchmarks with -O3 so that the generated inference code is +# auto-vectorized like in an optimized user build. More aggressive options +# (-march=native, -ffast-math) are deliberately not used: they would make the +# results machine-dependent and change the numerical behavior of operators +# that rely on infinities. +target_compile_options(SOFIEInference PRIVATE -O3) +target_compile_options(RDF_SOFIE_Inference PRIVATE -O3) + +# Optional comparison benchmark using ONNXRuntime on the same models. To help +# CMake find ONNXRuntime, configure with +# -DONNXRuntime_INCLUDE_DIRS=/include -DONNXRuntime_LIBRARIES=/lib +find_package(ONNXRuntime) +if(ONNXRuntime_FOUND) + message(STATUS "Found ONNXRuntime (library is ${ONNXRuntime_LIBRARY}, libraries ${ONNXRuntime_LIBRARIES})") + + # Generate one benchmark registration per model from + # ONNXRuntimeInference_Template.cxx.in + set(FUNC_NAME "BM_ONNXRuntime_Inference") + set(CAPTURE_STR "BENCHMARK_CAPTURE(${FUNC_NAME}, @1,\t@2)@3") + set(HEAD_COMMENT "Automatically configured by CMake") + set(ALL_CAPTURES "") + foreach(name ${sofie_models}) + string(REPLACE "@1" ${name} cap ${CAPTURE_STR}) + string(REPLACE "@2" "\"input_models/${name}.onnx\"" cap ${cap}) + list(APPEND ALL_CAPTURES ${cap}) + endforeach() + string(REPLACE ";" "\n" BENCHMARK_CAPTURES "${ALL_CAPTURES}") + string(REPLACE "@3" "->Unit(benchmark::kMillisecond);" BENCHMARK_CAPTURES "${BENCHMARK_CAPTURES}") + configure_file(ONNXRuntimeInference_Template.cxx.in ONNXRuntimeInference.cxx @ONLY) + + RB_ADD_GBENCHMARK(ONNXRuntimeInference + ONNXRuntimeInference.cxx + LABEL short + DEPENDS SofieCompileModels + LIBRARIES Core ${ONNXRuntime_LIBRARIES}) + target_link_directories(ONNXRuntimeInference PRIVATE ${ONNXRuntime_LIBRARIES}) + target_include_directories(ONNXRuntimeInference PRIVATE ${ONNXRuntime_INCLUDE_DIR}) + + RB_ADD_GBENCHMARK(RDF_ONNXRuntime_Inference + RDF_ONNXRuntime_Inference.cxx + LABEL short + DEPENDS SofieCompileModels + LIBRARIES Core Hist Imt MathCore RIO Tree TreePlayer ROOTDataFrame ROOTVecOps ${ONNXRuntime_LIBRARIES}) + target_link_directories(RDF_ONNXRuntime_Inference PRIVATE ${ONNXRuntime_LIBRARIES}) + target_include_directories(RDF_ONNXRuntime_Inference PRIVATE ${ONNXRuntime_INCLUDE_DIR}) +else() + message(STATUS "ONNXRuntime not found: disabling the ONNXRuntime benchmarks") +endif() diff --git a/root/tmva/sofie/EmitFromONNX.cxx b/root/tmva/sofie/EmitFromONNX.cxx new file mode 100644 index 00000000..e4edbe60 --- /dev/null +++ b/root/tmva/sofie/EmitFromONNX.cxx @@ -0,0 +1,29 @@ +// Author: Federico Sossai +// Last modified: 2021/07/30 +// Description: +// SOFIE command line compiler. +// This program is automatically run when the corresponding test target is built. +// Usage example: $./EmitFromONNX indir/mymodel.onnx outdir/myname.hxx + +#include + +#include "TMVA/RModel.hxx" +#include "TMVA/RModelParser_ONNX.hxx" + +using namespace TMVA::Experimental::SOFIE; + +int main(int argc, char *argv[]){ + if (argc < 2) { + std::cerr << "ERROR: missing input file\n"; + return -1; + } + + std::string outname= (argc > 2) ? argv[2] : ""; + RModelParser_ONNX parser; + std::cout << "Parsing file " << argv[1] << std::endl; + RModel model = parser.Parse(argv[1]); + model.Generate(Options::kDefault, 1); + model.PrintRequiredInputTensors(); + model.OutputGenerated(outname); + return 0; +} diff --git a/root/tmva/sofie/ONNXRuntimeInference_Template.cxx.in b/root/tmva/sofie/ONNXRuntimeInference_Template.cxx.in new file mode 100644 index 00000000..c09baf23 --- /dev/null +++ b/root/tmva/sofie/ONNXRuntimeInference_Template.cxx.in @@ -0,0 +1,197 @@ +// @HEAD_COMMENT@ +// Author: Federico Sossai (fsossai), 2021 + +#include +//#include +#include + +#include +#include +#include +#include +#include +#include +#include + +using namespace std; + +bool testOutput = true; + +static void @FUNC_NAME@(benchmark::State& state, string model_path) +{ + Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "benchmark"); + + Ort::SessionOptions session_options; + session_options.SetIntraOpNumThreads(1); + session_options.SetInterOpNumThreads(1); + session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED); + + //std::cout << "benchmarking model " << model_path << std::endl; + Ort::Session session(env, model_path.c_str(), session_options); + + int nin = session.GetInputCount(); + int nout = 1; + + vector input_node_names(nin); + vector output_node_names(nout); + vector inputStrings(nin); + vector outputStrings(nout); + + Ort::AllocatorWithDefaultOptions allocator; + for (int i = 0; i < nin; i++) { +#if ORT_API_VERSION > 12 + inputStrings[i] = session.GetInputNameAllocated(i, allocator).get(); +#else + inputStrings[i] = session.GetInputName(i, allocator); +#endif + input_node_names[i] = inputStrings[i].c_str(); + } + for (int i = 0; i < nout; i++) { +#if ORT_API_VERSION > 12 + outputStrings[i] = session.GetOutputNameAllocated(i, allocator).get(); +#else + outputStrings[i] = session.GetOutputName(i, allocator); +#endif + output_node_names[i] = outputStrings[i].c_str(); + } + // Getting the shapes + vector> input_node_dims(nin); + vector> output_node_dims(nout); + + for (int i = 0; i < nin; i++) + input_node_dims[i] = session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape(); + for (int i = 0; i < nout; i++) + output_node_dims[i] = session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape(); + + // for (int i = 0; i < nin; i++) { + // std::cout << "input " << input_node_names[i] << " shape : "; + // for (int j = 0; j < input_node_dims[i].size(); j++) + // std::cout << " " << input_node_dims[i][j]; + // std::cout << std::endl; + // } + // fix negative shapes + for (int i = 0; i < nin; i++) { + for (int j = 0; j < input_node_dims[i].size(); j++) { + if (input_node_dims[i][j] < 0) input_node_dims[i][j] = - input_node_dims[i][j]; + } + } + + + // Calculating the dimension of the input tensor + int nevts = 64; + int bsize = input_node_dims[0][0]; // assume this + //std::cout << "Using bsize = " << bsize << std::endl; + int nbatches = nevts / bsize; + + std::vector> inputData(nin); + std::vector inputSizes(nin); + + for (int i = 0; i < nin; i++) { + size_t input_tensor_size = accumulate(input_node_dims[i].begin(), input_node_dims[i].end(), 1, multiplies()); + inputSizes[i] = input_tensor_size; + auto &input_tensor_values = inputData[i]; + input_tensor_values.resize(input_tensor_size * nbatches); + // std::cout << "input tensor size " << input_tensor_size << " " << input_tensor_values.size() << std::endl; + + // Input tensor initialization + + if (testOutput) + fill_n(input_tensor_values.begin(), input_tensor_values.size(), float(i)+1.); + else { + static std::uniform_real_distribution distribution(-1, 1); + static std::default_random_engine generator; + std::generate(input_tensor_values.begin(), input_tensor_values.end(), []() { return distribution(generator); }); + } + } + + auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + // Ort::Value input_tensor = Ort::Value::CreateTensor(memory_info, + // input_tensor_values.data(), input_tensor_size, + // input_node_dims.data(), input_node_dims.size()); + + // Running the model + float *floatarr = nullptr; + + std::vector input_tensors; + + size_t osize = 1; + for (int d : output_node_dims[0]) { + if (d > 0) osize *= d; // first dim(batch size) can be -1 + } + std::vector yOut(osize); + + double totDuration = 0; + int ntimes = 0; + for (auto _ : state) { + auto t1 = std::chrono::high_resolution_clock::now(); + std::vector input_offset(nin); + for (int i = 0; i < nevts; i += bsize) { + // if (input_offset > input_tensor_values.size()) { + // std::cout << "Error in input size " << i << " " << nevts << " " << model_path << std::endl; + // throw std::runtime_error("Bad input size "); + // } + for (int k = 0; k < nin; k++) { + input_tensors.emplace_back(Ort::Value::CreateTensor(memory_info, inputData[k].data() + input_offset[k], + inputSizes[k], input_node_dims[k].data(), input_node_dims[k].size())); + } + auto output_tensors = session.Run(Ort::RunOptions{nullptr}, input_node_names.data(), input_tensors.data(), nin, + output_node_names.data(), nout); + floatarr = output_tensors.front().GetTensorMutableData(); + for (int k = 0; k < nin; k++) { + input_offset[k] += inputSizes[k]; + } + if (testOutput && i == 0) + std::copy(floatarr, floatarr + osize, yOut.begin()); + } + + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + totDuration += duration / 1.E3; // in milliseconds + ntimes++; + if (testOutput) { + std::string filename = model_path + ".ort.out"; + //std::cout << "writing file" << filename << std::endl; + ofstream f; + f.open(filename); + f << yOut.size(); + for (size_t i = 0; i < yOut.size(); i++) { + if ((i % 10) == 0) f << "\n"; // add endline every 10 + f << yOut[i] << " "; + } + f << std::endl; + f.close(); + } + } + //for (int i = 0; i < 10; i++) + // printf("%f\t", i, floatarr[i]); + state.counters["time/evt(ms)"] = totDuration / double(ntimes * nevts); + +} +@BENCHMARK_CAPTURES@ + +//BENCHMARK_MAIN(); + +// define main to pass some convenient command line parameters +int main(int argc, char **argv) { + + // Parse command line arguments + for (int i = 1; i < argc; i++) { + std::string arg = argv[i]; + if (arg == "-v") { + //std::cout << "---running in verbose mode" << std::endl; + //verbose = true; + } else if ((arg == "-d" || arg == "--dir") && argc > i+1) { + std::string pathDir = argv[i+1]; + std::filesystem::path path(pathDir); + std::filesystem::current_path(path); + i++; + } + } + + std::cout << "running benchmark from current directory " << std::filesystem::current_path() << std::endl; + + ::benchmark::Initialize(&argc, argv); + ::benchmark::RunSpecifiedBenchmarks(); + + return 0; +} \ No newline at end of file diff --git a/root/tmva/sofie/RDF_ONNXRuntime_Inference.cxx b/root/tmva/sofie/RDF_ONNXRuntime_Inference.cxx new file mode 100644 index 00000000..dce92574 --- /dev/null +++ b/root/tmva/sofie/RDF_ONNXRuntime_Inference.cxx @@ -0,0 +1,180 @@ +#include +#include "TROOT.h" +#include "TSystem.h" +#include "ROOT/RDataFrame.hxx" +#include "TMath.h" + +#include + + +#include +#include +#include + +#include + +// template +struct ONNXFunctor { + + // std::vector input; + // std::vector> sessions; + + std::shared_ptr session; + + //td::vector input_tensors; + + //Ort::Value * ort_input = nullptr; + + //float *inputArray = nullptr; + + std::vector input_node_names; + std::vector output_node_names; + std::vector input_node_str; + std::vector output_node_str; + + std::vector input_tensor_values; + + std::vector input_node_dims; + std::vector output_node_dims; + + Ort::Value inputTensor{nullptr}; + + float *inputArray = nullptr; + + ONNXFunctor(unsigned nslots) + { + + Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "benchmark"); + + std::string model_path = "higgs_model_dense.onnx"; + + Ort::SessionOptions session_options; + session_options.SetIntraOpNumThreads(1); + session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED); + + // std::cout << "benchmarking model " << model_path << std::endl; + session = std::make_shared(env, model_path.c_str(), session_options); + + + + Ort::AllocatorWithDefaultOptions allocator; + #if ORT_API_VERSION > 12 + input_node_str.push_back(session->GetInputNameAllocated(0, allocator).get()); + output_node_str.push_back(session->GetOutputNameAllocated(0, allocator).get()); + #else + input_node_str.push_back(session->GetInputName(0, allocator)); + output_node_str.push_back( session->GetOutputName(0, allocator)); + #endif + input_node_names.push_back(input_node_str.back().c_str()); + output_node_names.push_back(output_node_str.back().c_str()); + // Getting the shapes + + input_node_dims = session->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); + output_node_dims = session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); + + // Calculating the dimension of the input tensor + + + size_t input_tensor_size = std::accumulate(input_node_dims.begin(), input_node_dims.end(), 1, std::multiplies()); + //std::vector input_tensor_values(input_tensor_size ); + + input_tensor_values.resize(input_tensor_size); + + auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + + inputTensor = + Ort::Value::CreateTensor(memory_info, input_tensor_values.data(), input_tensor_values.size(), + input_node_dims.data(), input_node_dims.size()); + + inputArray = inputTensor.GetTensorMutableData(); + } + + double operator()(unsigned nslots, float x0, float x1, float x2, float x3, float x4, float x5, float x6) + { + + + int off = 0; + inputArray[off] = x0; + inputArray[off + 1] = x1; + inputArray[off + 2] = x2; + inputArray[off + 3] = x3; + inputArray[off + 4] = x4; + inputArray[off + 5] = x5; + inputArray[off + 6] = x6; + + + + auto output_tensors = session->Run(Ort::RunOptions{nullptr}, input_node_names.data(), &inputTensor, 1, output_node_names.data(), 1); + float * floatarr = output_tensors.front().GetTensorMutableData(); + return floatarr[0]; + } + + // need copy ctor for ONNXruntime + // because I cannot copy Ort::Value + ONNXFunctor(const ONNXFunctor & rhs) { + session = rhs.session; + input_node_names = rhs.input_node_names; + output_node_names = rhs.output_node_names; + + input_tensor_values = rhs.input_tensor_values; + + input_node_dims = rhs.input_node_dims; + output_node_dims = rhs.output_node_dims; + + auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + inputTensor = Ort::Value::CreateTensor(memory_info, input_tensor_values.data(), input_tensor_values.size(), + input_node_dims.data(), input_node_dims.size()); + inputArray = inputTensor.GetTensorMutableData(); + } +}; + +void BM_RDF_ONNX_Inference(benchmark::State &state) +{ + + int nslot = 1; + if (nslot > 1) + ROOT::EnableImplicitMT(nslot); + + auto fileName = "Higgs_data_full.root"; + // file is available at "https://cernbox.cern.ch/index.php/s/YuSHwTXBa0UBEhD/download"; + // do curl https://cernbox.cern.ch/index.php/s/XaPBtaGrnN38wU0 -o Higgs_data_full.root + // https://cernbox.cern.ch/s/vLOqclhWirZEWpj + std::string directLink = "https://cernbox.cern.ch/remote.php/dav/public-files/vLOqclhWirZEWpj/Higgs_data_full.root"; + if (gSystem->AccessPathName(fileName)) { + std::string cmd = "curl " + directLink + " -o "; + cmd += fileName; + gSystem->Exec(cmd.c_str()); + } + auto treeName = "test_tree"; + ROOT::RDataFrame df(treeName, fileName); + + ONNXFunctor functor(nslot); + + std::vector durations; + double ntot = 0; + + for (auto _ : state) { + + auto h1 = df.DefineSlot("DNN_Value", functor, {"m_jj", "m_jjj", "m_lv", "m_jlv", "m_bb", "m_wbb", "m_wwbb"}) + .Histo1D("DNN_Value"); + + auto t1 = std::chrono::high_resolution_clock::now(); + + auto n = h1->GetEntries(); + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + + durations.push_back(duration/1.E6); + ntot += n; + // std::cout << " Processed " << n << " entries " + // << " time = " << duration / 1.E6 << " (sec) time/event = " << duration / double(n) << " musec" + // << std::endl; + } + double avgDuration = TMath::Mean(durations.begin(), durations.end()); + state.counters["avg-time(s)"] = avgDuration; + state.counters["time/evt(s)"] = avgDuration * double(durations.size()) / ntot; +} + +BENCHMARK(BM_RDF_ONNX_Inference)->Unit(benchmark::kMillisecond); + +BENCHMARK_MAIN(); diff --git a/root/tmva/sofie/RDF_SOFIE_Inference.cxx b/root/tmva/sofie/RDF_SOFIE_Inference.cxx new file mode 100644 index 00000000..7a910595 --- /dev/null +++ b/root/tmva/sofie/RDF_SOFIE_Inference.cxx @@ -0,0 +1,124 @@ +#include "higgs_model_dense.hxx" +#include +#include +#include "TROOT.h" +#include "TSystem.h" +#include "ROOT/RDataFrame.hxx" + +#include + +// Functor to wrap SOFIE session to RDF functor signature + +template +class SofieFunctorHelper; + +template +class SofieFunctorHelper, S, T> { + /// this is the magic to defined the operator () with N fixed parameter arguments + template + using AlwaysT = T; + + std::vector> fInput; + std::vector> fSessions; + +public: + + SofieFunctorHelper(int nslots) : + fInput(nslots) + { + for (int i = 0; i < nslots; i++) { + fSessions.emplace_back(std::make_shared()); + } + } + + double operator()(unsigned slot, AlwaysT... args) { + fInput[slot] = {args...}; + auto y = fSessions[slot]->infer(fInput[slot].data()); + return y[0]; + } + + +}; + +template +auto SofieFunctor(int nslot) -> SofieFunctorHelper, F, float> +{ + return SofieFunctorHelper, F, float>(nslot); +} + + +int NEVTS = -1; +void BM_RDF_SOFIE_Inference(benchmark::State &state) +{ + int nslot = state.range(0); + + if (nslot > 1) + ROOT::EnableImplicitMT(nslot); + auto fileName = "Higgs_data_full.root"; + //file is available at "https://cernbox.cern.ch/index.php/s/YuSHwTXBa0UBEhD/download"; + // do curl https://cernbox.cern.ch/index.php/s/XaPBtaGrnN38wU0 -o Higgs_data_full.root + std::string directLink = "https://cernbox.cern.ch/remote.php/dav/public-files/vLOqclhWirZEWpj/Higgs_data_full.root"; + if (gSystem->AccessPathName(fileName)) { + std::string cmd = "curl " + directLink + " -o "; + cmd += fileName; + gSystem->Exec(cmd.c_str()); + } + auto treeName = "test_tree"; + ROOT::RDataFrame df(treeName, fileName); + + + //auto functor = SofieFunctor(nslot); + // auto rdf_functor = [&](int slot, float x1, float x2, float x3, float x4, float x5, float x6, float x7){ + // return functor(slot, x1,x2,x3,x4,x5,x6,x7); + // }; + SofieFunctorHelper, TMVA_SOFIE_higgs_model_dense::Session, float> functor(nslot); + + // test + auto y = functor(0,1.,2.,3.,4.,5.,6.,7.); + std::cout << y << std::endl; + + std::vector durations; + + double ntot = 0; + + for (auto _ : state) { + + auto h1 = df.DefineSlot("DNN_Value", SofieFunctor<7,TMVA_SOFIE_higgs_model_dense::Session>(nslot), {"m_jj", "m_jjj", "m_lv", "m_jlv", "m_bb", "m_wbb", "m_wwbb"}) + .Histo1D("DNN_Value"); + // auto h1 = df.Define("DNN_Value", "functor(m_jj, m_jjj, m_lv,m_jlv, m_bb, m_wbb, m_wwbb)") + // .Histo1D("DNN_Value"); + + auto t1 = std::chrono::high_resolution_clock::now(); + + auto n = h1->GetEntries(); + //int n = 100; + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + + durations.push_back(duration / 1.E6); + NEVTS = n; + ntot += n; + // std::cout << " Processed " << n << " entries " + // << " time = " << duration / 1.E6 << " (sec) time/event = " << duration / double(n) << " musec" + // << std::endl; + } + + double avgDuration = TMath::Mean(durations.begin(), durations.end()); + state.counters["avg-time(s)"] = avgDuration; + if (durations.size() > 1) + state.counters["+/-"] = TMath::StdDev(durations.begin(), durations.end()) / sqrt(durations.size() - 1); + state.counters["time/evt(s)"] = avgDuration *double(durations.size()) / ntot; + // h1->DrawClone(); +} + +BENCHMARK(BM_RDF_SOFIE_Inference) + ->Unit(benchmark::kMillisecond) + // ->ComputeStatistics("Time/evt", + // [](const std::vector &v) -> double { + // return std::accumulate(v.begin(), v.end(), 0.) / (v.size() * NEVTS);} + // , benchmark::StatisticUnit::kTime) + ->Arg(1) + ->Arg(2) + ->Arg(4); + +BENCHMARK_MAIN(); \ No newline at end of file diff --git a/root/tmva/sofie/SOFIEInference.cxx b/root/tmva/sofie/SOFIEInference.cxx new file mode 100644 index 00000000..010fe8fb --- /dev/null +++ b/root/tmva/sofie/SOFIEInference.cxx @@ -0,0 +1,227 @@ +// Author: Federico Sossai (fsossai), 2021 + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "Linear_event.hxx" +#include "Linear_16.hxx" +#include "Linear_32.hxx" +#include "Linear_64.hxx" +#include "Generator_B1.hxx" +#include "Generator_B64.hxx" +#include "Conv_d100_L1_B1.hxx" +#include "Conv_d100_L14_B1.hxx" +#include "Conv_d100_L14_B32.hxx" +#include "Conv3d_d32_L4_B1.hxx" +#include "higgs_model_dense.hxx" +#include "ConvTrans2dModel_B1.hxx" +#include "ConvTModel_G4.hxx" +#include "SimpleNN_Alice.hxx" + +// The following models from PR #239 are not benchmarked because they are +// not supported by the current version of SOFIE: +// * RNN_d10_L20_h8_B1, GRU_d10_L20_h8_B1, LSTM_d10_L20_h8_B1, DDB_B1: +// the generated code does not compile +// * Conv2DTranspose_Relu_Sigmoid: dynamic tensor error when parsing +// * resnet18v1: "intermediate tensor already exists" error when parsing + +#include "TMath.h" + + +using namespace std; +bool verbose = false; +// use fixed instead of random input data, so that inference outputs are reproducible +bool testOutput = true; +// write the first inference output of each model to a file (for validating +// the results across ROOT versions), enabled with the -o command line option +bool writeOutput = false; + + +template +void BM_SOFIE_Inference(benchmark::State &state) +{ + size_t inputSize = state.range(0); // input size (without batch size) + size_t bsize = (state.range(1) > 0) ? state.range(1) : 1; + size_t nevts = 64; + size_t nrep = nevts / bsize; + + vector input(inputSize*nevts); + + if (testOutput) { + input = std::vector(input.size(),1.); + } + else { + static std::uniform_real_distribution distribution(-1, 1); + static std::default_random_engine generator; + std::generate(input.begin(), input.end(), []() { return distribution(generator); }); + } + float *input_ptr = input.data(); + // construct session (no need to pass filename, use default value) + S s; + + double totDuration = 0; + int ntimes = 0; + std::vector yOut; + bool first = true; + bool doWrite = writeOutput; + for (auto _ : state) { + auto t1 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < nevts; i += bsize) { + auto y = s.infer(input.data()+ inputSize*i); + if (first) { + //std::cout << std::string(typeid(s).name()) << " : " << y[0] << " " << y[1] << std::endl; + yOut = y; + first = false; + } + } + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + totDuration += duration / 1.E3; // in milliseconds + ntimes++; + if (doWrite) { + // write output for test + //std::cout << "write output " << std::endl; + std::ofstream f; + std::string filename = std::string(typeid(s).name()) + ".out"; + f.open(filename); + f << yOut.size(); + for (size_t i = 0; i < yOut.size(); i++) { + if ((i % 10) == 0) f << "\n"; // add endline every 10 + f << yOut[i] << " "; + } + f << std::endl; + f.close(); + doWrite = false; + } + } + + state.counters["time/evt(ms)"] = totDuration / double(ntimes * nevts); + // input[0] = -999; + // s.inf + // std::cout << "number of times " << s.itime << std::endl; + // int n = s.itime - 1; + // for (size_t i = 0; i < 5; ++i) { + // double mean = TMath::Mean(n, resTimes[i].data()); + // double rms = TMath::RMS(n, resfTimes[i].data()); + // std::cout << "elapsed time for " << i << " : " << mean << " +/- " << rms / sqrt(n) << std::endl; + // } + //if (verbose) std::cout << "output : " << output.size() << " : " << output.front() << " ......" << output.back() << std::endl; +} + +// inference for model with 3 inputs +template +void BM_SOFIE_Inference_3(benchmark::State &state) +{ + size_t bsize = state.range(0); // batch size + size_t inputSize1 = state.range(1); // input 1 size + size_t inputSize2 = state.range(2); // input 2 size + size_t inputSize3 = state.range(3); + + size_t nevts = 64; + size_t nrep = nevts / bsize; + + size_t eventSize = inputSize1 + inputSize2+inputSize3; + + vector input1(inputSize1*nevts); + vector input2(inputSize2*nevts); + vector input3(inputSize3*nevts); + + if (!testOutput) { + static std::uniform_real_distribution distribution(-1, 1); + static std::default_random_engine generator; + std::generate(input1.begin(), input1.end(), []() { return distribution(generator); }); + std::generate(input2.begin(), input2.end(), []() { return distribution(generator); }); + std::generate(input3.begin(), input3.end(), []() { return distribution(generator); }); + } + else { + // generate fixed data + input1 = vector(input1.size(),1.); + input2 = vector(input2.size(),2.); + input3 = vector(input3.size(),3.); + } + + // create session with default filename + S s{}; + + //std::cout << "init done - do benchmark \n"; + + double totDuration = 0; + int ntimes = 0; + for (auto _ : state) { + auto t1 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < nevts; i += bsize) { + float * p1 = input1.data()+ inputSize1*i; + float * p2 = input2.data()+ inputSize2*i; + float * p3 = input3.data()+ inputSize3*i; + auto y = s.infer(p1,p2,p3); + } + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + totDuration += duration / 1.E3; // in milliseconds + ntimes++; + } + + state.counters["time/evt(ms)"] = totDuration / double(ntimes * nevts); +} + +// Conv Transpose +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_ConvTModel_G4::Session)->Name("ConvTModel_G4")->Args({15,1})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_ConvTrans2dModel_B1::Session)->Name("ConvTrans2dModel_B1")->Args({4*4*4,1})->Unit(benchmark::kMillisecond); + +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_SimpleNN_Alice::Session)->Name("SimpleNN_Alice")->Args({16,1})->Unit(benchmark::kMillisecond); + +//Gemm benchmarks +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Linear_16::Session)->Name("Linear_16")->Args({100, 16})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Linear_32::Session)->Name("Linear_32")->Args({100, 32})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Linear_64::Session)->Name("Linear_64")->Args({100, 64})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Linear_event::Session)->Name("Linear_event")->Args({100, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Generator_B1::Session)->Name("Generator_B1")->Args({14, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Generator_B64::Session)->Name("Generator_B64")->Args({14, 64})->Unit(benchmark::kMillisecond); + +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_higgs_model_dense::Session)->Name("higgs_model_dense")->Args({7, 1})->Unit(benchmark::kMillisecond); + +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Conv_d100_L14_B1::Session)->Name( "Conv_d100_L14_B1")->Args({100*100, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Conv_d100_L14_B32::Session)->Name("Conv_d100_L14_B32")->Args({100*100, 32})->Unit(benchmark::kMillisecond); +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Conv_d100_L1_B1::Session)->Name( "Conv_d100_L1_B1")->Args({100*100, 1})->Unit(benchmark::kMillisecond); + +BENCHMARK_TEMPLATE(BM_SOFIE_Inference, TMVA_SOFIE_Conv3d_d32_L4_B1::Session)->Name( "Conv3d_d32_L4_B1")->Args({32*32*32, 1})->Unit(benchmark::kMillisecond); + +// default main +//BENCHMARK_MAIN(); + +// define main to pass some convenient command line parameters +int main(int argc, char **argv) { + + // Parse command line arguments + for (Int_t i = 1; i < argc; i++) { + std::string arg = argv[i]; + if (arg == "-v") { + std::cout << "---running in verbose mode" << std::endl; + verbose = true; + } else if (arg == "-o") { + std::cout << "---writing inference outputs to files" << std::endl; + writeOutput = true; + } else if ((arg == "-d" || arg == "--dir") && argc > i+1) { + std::string pathDir = argv[i+1]; + std::filesystem::path path(pathDir); + std::filesystem::current_path(path); + i++; + } + } + + std::cout << "running benchmark from current directory " << std::filesystem::current_path() << std::endl; + + ::benchmark::Initialize(&argc, argv); + ::benchmark::RunSpecifiedBenchmarks(); + + return 0; +} diff --git a/root/tmva/sofie/SOFIEInference_Reader.cxx b/root/tmva/sofie/SOFIEInference_Reader.cxx new file mode 100644 index 00000000..a0619a4d --- /dev/null +++ b/root/tmva/sofie/SOFIEInference_Reader.cxx @@ -0,0 +1,183 @@ +// Author: Federico Sossai (fsossai), 2021 + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "TMVA/RSofieReader.hxx" + +#include "TMath.h" + + +using namespace std; +bool verbose = false; +// use fixed instead of random input data, so that inference outputs are reproducible +bool testOutput = true; +// write the first inference output of each model to a file (for validating +// the results across ROOT versions) +bool writeOutput = false; + + +void BM_SOFIE_Inference(benchmark::State &state, std::string model_file) +{ + std::string model_path = "input_models/" + model_file; + size_t inputSize = state.range(0); // input size (without batch size) + size_t bsize = (state.range(1) > 0) ? state.range(1) : 1; + size_t nevts = 64; + size_t nrep = nevts / bsize; + + vector input(inputSize*nevts); + + if (testOutput) { + input = std::vector(input.size(),1.); + } + else { + static std::uniform_real_distribution distribution(-1, 1); + static std::default_random_engine generator; + std::generate(input.begin(), input.end(), []() { return distribution(generator); }); + } + float *input_ptr = input.data(); + + + // parse the model + TMVA::Experimental::RSofieReader r(model_path); + + double totDuration = 0; + int ntimes = 0; + std::vector yOut; + bool first = true; + bool doWrite = writeOutput; + for (auto _ : state) { + auto t1 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < nevts; i += bsize) { + std::vector x(input.begin()+inputSize*i, input.begin()+inputSize*(i+1)); + auto y = r.Compute(x); + if (first) { + //std::cout << std::string(typeid(s).name()) << " : " << y[0] << " " << y[1] << std::endl; + yOut = y; + first = false; + } + } + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + totDuration += duration / 1.E3; // in milliseconds + ntimes++; + if (doWrite) { + // write output for test + //std::cout << "write output " << std::endl; + std::ofstream f; + std::string filename = std::string(model_file) + ".out"; + f.open(filename); + f << yOut.size(); + for (size_t i = 0; i < yOut.size(); i++) { + if ((i % 10) == 0) f << "\n"; // add endline every 10 + f << yOut[i] << " "; + } + f << std::endl; + f.close(); + doWrite = false; + } + } + + state.counters["time/evt(ms)"] = totDuration / double(ntimes * nevts); + // input[0] = -999; + // s.inf + // std::cout << "number of times " << s.itime << std::endl; + // int n = s.itime - 1; + // for (size_t i = 0; i < 5; ++i) { + // double mean = TMath::Mean(n, resTimes[i].data()); + // double rms = TMath::RMS(n, resfTimes[i].data()); + // std::cout << "elapsed time for " << i << " : " << mean << " +/- " << rms / sqrt(n) << std::endl; + // } + //if (verbose) std::cout << "output : " << output.size() << " : " << output.front() << " ......" << output.back() << std::endl; +} +#if 0 +// inference for model with 3 inputs +template +void BM_SOFIE_Inference_3(benchmark::State &state) +{ + size_t bsize = state.range(0); // batch size + size_t inputSize1 = state.range(1); // input 1 size + size_t inputSize2 = state.range(2); // input 2 size + size_t inputSize3 = state.range(3); + + size_t nevts = 64; + size_t nrep = nevts / bsize; + + size_t eventSize = inputSize1 + inputSize2+inputSize3; + + vector input1(inputSize1*nevts); + vector input2(inputSize2*nevts); + vector input3(inputSize3*nevts); + + if (!testOutput) { + static std::uniform_real_distribution distribution(-1, 1); + static std::default_random_engine generator; + std::generate(input1.begin(), input1.end(), []() { return distribution(generator); }); + std::generate(input2.begin(), input2.end(), []() { return distribution(generator); }); + std::generate(input3.begin(), input3.end(), []() { return distribution(generator); }); + } + else { + // generate fixed data + input1 = vector(input1.size(),1.); + input2 = vector(input2.size(),2.); + input3 = vector(input3.size(),3.); + } + + S s(""); + + //std::cout << "init done - do benchmark \n"; + + double totDuration = 0; + int ntimes = 0; + for (auto _ : state) { + auto t1 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < nevts; i += bsize) { + float * p1 = input1.data()+ inputSize1*i; + float * p2 = input2.data()+ inputSize2*i; + float * p3 = input3.data()+ inputSize3*i; + auto y = s.infer(p1,p2,p3); + } + auto t2 = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(t2 - t1).count(); + totDuration += duration / 1.E3; // in milliseconds + ntimes++; + } + + state.counters["time/evt(ms)"] = totDuration / double(ntimes * nevts); +} +#endif + +// Some models from PR #239 are not benchmarked because they are not +// supported by the current version of SOFIE: +// * RNN_d10_L20_h8_B1, GRU_d10_L20_h8_B1, LSTM_d10_L20_h8_B1, DDB_B1: +// the generated code does not compile +// * Conv2DTranspose_Relu_Sigmoid: dynamic tensor error when parsing +// * resnet18v1: "intermediate tensor already exists" error when parsing + +BENCHMARK_CAPTURE(BM_SOFIE_Inference,higgs_model_dense,"higgs_model_dense.onnx")->Args({7, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, ConvTrans2dModel_B1,"ConvTrans2dModel_B1.onnx")->Args({4*4*4,1})->Unit(benchmark::kMillisecond); + +BENCHMARK_CAPTURE(BM_SOFIE_Inference, SimpleNN_Alice,"SimpleNN_Alice.onnx")->Args({16,1})->Unit(benchmark::kMillisecond); + +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Linear_16,"Linear_16.onnx")->Args({100, 16})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Linear_32,"Linear_32.onnx")->Args({100, 32})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Linear_64,"Linear_64.onnx")->Args({100, 64})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Linear_event,"Linear_event.onnx")->Args({100, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Generator_B1,"Generator_B1.onnx")->Args({14, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Generator_B64,"Generator_B64.onnx")->Args({14, 64})->Unit(benchmark::kMillisecond); + +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Conv_d100_L14_B1,"Conv_d100_L14_B1.onnx")->Args({100*100, 1})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Conv_d100_L14_B32,"Conv_d100_L14_B32.onnx")->Args({100*100, 32})->Unit(benchmark::kMillisecond); +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Conv_d100_L1_B1,"Conv_d100_L1_B1.onnx")->Args({100*100, 1})->Unit(benchmark::kMillisecond); + +BENCHMARK_CAPTURE(BM_SOFIE_Inference, Conv3d_d32_L4_B1,"Conv3d_d32_L4_B1.onnx")->Args({32*32*32, 1})->Unit(benchmark::kMillisecond); + +BENCHMARK_MAIN(); diff --git a/root/tmva/sofie/make_input_models.py b/root/tmva/sofie/make_input_models.py new file mode 100644 index 00000000..1fa9aad9 --- /dev/null +++ b/root/tmva/sofie/make_input_models.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +"""Generate the ONNX input models for the SOFIE benchmarks. + +The models are built directly with the onnx helper API, using seeded random +weights: only the network architecture matters for benchmarking the +inference speed, so nothing needs to be trained and no binary model files +need to be stored in the repository or downloaded. + +The architectures reproduce the models that were originally benchmarked in +PR #239 (and presented at ACAT 2021): parametrized dense and convolutional +networks exported from PyTorch generators (Linear_*, Conv_*, Conv3d_*, +ConvTrans2dModel_*, *_d10_L20_h8_B1), the higgs_model_dense classifier from +the TMVA tutorials, a small ALICE network (SimpleNN_Alice), and fast +simulation models (ConvTModel_G4 and the Generator GAN). + +Usage: + make_input_models.py [--outdir DIR] [model ...] + +Without model arguments, all models are generated. --list prints the +available model names. +""" + +import argparse +import os +import sys + +try: + import numpy as np + import onnx + from onnx import TensorProto, helper, numpy_helper +except ImportError as e: + print(f"ERROR: missing Python package: {e.name}", file=sys.stderr) + sys.exit(1) + + +def _tensor(rng, name, *dims): + """Random float32 weight initializer.""" + return numpy_helper.from_array( + rng.standard_normal(dims).astype(np.float32) * 0.1, name=name) + + +def _pos_tensor(rng, name, *dims): + """Random strictly positive float32 initializer (e.g. batchnorm variance).""" + return numpy_helper.from_array( + rng.uniform(0.5, 1.5, dims).astype(np.float32), name=name) + + +def _model(name, nodes, inputs, outputs, initializers, opset=9): + graph = helper.make_graph(nodes, name, inputs, outputs, initializers) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", opset)]) + onnx.checker.check_model(model) + return model + + +def _finfo(name, shape): + return helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + + +def make_linear(batch_size): + """Dense network with ten Gemm+Relu layers: 100 -> 8x50 -> 10.""" + rng = np.random.default_rng(16) + widths = [100] + 9 * [50] + [10] + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out) in enumerate(zip(widths[:-1], widths[1:])): + inits += [_tensor(rng, f"w{i}", n_out, n_in), _tensor(rng, f"b{i}", n_out)] + last = i == len(widths) - 2 + out = "output" if last else f"gemm{i}" + nodes.append(helper.make_node("Gemm", [x, f"w{i}", f"b{i}"], [out], transB=1)) + if not last: + nodes.append(helper.make_node("Relu", [out], [f"relu{i}"])) + x = f"relu{i}" + return _model("Linear", nodes, + [_finfo("input", [batch_size, 100])], + [_finfo("output", [batch_size, 10])], inits) + + +def make_generator(batch_size): + """Dense GAN generator: 14 -> 14 -> 20 -> 50 -> 100 -> 40500, with + batch normalization after each hidden layer and a Sigmoid output.""" + rng = np.random.default_rng(17) + widths = [14, 14, 20, 50, 100, 40500] + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out) in enumerate(zip(widths[:-1], widths[1:])): + # keras2onnx-style Gemm with the weight matrix stored as (in, out) + inits += [_tensor(rng, f"w{i}", n_in, n_out), _tensor(rng, f"b{i}", n_out)] + nodes.append(helper.make_node("Gemm", [x, f"w{i}", f"b{i}"], [f"gemm{i}"])) + x = f"gemm{i}" + if i < len(widths) - 2: + nodes.append(helper.make_node("Relu", [x], [f"relu{i}"])) + inits += [_tensor(rng, f"scale{i}", n_out), _tensor(rng, f"beta{i}", n_out), + _tensor(rng, f"mean{i}", n_out), _pos_tensor(rng, f"var{i}", n_out)] + nodes.append(helper.make_node( + "BatchNormalization", + [f"relu{i}", f"scale{i}", f"beta{i}", f"mean{i}", f"var{i}"], + [f"bn{i}"], epsilon=1e-6)) + x = f"bn{i}" + nodes.append(helper.make_node("Sigmoid", [x], ["output"])) + return _model("Generator", nodes, + [_finfo("input", [batch_size, 14])], + [_finfo("output", [batch_size, 40500])], inits) + + +def make_conv2d(nlayers, batch_size, pads): + """Chain of 5x5 Conv+Relu layers on a 100x100 image, with the channel + count doubling up to 128 in the middle of the chain and halving back to + one at the end (for nlayers=14), or a single 1->2 channel layer.""" + rng = np.random.default_rng(18) + if nlayers == 1: + channels = [1, 2] + else: + channels = [1, 2, 4, 8, 16, 32, 64, 128, 64, 32, 16, 8, 4, 2, 1] + assert len(channels) == nlayers + 1 + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out) in enumerate(zip(channels[:-1], channels[1:])): + inits += [_tensor(rng, f"w{i}", n_out, n_in, 5, 5), _tensor(rng, f"b{i}", n_out)] + nodes.append(helper.make_node( + "Conv", [x, f"w{i}", f"b{i}"], [f"conv{i}"], + kernel_shape=[5, 5], pads=4 * [pads], strides=[1, 1])) + out = "output" if i == nlayers - 1 else f"relu{i}" + nodes.append(helper.make_node("Relu", [f"conv{i}"], [out])) + x = out + d_out = 100 if pads == 2 else 100 - nlayers * 4 + return _model("Conv2d", nodes, + [_finfo("input", [batch_size, 1, 100, 100])], + [_finfo("output", [batch_size, channels[-1], d_out, d_out])], inits) + + +def make_conv3d(): + """3d convolutional network on a 32x32x32 volume: four 5x5x5 Conv+Relu + layers, a strided 6x6x6 pooling convolution, and a dense layer.""" + rng = np.random.default_rng(19) + channels = [1, 32, 8, 8, 8] + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out) in enumerate(zip(channels[:-1], channels[1:])): + inits += [_tensor(rng, f"w{i}", n_out, n_in, 5, 5, 5), _tensor(rng, f"b{i}", n_out)] + nodes.append(helper.make_node( + "Conv", [x, f"w{i}", f"b{i}"], [f"conv{i}"], + kernel_shape=[5, 5, 5], pads=6 * [1], strides=[1, 1, 1])) + nodes.append(helper.make_node("Relu", [f"conv{i}"], [f"relu{i}"])) + x = f"relu{i}" + inits += [_tensor(rng, "wpool", 4, 8, 6, 6, 6), _tensor(rng, "bpool", 4)] + nodes.append(helper.make_node( + "Conv", [x, "wpool", "bpool"], ["convpool"], + kernel_shape=[6, 6, 6], pads=6 * [0], strides=[6, 6, 6])) + nodes.append(helper.make_node("Relu", ["convpool"], ["relupool"])) + nodes.append(helper.make_node("Flatten", ["relupool"], ["flat"], axis=1)) + inits += [_tensor(rng, "wfc", 8, 256), _tensor(rng, "bfc", 8)] + nodes.append(helper.make_node("Gemm", ["flat", "wfc", "bfc"], ["output"], transB=1)) + return _model("Conv3d", nodes, + [_finfo("input", [1, 1, 32, 32, 32])], + [_finfo("output", [1, 8])], inits) + + +def make_convtrans2d(): + """Small chain of four ConvTranspose layers with Relu in between.""" + rng = np.random.default_rng(20) + # (in channels, out channels, kernel size, pads, strides) + layers = [(1, 4, 2, 1, 1), (4, 8, 3, 1, 1), (8, 4, 3, 1, 1), (4, 1, 2, 0, 2)] + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out, k, p, s) in enumerate(layers): + inits += [_tensor(rng, f"w{i}", n_in, n_out, k, k), _tensor(rng, f"b{i}", n_out)] + last = i == len(layers) - 1 + out = "output" if last else f"conv{i}" + nodes.append(helper.make_node( + "ConvTranspose", [x, f"w{i}", f"b{i}"], [out], + kernel_shape=[k, k], pads=4 * [p], strides=[s, s])) + if not last: + nodes.append(helper.make_node("Relu", [out], [f"relu{i}"])) + x = f"relu{i}" + return _model("ConvTrans2d", nodes, + [_finfo("input", [1, 1, 4, 4])], + [_finfo("output", [1, 1, 6, 6])], inits) + + +def make_convt_g4(): + """Fast simulation model: a dense layer inflates 15 inputs to a 3x11 + image with 180 channels, which three ConvTranspose layers upscale to + 18 channels of 18x50 (in channels-last output layout).""" + rng = np.random.default_rng(21) + nodes = [helper.make_node("Gemm", ["input", "wfc", "bfc"], ["fc"], transB=1), + helper.make_node("Relu", ["fc"], ["fcrelu"]), + helper.make_node("Reshape", ["fcrelu", "shape"], ["reshaped"]), + helper.make_node("Transpose", ["reshaped"], ["nchw"], perm=[0, 3, 1, 2])] + inits = [_tensor(rng, "wfc", 5940, 15), _tensor(rng, "bfc", 5940), + numpy_helper.from_array(np.array([1, 3, 11, 180], dtype=np.int64), name="shape")] + # (in channels, out channels, kernel size, strides) + layers = [(180, 180, 3, 2), (180, 90, 3, 2), (90, 45, 4, 1)] + x = "nchw" + for i, (n_in, n_out, k, s) in enumerate(layers): + inits += [_tensor(rng, f"w{i}", n_in, n_out, k, k), _tensor(rng, f"b{i}", n_out)] + nodes.append(helper.make_node( + "ConvTranspose", [x, f"w{i}", f"b{i}"], [f"conv{i}"], + kernel_shape=[k, k], pads=[0, 0, 0, 0], strides=[s, s])) + x = f"conv{i}" + if i < len(layers) - 1: + nodes.append(helper.make_node("Relu", [x], [f"relu{i}"])) + x = f"relu{i}" + nodes.append(helper.make_node("Sigmoid", [x], ["sigmoid"])) + nodes.append(helper.make_node("Transpose", ["sigmoid"], ["output"], perm=[0, 3, 1, 2])) + return _model("ConvTModel_G4", nodes, + [_finfo("input", [1, 15])], + [_finfo("output", [1, 50, 45, 18])], inits) + + +def make_simplenn_alice(): + """Small dense network with LeakyRelu activations and no batch dimension.""" + rng = np.random.default_rng(22) + widths = [16, 100, 50, 1] + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out) in enumerate(zip(widths[:-1], widths[1:])): + inits += [_tensor(rng, f"w{i}", n_in, n_out), _tensor(rng, f"b{i}", n_out)] + last = i == len(widths) - 2 + out = "output" if last else f"add{i}" + nodes.append(helper.make_node("MatMul", [x, f"w{i}"], [f"matmul{i}"])) + nodes.append(helper.make_node("Add", [f"b{i}", f"matmul{i}"], [out])) + if not last: + nodes.append(helper.make_node("LeakyRelu", [out], [f"lrelu{i}"], alpha=0.01)) + x = f"lrelu{i}" + return _model("SimpleNN_Alice", nodes, + [_finfo("input", [16])], [_finfo("output", [1])], inits) + + +def make_higgs_model_dense(): + """The dense Higgs classifier from the TMVA tutorials: 7 -> 5x100 -> 2 + with Relu activations and a Sigmoid output.""" + rng = np.random.default_rng(23) + widths = [7] + 5 * [100] + [2] + nodes, inits = [], [] + x = "input" + for i, (n_in, n_out) in enumerate(zip(widths[:-1], widths[1:])): + # keras2onnx-style Gemm with the weight matrix stored as (in, out) + inits += [_tensor(rng, f"w{i}", n_in, n_out), _tensor(rng, f"b{i}", n_out)] + nodes.append(helper.make_node("Gemm", [x, f"w{i}", f"b{i}"], [f"gemm{i}"])) + x = f"gemm{i}" + if i < len(widths) - 2: + nodes.append(helper.make_node("Relu", [x], [f"relu{i}"])) + x = f"relu{i}" + nodes.append(helper.make_node("Sigmoid", [x], ["output"])) + return _model("higgs_model_dense", nodes, + [_finfo("input", [1, 7])], [_finfo("output", [1, 2])], inits) + + +def make_recurrent(op_type): + """Recurrent network (RNN, GRU, or LSTM) with 10 inputs, 20 time steps, + and a hidden size of 8, followed by a dense layer on the last step.""" + rng = np.random.default_rng(24) + d, t, h = 10, 20, 8 + ngates = {"RNN": 1, "GRU": 3, "LSTM": 4}[op_type] + rec_inputs = ["xt", "w", "r", "bias", "", "h0"] + rec_outputs = ["y", "yh"] + inits = [_tensor(rng, "w", 1, ngates * h, d), _tensor(rng, "r", 1, ngates * h, h), + _tensor(rng, "bias", 1, 2 * ngates * h), _tensor(rng, "h0", 1, 1, h), + _tensor(rng, "wfc", 2, h), _tensor(rng, "bfc", 2)] + kwargs = {"hidden_size": h} + if op_type == "GRU": + kwargs["linear_before_reset"] = 1 + if op_type == "LSTM": + rec_inputs.append("c0") + rec_outputs.append("yc") + inits.append(_tensor(rng, "c0", 1, 1, h)) + nodes = [ + helper.make_node("Transpose", ["input"], ["xt"], perm=[1, 0, 2]), + helper.make_node(op_type, rec_inputs, rec_outputs, **kwargs), + helper.make_node("Squeeze", ["y"], ["squeezed"], axes=[1]), + helper.make_node("Transpose", ["squeezed"], ["batchfirst"], perm=[1, 0, 2]), + helper.make_node("Slice", ["batchfirst"], ["laststep"], + axes=[1], starts=[-1], ends=[np.iinfo(np.int64).max]), + helper.make_node("Squeeze", ["laststep"], ["flat"], axes=[1]), + helper.make_node("Gemm", ["flat", "wfc", "bfc"], ["output"], transB=1), + ] + return _model(op_type, nodes, + [_finfo("input", [1, t, d])], [_finfo("output", [1, 2])], inits) + + +MODELS = { + "Linear_16": lambda: make_linear(16), + "Linear_32": lambda: make_linear(32), + "Linear_64": lambda: make_linear(64), + "Linear_event": lambda: make_linear(1), + "Generator_B1": lambda: make_generator(1), + "Generator_B64": lambda: make_generator(64), + "Conv_d100_L1_B1": lambda: make_conv2d(1, 1, pads=2), + "Conv_d100_L14_B1": lambda: make_conv2d(14, 1, pads=2), + "Conv_d100_L14_B32": lambda: make_conv2d(14, 32, pads=0), + "Conv3d_d32_L4_B1": make_conv3d, + "ConvTrans2dModel_B1": make_convtrans2d, + "ConvTModel_G4": make_convt_g4, + "SimpleNN_Alice": make_simplenn_alice, + "higgs_model_dense": make_higgs_model_dense, + "RNN_d10_L20_h8_B1": lambda: make_recurrent("RNN"), + "GRU_d10_L20_h8_B1": lambda: make_recurrent("GRU"), + "LSTM_d10_L20_h8_B1": lambda: make_recurrent("LSTM"), +} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("models", nargs="*", help="models to generate (default: all)") + parser.add_argument("--outdir", default=".", help="output directory") + parser.add_argument("--list", action="store_true", help="list available models") + args = parser.parse_args() + + if args.list: + print("\n".join(MODELS)) + return + + names = args.models or list(MODELS) + for name in names: + if name not in MODELS: + parser.error(f"unknown model {name} (--list shows the available ones)") + + os.makedirs(args.outdir, exist_ok=True) + for name in names: + path = os.path.join(args.outdir, f"{name}.onnx") + onnx.save(MODELS[name](), path) + print(f"generated {path}") + + +if __name__ == "__main__": + main()