Optimize Output layer top-k implementation (#210)

aobolensk · web-flow · commit bddbfda45554 · 2025-09-18T12:10:42.000+02:00
diff --git a/include/layers/OutputLayer.hpp b/include/layers/OutputLayer.hpp
@@ -104,7 +104,7 @@ std::pair<std::vector<std::string>, std::vector<ValueType>> top_k_vec(
   for (size_t i = 0; i < input.size(); i++) {
     sort_buf[i] = std::make_pair(labels[i], input[i]);
   }
-  std::sort(sort_buf.begin(), sort_buf.end(), compare_pair<ValueType>);
+  std::partial_sort(sort_buf.begin(), sort_buf.begin() + k, sort_buf.end(), compare_pair<ValueType>);
   std::vector<std::string> res_labels(k);
   std::vector<ValueType> res_input(k);
   for (size_t i = 0; i < k; i++) {

Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,7 @@ std::pair<std::vector<std::string>, std::vector<ValueType>> top_k_vec(`
`104`	`104`	`for (size_t i = 0; i < input.size(); i++) {`
`105`	`105`	`sort_buf[i] = std::make_pair(labels[i], input[i]);`
`106`	`106`	`}`
`107`		`- std::sort(sort_buf.begin(), sort_buf.end(), compare_pair<ValueType>);`
	`107`	`+ std::partial_sort(sort_buf.begin(), sort_buf.begin() + k, sort_buf.end(), compare_pair<ValueType>);`
`108`	`108`	`std::vector<std::string> res_labels(k);`
`109`	`109`	`std::vector<ValueType> res_input(k);`
`110`	`110`	`for (size_t i = 0; i < k; i++) {`