CODEWITHSUNDEEP
pythonmachine-learningscikit-learndecision-tree
Danissimo
Danissimo

Reputation: 21

Is it possible to get dict {class: proba} from Decision Tree Classifier predict_proba()?

My model have more than 1k classes, and the method returns an array with probabilities, most of which are 0. I want to get top 3 predictions with their probabilities. How can i implement this? I expect to get something like this:

[{class: proba}, {class: proba}, {class: proba}]

Upvotes: 0

Views: 152

Answers (1)

dx2-66
dx2-66

Reputation: 2851

Assuming it's one sample:

{k: probs[k] for k in reversed(np.argpartition(probs,-3)[-3:])}

For a larger test set it'll grow into something like:

[{k: p[k] for k in reversed(np.argpartition(p,-3)[-3:])} for p in probs]

argpartition is like argsort but would only sort out 3 indices with largest values in this notation, saving you some complexity.

Or, if you don't mind importing torch just for this, it's got a nice .topk() method for tensors.

Upvotes: 2

Related Questions