Improve fetching, but have to pass settings by value, not const ref

Author: ChSonnabend

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -255,6 +255,7 @@ AddOption(nnSigmoidTrafoClassThreshold, int, 1, "", 0, "If true (default), then
 AddOption(nnLoadFromCCDB, int, 1, "", 0, "If 1 networks are fetched from ccdb, else locally")
 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/Global/GPUChainTrackingClusterizer.cxx

@@ -40,6 +40,7 @@
 #endif
 
 #ifdef GPUCA_HAS_ONNX
+#include <CommonUtils/StringUtils.h>
 #include "GPUTPCNNClusterizerKernels.h"
 #include "GPUTPCNNClusterizerHost.h"
 #endif
@@ -612,7 +613,7 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
   }
 
 #ifdef GPUCA_HAS_ONNX
-  const GPUSettingsProcessingNNclusterizer& nn_settings = GetProcessingSettings().nn;
+  GPUSettingsProcessingNNclusterizer nn_settings = GetProcessingSettings().nn;
   GPUTPCNNClusterizerHost nnApplication; // potentially this needs to be GPUTPCNNClusterizerHost nnApplication[NSECTORS]; Technically ONNX ->Run() is threadsafe at inference time since its read-only
   if (GetProcessingSettings().nn.applyNNclusterizer) {
     if(nn_settings.nnLoadFromCCDB) {
@@ -626,17 +627,35 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
         {"nnCCDBInteractionRate", std::to_string(nn_settings.nnCCDBInteractionRate)}
       };
 
+      std::string nnFetchFolder = "";
+      std::vector<std::string> fetchMode = o2::utils::Str::tokenize(nn_settings.nnCCDBFetchMode, ':');
       std::map<std::string, std::string> networkRetrieval = ccdbSettings;
 
-      networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBClassificationLayerType;
-      networkRetrieval["nnCCDBEvalType"] = "classification_c1";
-      networkRetrieval["outputFile"] = "net_classification_c1.onnx";
-      nnApplication.loadFromCCDB(networkRetrieval);
+      if (fetchMode[0] == "c1") {
+        networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBClassificationLayerType;
+        networkRetrieval["nnCCDBEvalType"] = "classification_c1";
+        networkRetrieval["outputFile"] = nnFetchFolder + "net_classification_c1.onnx";
+        nnApplication.loadFromCCDB(networkRetrieval);
+      } else if (fetchMode[0] == "c2") {
+        networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBClassificationLayerType;
+        networkRetrieval["nnCCDBEvalType"] = "classification_c2";
+        networkRetrieval["outputFile"] = nnFetchFolder + "net_classification_c2.onnx";
+        nnApplication.loadFromCCDB(networkRetrieval);
+      }
+      nn_settings.nnClassificationPath = networkRetrieval["outputFile"]; // Setting the proper path from the where the models will be initialized locally
 
       networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBRegressionLayerType;
       networkRetrieval["nnCCDBEvalType"] = "regression_c1";
-      networkRetrieval["outputFile"] = "net_regression_c1.onnx";
+      networkRetrieval["outputFile"] = nnFetchFolder + "net_regression_c1.onnx";
       nnApplication.loadFromCCDB(networkRetrieval);
+      nn_settings.nnRegressionPath = networkRetrieval["outputFile"];
+      if (fetchMode[1] == "r2") {
+        networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBRegressionLayerType;
+        networkRetrieval["nnCCDBEvalType"] = "regression_c2";
+        networkRetrieval["outputFile"] = nnFetchFolder + "net_regression_c2.onnx";
+        nnApplication.loadFromCCDB(networkRetrieval);
+        nn_settings.nnRegressionPath += ":", networkRetrieval["outputFile"];
+      }
     }
 
     uint32_t maxClusters = 0;

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx

@@ -22,7 +22,7 @@
 
 using namespace o2::gpu;
 
-GPUTPCNNClusterizerHost::GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer& settings)
+GPUTPCNNClusterizerHost::GPUTPCNNClusterizerHost(GPUSettingsProcessingNNclusterizer settings)
 {
   init(settings);
 }
@@ -53,7 +53,7 @@ void GPUTPCNNClusterizerHost::loadFromCCDB(std::map<std::string, std::string> se
   }
 }
 
-void GPUTPCNNClusterizerHost::init(const GPUSettingsProcessingNNclusterizer& settings)
+void GPUTPCNNClusterizerHost::init(GPUSettingsProcessingNNclusterizer settings)
 {
   OrtOptions = {
     {"model-path", settings.nnClassificationPath},

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.h

@@ -37,9 +37,9 @@ class GPUTPCNNClusterizerHost
 {
  public:
   GPUTPCNNClusterizerHost() = default;
-  GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer&);
+  GPUTPCNNClusterizerHost(GPUSettingsProcessingNNclusterizer);
 
-  void init(const GPUSettingsProcessingNNclusterizer&);
+  void init(GPUSettingsProcessingNNclusterizer);
   void initClusterizer(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
   void loadFromCCDB(std::map<std::string, std::string>);