Using a custom threshold value with tf.contrib.lea

2019-09-18 13:11发布

I'm working on a binary classification problem and I'm using the tf.contrib.learn.DNNClassifier class within TensorFlow. When invoking this estimator for only 2 classes, it uses a threshold value of 0.5 as the cutoff between the 2 classes. I'd like to know if there's a way to use a custom threshold value since this might improve the model's accuracy.

I've searched all around the web and apparently there isn't a way to do this.

Any help will be greatly appreciated, thank you.

1条回答
Ridiculous、
2楼-- · 2019-09-18 14:15

The tf.contrib.learn.DNNClassifier class has a method called predict_proba which returns the probabilities belonging to each class for the given inputs. Then you can use something like, tf.round(prob+thres) for binary thresholding with the custom parameter thres.

查看更多
登录 后发表回答