当前位置:   article > 正文

Bert文本分类run_classifier的预测模块修改_runclassifier.py修改

runclassifier.py修改

修改位置1:run_classifier.py model_fn() 函数中

源码1:

else:
	output_spec = tf.contrib.tpu.TPUEstimatorSpec(
		mode=mode, predictions=probabilities, scaffold_fn=scaffold_fn)
  • 1
  • 2
  • 3

替换源码1:

elif mode == tf.estimator.ModeKeys.PREDICT:
    def metric_fn(logits,probabilities):
        predicted_classes = tf.argmax(logits, axis=1,output_type=tf.int32)
        return {
             'pred_class_ids': predicted_classes[:, tf.newaxis],
             'probabilities':probabilities,
             'logits': logits}                

    pred_metrics = metric_fn(logits,probabilities)   
    output_spec = tf.estimator.EstimatorSpec(
        mode=mode,predictions=pred_metrics)         
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11

修改位置2:run_classifier.py main()函数中

源码2:

with tf.gfile.GFile(output_predict_file, "w") as writer:
   
  • 1
声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/小丑西瓜9/article/detail/312363
推荐阅读
相关标签
  

闽ICP备14008679号