[Treelite] multi:softmax output shape fix (#406)

python/dlr/dlr_model.py

@@ -325,13 +325,18 @@ def _set_input(self, name, data):
                                      shape.ctypes.data_as(POINTER(c_longlong)),
                                      in_data_pointer,
                                      c_int(len(shape))))
-        if self.backend == 'treelite':
+        # This helps to determine output batch size only
+        # Treelite model output shape will be know only after the model run
+        if self.backend == "treelite":
             self._lazy_init_output_shape()
 
     def _run(self):
         """A light wrapper to call run in the DLR backend."""
         self._check_call(self._lib.RunDLRModel(byref(self.handle)))
-        if self.backend == "relayvm":
+        # Treelite model output shape will be know only after the model run
+        # If model uses objective multi:softmax then output shape will be (batch, 1)
+        # because predictor will execute predictor_transform max_index
+        if self.backend in ["relayvm", "treelite"]:
             self._lazy_init_output_shape()
 
     def _get_num_outputs(self):

src/dlr_treelite.cc

@@ -228,6 +228,10 @@ void TreeliteModel::Run() {
                                          &out_result_size),
            0)
       << TreeliteGetLastError();
+  // Treelite model output shape will be know only after the model run
+  // If model uses objective multi:softmax then output shape will be (batch, 1)
+  // because predictor will execute predictor_transform max_index
+  treelite_output_size_ = out_result_size / treelite_input_->num_row;
 }
 
 void TreeliteModel::SetNumThreads(int threads) {

tests/python/integration/load_and_run_treelite_model.py

@@ -16,13 +16,13 @@ def test_mnist():
     data_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost', 'mnist.libsvm')
     model = DLRModel(model_path, 'cpu', 0)
 
-    X, _ = load_svmlight_file(data_file)
+    X, y = load_svmlight_file(data_file)
+    assert y.shape == (8,)
     print('Testing inference on XGBoost MNIST...')
     res = model.run(X.toarray())[0]
-    # TODO investigate why output shape size in (1,10) instead of (1,1)
-    # the model uses multi:softmax objective which outputs one class with the maximum probability
-    # assert res.shape == (1,1)
-    assert res[0][0] == 7.0
+    # mnist model uses multi:softmax objective which outputs one class with the maximum probability
+    assert res.shape == (8, 1)
+    assert np.allclose(res.flatten(), y)
 
 def test_iris():
     model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-iris-1.10.0')