neo-ai / neo-ai-dlr

Neo-AI-DLR is a common runtime for machine learning models compiled by AWS SageMaker Neo, TVM, or TreeLite.
Apache License 2.0
492 stars 106 forks source link

[Treelite] multi:softmax output shape fix #406

Closed apivovarov closed 2 years ago

apivovarov commented 2 years ago

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.

mnist model uses multi:softmax objective which outputs one class with the maximum probability.

Related discussion: https://github.com/dmlc/treelite/issues/351