CCDB fetching of NNs ported to GPUWorkflowSpec

Author: ChSonnabend

Detectors/TPC/calibration/CMakeLists.txt

@@ -25,6 +25,7 @@ o2_add_library(TPCCalibration
                        src/CalibPadGainTracksBase.cxx
                        src/CalibLaserTracks.cxx
                        src/LaserTracksCalibrator.cxx
+                       src/NeuralNetworkClusterizer.cxx
                        src/SACDecoder.cxx
                        src/IDCAverageGroup.cxx
                        src/IDCAverageGroupBase.cxx
@@ -82,6 +83,7 @@ o2_target_root_dictionary(TPCCalibration
                                   include/TPCCalibration/FastHisto.h
                                   include/TPCCalibration/CalibLaserTracks.h
                                   include/TPCCalibration/LaserTracksCalibrator.h
+                                  include/TPCCalibration/NeuralNetworkClusterizer.h
                                   include/TPCCalibration/SACDecoder.h
                                   include/TPCCalibration/IDCAverageGroup.h
                                   include/TPCCalibration/IDCAverageGroupBase.h

Detectors/TPC/calibration/include/TPCCalibration/NeuralNetworkClusterizer.h

@@ -0,0 +1,39 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file   NeuralNetworkClusterizer.h
+/// \brief  Fetching neural networks for clusterization from CCDB
+/// \author Christian Sonnabend
+
+#ifndef AliceO2_TPC_NeuralNetworkClusterizer_h
+#define AliceO2_TPC_NeuralNetworkClusterizer_h
+
+#include "CCDB/CcdbApi.h"
+
+namespace o2::tpc
+{
+
+class NeuralNetworkClusterizer
+{
+    public:
+        NeuralNetworkClusterizer() = default;
+        void initCcdbApi(std::string url);
+        void loadIndividualFromCCDB(std::map<std::string, std::string> settings);
+
+    private:
+        o2::ccdb::CcdbApi ccdbApi;
+        std::map<std::string, std::string> metadata;
+        std::map<std::string, std::string> headers;
+
+};
+
+} // namespace o2::tpc
+#endif

Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx

@@ -0,0 +1,47 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file   NeuralNetworkClusterizer.cxx
+/// \brief  Fetching neural networks for clusterization from CCDB
+/// \author Christian Sonnabend
+
+#include <CommonUtils/StringUtils.h>
+#include "TPCCalibration/NeuralNetworkClusterizer.h"
+
+using namespace o2::tpc;
+
+void NeuralNetworkClusterizer::initCcdbApi(std::string url) {
+  ccdbApi.init(url);
+}
+
+void NeuralNetworkClusterizer::loadIndividualFromCCDB(std::map<std::string, std::string> settings)
+{
+  metadata["inputDType"] = settings["inputDType"];
+  metadata["outputDType"] = settings["outputDType"];
+  metadata["nnCCDBEvalType"] = settings["nnCCDBEvalType"];         // classification_1C, classification_2C, regression_1C, regression_2C
+  metadata["nnCCDBWithMomentum"] = settings["nnCCDBWithMomentum"]; // 0, 1 -> Only for regression model
+  metadata["nnCCDBLayerType"] = settings["nnCCDBLayerType"];       // FC, CNN
+  if (settings["nnCCDBInteractionRate"] != "" && std::stoi(settings["nnCCDBInteractionRate"]) > 0) {
+    metadata["nnCCDBInteractionRate"] = settings["nnCCDBInteractionRate"];
+  }
+  if (settings["nnCCDBBeamType"] != "") {
+    metadata["nnCCDBBeamType"] = settings["nnCCDBBeamType"];
+  }
+
+  bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnCCDBPath"], settings["outputFolder"], metadata, 1, false, settings["outputFile"]);
+  // headers = ccdbApi.retrieveHeaders(settings["nnPathCCDB"], metadata, 1); // potentially needed to init some local variables
+
+  if (retrieveSuccess) {
+    LOG(info) << "Network " << settings["nnCCDBPath"] << " retrieved from CCDB, stored at " << settings["outputFile"];
+  } else {
+    LOG(error) << "Failed to retrieve network from CCDB";
+  }
+}

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -268,11 +268,12 @@ AddOption(nnClassificationPath, std::string, "network_class.onnx", "", 0, "The c
 AddOption(nnClassThreshold, float, 0.5, "", 0, "The cutoff at which clusters will be accepted / rejected.")
 AddOption(nnRegressionPath, std::string, "network_reg.onnx", "", 0, "The regression network path")
 AddOption(nnSigmoidTrafoClassThreshold, int, 1, "", 0, "If true (default), then the classification threshold is transformed by an inverse sigmoid function. This depends on how the network was trained (with a sigmoid as acitvation function in the last layer or not).")
+AddOption(nnEvalMode, std::string, "c1:r1", "", 0, "Concatention of modes, e.g. c1:r1 (classification class 1, regression class 1)")
 // CCDB
 AddOption(nnLoadFromCCDB, int, 1, "", 0, "If 1 networks are fetched from ccdb, else locally")
+AddOption(nnLocalFolder, std::string, ".", "", 0, "Local folder in which the networks will be fetched")
 AddOption(nnCCDBURL, std::string, "http://ccdb-test.cern.ch:8080", "", 0, "The CCDB URL from where the network files are fetched")
 AddOption(nnCCDBPath, std::string, "Users/c/csonnabe/TPC/Clusterization", "", 0, "Folder path containing the networks")
-AddOption(nnCCDBFetchMode, std::string, "c1:r1", "", 0, "Concatention of modes, e.g. c1:r1 (classification class 1, regression class 1)")
 AddOption(nnCCDBWithMomentum, int, 1, "", 0, "Distinguishes between the network with and without momentum output for the regression")
 AddOption(nnCCDBClassificationLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
 AddOption(nnCCDBRegressionLayerType, std::string, "CNN", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx

@@ -16,82 +16,27 @@
 
 #include "GPUTPCNNClusterizerHost.h"
 #include "GPUTPCNNClusterizer.h"
-#include "CCDB/CcdbApi.h"
 #include "GPUSettings.h"
 #include "ML/3rdparty/GPUORTFloat16.h"
 
 using namespace o2::gpu;
 
-void GPUTPCNNClusterizerHost::loadFromCCDB(std::map<std::string, std::string> settings)
-{
-  o2::ccdb::CcdbApi ccdbApi;
-  ccdbApi.init(settings["nnCCDBURL"]);
-
-  metadata["inputDType"] = settings["inputDType"];
-  metadata["outputDType"] = settings["outputDType"];
-  metadata["nnCCDBEvalType"] = settings["nnCCDBEvalType"];         // classification_1C, classification_2C, regression_1C, regression_2C
-  metadata["nnCCDBWithMomentum"] = settings["nnCCDBWithMomentum"]; // 0, 1 -> Only for regression model
-  metadata["nnCCDBLayerType"] = settings["nnCCDBLayerType"];       // FC, CNN
-  if (settings["nnCCDBInteractionRate"] != "" && std::stoi(settings["nnCCDBInteractionRate"]) > 0) {
-    metadata["nnCCDBInteractionRate"] = settings["nnCCDBInteractionRate"];
-  }
-  if (settings["nnCCDBBeamType"] != "") {
-    metadata["nnCCDBBeamType"] = settings["nnCCDBBeamType"];
-  }
-
-  bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnCCDBPath"], ".", metadata, 1, false, settings["outputFile"]);
-  // headers = ccdbApi.retrieveHeaders(settings["nnPathCCDB"], metadata, 1); // potentially needed to init some local variables
-
-  if (retrieveSuccess) {
-    LOG(info) << "Network " << settings["nnCCDBPath"] << " retrieved from CCDB, stored at " << settings["outputFile"];
-  } else {
-    LOG(error) << "Failed to retrieve network from CCDB";
-  }
-}
-
 void GPUTPCNNClusterizerHost::init(const GPUSettingsProcessingNNclusterizer& settings)
 {
   std::string class_model_path = settings.nnClassificationPath, reg_model_path = settings.nnRegressionPath;
   std::vector<std::string> reg_model_paths;
-
-  if (settings.nnLoadFromCCDB) {
-    std::map<std::string, std::string> ccdbSettings = {
-      {"nnCCDBURL", settings.nnCCDBURL},
-      {"nnCCDBPath", settings.nnCCDBPath},
-      {"inputDType", settings.nnInferenceInputDType},
-      {"outputDType", settings.nnInferenceOutputDType},
-      {"nnCCDBWithMomentum", std::to_string(settings.nnCCDBWithMomentum)},
-      {"nnCCDBBeamType", settings.nnCCDBBeamType},
-      {"nnCCDBInteractionRate", std::to_string(settings.nnCCDBInteractionRate)}};
-
-    std::string nnFetchFolder = "";
-    std::vector<std::string> fetchMode = o2::utils::Str::tokenize(settings.nnCCDBFetchMode, ':');
-    std::map<std::string, std::string> networkRetrieval = ccdbSettings;
-
-    if (fetchMode[0] == "c1") {
-      networkRetrieval["nnCCDBLayerType"] = settings.nnCCDBClassificationLayerType;
-      networkRetrieval["nnCCDBEvalType"] = "classification_c1";
-      networkRetrieval["outputFile"] = nnFetchFolder + "net_classification_c1.onnx";
-      loadFromCCDB(networkRetrieval);
-    } else if (fetchMode[0] == "c2") {
-      networkRetrieval["nnCCDBLayerType"] = settings.nnCCDBClassificationLayerType;
-      networkRetrieval["nnCCDBEvalType"] = "classification_c2";
-      networkRetrieval["outputFile"] = nnFetchFolder + "net_classification_c2.onnx";
-      loadFromCCDB(networkRetrieval);
+  std::vector<std::string> evalMode = o2::utils::Str::tokenize(settings.nnEvalMode, ':');
+
+  if(settings.nnLoadFromCCDB) {
+    reg_model_path = settings.nnLocalFolder + "/net_regression_c1.onnx"; // Needs to be set identical to NeuralNetworkClusterizer.cxx, otherwise the networks might be loaded from the wrong place
+    if (evalMode[0] == "c1") {
+      class_model_path = settings.nnLocalFolder + "/net_classification_c1.onnx";
+    } else if (evalMode[0] == "c2") {
+      class_model_path = settings.nnLocalFolder + "/net_classification_c2.onnx";
     }
-    class_model_path = networkRetrieval["outputFile"]; // Setting the proper path from the where the models will be initialized locally
 
-    networkRetrieval["nnCCDBLayerType"] = settings.nnCCDBRegressionLayerType;
-    networkRetrieval["nnCCDBEvalType"] = "regression_c1";
-    networkRetrieval["outputFile"] = nnFetchFolder + "net_regression_c1.onnx";
-    loadFromCCDB(networkRetrieval);
-    reg_model_path = networkRetrieval["outputFile"];
-    if (fetchMode[1] == "r2") {
-      networkRetrieval["nnCCDBLayerType"] = settings.nnCCDBRegressionLayerType;
-      networkRetrieval["nnCCDBEvalType"] = "regression_c2";
-      networkRetrieval["outputFile"] = nnFetchFolder + "net_regression_c2.onnx";
-      loadFromCCDB(networkRetrieval);
-      reg_model_path += ":", networkRetrieval["outputFile"];
+    if (evalMode[1] == "r2") {
+      reg_model_path += ":" + settings.nnLocalFolder + "/net_regression_c2.onnx";
     }
   }
 

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.h

@@ -48,10 +48,6 @@ class GPUTPCNNClusterizerHost
   std::vector<bool> modelsUsed = {false, false, false}; // 0: class, 1: reg_1, 2: reg_2
   int32_t deviceId = -1;
   std::vector<std::string> reg_model_paths;
-
- private:
-  std::map<std::string, std::string> metadata;
-  std::map<std::string, std::string> headers;
 }; // class GPUTPCNNClusterizerHost
 
 } // namespace o2::gpu

GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h

@@ -83,6 +83,7 @@ class GPUO2Interface;
 struct TPCPadGainCalib;
 struct TPCZSLinkMapping;
 struct GPUSettingsO2;
+struct GPUSettingsProcessingNNclusterizer;
 class GPUO2InterfaceQA;
 struct GPUTrackingInOutPointers;
 struct GPUTrackingInOutZS;
@@ -224,6 +225,8 @@ class GPURecoWorkflowSpec : public o2::framework::Task
   uint32_t mNextThreadIndex = 0;
   bool mUpdateGainMapCCDB = true;
   std::unique_ptr<o2::gpu::GPUSettingsTF> mTFSettings;
+  std::unique_ptr<o2::gpu::GPUSettingsProcessingNNclusterizer> mNNClusterizerSettings;
+
   Config mSpecConfig;
   std::shared_ptr<o2::base::GRPGeomRequest> mGGR;
   bool mGRPGeomUpdated = false;

GPU/Workflow/src/GPUWorkflowSpec.cxx

@@ -78,6 +78,7 @@
 #include "DetectorsRaw/RDHUtils.h"
 #include "ITStracking/TrackingInterface.h"
 #include "GPUWorkflowInternal.h"
+#include "TPCCalibration/NeuralNetworkClusterizer.h"
 // #include "Framework/ThreadPool.h"
 
 #include <TStopwatch.h>
@@ -118,6 +119,7 @@ GPURecoWorkflowSpec::GPURecoWorkflowSpec(GPURecoWorkflowSpec::CompletionPolicyDa
   mConfig.reset(new GPUO2InterfaceConfiguration);
   mConfParam.reset(new GPUSettingsO2);
   mTFSettings.reset(new GPUSettingsTF);
+  mNNClusterizerSettings.reset(new GPUSettingsProcessingNNclusterizer);
   mTimer.reset(new TStopwatch);
   mPipeline.reset(new GPURecoWorkflowSpec_PipelineInternals);
 
@@ -133,6 +135,49 @@ void GPURecoWorkflowSpec::init(InitContext& ic)
   GRPGeomHelper::instance().setRequest(mGGR);
   GPUO2InterfaceConfiguration& config = *mConfig.get();
 
+  if (mNNClusterizerSettings->nnLoadFromCCDB){
+    LOG(info) << "Loading neural networks from CCDB";
+    o2::tpc::NeuralNetworkClusterizer nnClusterizerFetcher;
+    nnClusterizerFetcher.initCcdbApi(mNNClusterizerSettings->nnCCDBURL);
+    std::map<std::string, std::string> ccdbSettings = {
+    {"nnCCDBURL", mNNClusterizerSettings->nnCCDBURL},
+    {"nnCCDBPath", mNNClusterizerSettings->nnCCDBPath},
+    {"inputDType", mNNClusterizerSettings->nnInferenceInputDType},
+    {"outputDType", mNNClusterizerSettings->nnInferenceOutputDType},
+    {"outputFolder", mNNClusterizerSettings->nnLocalFolder},
+    {"nnCCDBPath", mNNClusterizerSettings->nnCCDBPath},
+    {"nnCCDBWithMomentum", std::to_string(mNNClusterizerSettings->nnCCDBWithMomentum)},
+    {"nnCCDBBeamType", mNNClusterizerSettings->nnCCDBBeamType},
+    {"nnCCDBInteractionRate", std::to_string(mNNClusterizerSettings->nnCCDBInteractionRate)}};
+
+    std::string nnFetchFolder = mNNClusterizerSettings->nnLocalFolder;
+    std::vector<std::string> evalMode = o2::utils::Str::tokenize(mNNClusterizerSettings->nnEvalMode, ':');
+
+    if (evalMode[0] == "c1") {
+      ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings->nnCCDBClassificationLayerType;
+      ccdbSettings["nnCCDBEvalType"] = "classification_c1";
+      ccdbSettings["outputFile"] = "net_classification_c1.onnx";
+      nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
+    } else if (evalMode[0] == "c2") {
+      ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings->nnCCDBClassificationLayerType;
+      ccdbSettings["nnCCDBEvalType"] = "classification_c2";
+      ccdbSettings["outputFile"] = "net_classification_c2.onnx";
+      nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
+    }
+
+    ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings->nnCCDBRegressionLayerType;
+    ccdbSettings["nnCCDBEvalType"] = "regression_c1";
+    ccdbSettings["outputFile"] = "net_regression_c1.onnx";
+    nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
+    if (evalMode[1] == "r2") {
+      ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings->nnCCDBRegressionLayerType;
+      ccdbSettings["nnCCDBEvalType"] = "regression_c2";
+      ccdbSettings["outputFile"] = "net_regression_c2.onnx";
+      nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
+    }
+    LOG(info) << "Neural network loading done!";
+  }
+
   // Create configuration object and fill settings
   mConfig->configGRP.solenoidBzNominalGPU = 0;
   mTFSettings->hasSimStartOrbit = 1;