feat(ml): tensorrt support for regression models

jolibrain · Oct 21, 2020 · 77a016b · 77a016b
1 parent a8b81f2
commit 77a016b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/docs/api.md b/docs/api.md
@@ -1223,6 +1223,7 @@ measure              | array  | yes      | empty   | Output measures requested,
 template             | string | yes      | empty   | Output template in Mustache format
 confidence_threshold | double | yes      | 0.0     | only returns classifications or detections with probability strictly above threshold
 bbox                 | bool   | yes      | false   | returns bounding boxes around object when using an object detection model
+regression           | bool   | yes      | false   | whether the output of a model is a regression target (i.e. vector of one or more floats)
 
 The variables that are usable in the output template format are those from the standard JSON output. See the [output template](#output-templates) dedicated section for more details and examples.
 

diff --git a/src/backends/tensorrt/tensorrtlib.cc b/src/backends/tensorrt/tensorrtlib.cc
@@ -94,10 +94,11 @@ namespace dd
     _readEngine = tl._readEngine;
     _writeEngine = tl._writeEngine;
     _TRTContextReady = tl._TRTContextReady;
-    _timeserie = tl._timeserie;
     _buffers = tl._buffers;
     _bbox = tl._bbox;
     _ctc = tl._ctc;
+    _timeserie = tl._timeserie;
+    _regression = tl._regression;
     _inputIndex = tl._inputIndex;
     _outputIndex0 = tl._outputIndex0;
     _outputIndex1 = tl._outputIndex1;
@@ -383,6 +384,9 @@ namespace dd
       {
         if (ad_output.has("bbox"))
           _bbox = ad_output.get("bbox").get<bool>();
+        if (ad_output.has("regression"))
+          _regression = ad_output.get("regression").get<bool>();
+
         // Ctc model
         if (ad_output.has("ctc"))
           {
@@ -410,6 +414,8 @@ namespace dd
             throw MLLibBadParamException(
                 "timeseries not yet implemented over tensorRT backend");
           }
+        else if (_regression)
+          out_blob = "pred";
 
         if (_nclasses == 0)
           {
@@ -529,9 +535,12 @@ namespace dd
             throw MLLibBadParamException(
                 "timeseries not yet implemented over tensorRT backend");
           }
-        else // classification
+        else // classification / regression
           {
-            _buffers.resize(2);
+            if (_regression)
+              _buffers.resize(1);
+            else
+              _buffers.resize(2);
             _floatOut.resize(_max_batch_size * this->_nclasses);
             if (inputc._bw)
               cudaMalloc(&_buffers.data()[_inputIndex],
@@ -612,7 +621,7 @@ namespace dd
                 throw MLLibBadParamException(
                     "timeseries not yet implemented over tensorRT backend");
               }
-            else // classification
+            else // classification / regression
               {
                 if (inputc._bw)
                   cudaMemcpyAsync(_buffers.data()[_inputIndex], inputc.data(),
@@ -745,7 +754,7 @@ namespace dd
             throw MLLibBadParamException(
                 "timeseries not yet implemented over tensorRT backend");
           }
-        else // classification
+        else // classification / regression
           {
             for (int j = 0; j < num_processed; j++)
               {
@@ -761,7 +770,7 @@ namespace dd
                 for (int i = 0; i < _nclasses; i++)
                   {
                     double prob = _floatOut.at(j * _nclasses + i);
-                    if (prob < confidence_threshold)
+                    if (prob < confidence_threshold && !_regression)
                       continue;
                     probs.push_back(prob);
                     cats.push_back(this->_mlmodel.get_hcorresp(i));
@@ -782,6 +791,8 @@ namespace dd
     out.add("nclasses", this->_nclasses);
     if (_bbox)
       out.add("bbox", true);
+    if (_regression)
+      out.add("regression", true);
     out.add("roi", false);
     out.add("multibox_rois", false);
     tout.finalize(ad.getobj("parameters").getobj("output"), out,

diff --git a/src/backends/tensorrt/tensorrtlib.h b/src/backends/tensorrt/tensorrtlib.h
@@ -127,11 +127,12 @@ namespace dd
 
     bool _bbox = false;
     bool _ctc = false;
+    bool _regression = false;
+    bool _timeserie = false;
 
     std::vector<void *> _buffers;
 
     bool _TRTContextReady = false;
-    bool _timeserie = false;
 
     int _inputIndex;
     int _outputIndex0;