Merge pull request #36 from tomtung/limit-candidate-per-leaf

Limit number of label candidates per leaf for prediction

Cargo.toml

@@ -33,6 +33,7 @@ serde_cbor = "0.11.*"
 serde_json = "1.0.*"
 simple_logger = { version = "1.15.*", features = ["stderr"], optional = true }
 sprs = { version = "0.9.*", features = ["serde"] }
+pdqselect = "0.1.*"
 
 [dev-dependencies]
 assert_approx_eq = "1.1.*"
@@ -47,4 +48,4 @@ cli = ["simple_logger", "clap"]
 
 [profile.release]
 lto = true
-codegen-units = 1
+codegen-units = 1

src/model/mod.rs

@@ -324,7 +324,9 @@ impl TreeNode {
 
  swap(&mut curr_level, &mut next_level);
  if curr_level.len() > beam_size {
- curr_level.sort_unstable_by_key(|&(_, score)| Reverse(NotNan::new(score).unwrap()));
+ pdqselect::select_by_key(curr_level.as_mut_slice(), beam_size, |&(_, score)| {
+ Reverse(NotNan::new(score).unwrap())
+ });
  curr_level.truncate(beam_size);
  }
  }
@@ -336,11 +338,19 @@ impl TreeNode {
  let mut label_scores =
  liblinear::predict(weights, classifier_loss_type, feature_vec);
  label_scores.mapv_inplace(|v| (v + leaf_score).exp());
- labels
+
+ let mut label_score_pairs = labels
  .iter()
  .cloned()
  .zip_eq(label_scores.into_iter().cloned())
- .collect_vec()
+ .collect_vec();
+ pdqselect::select_by_key(
+ label_score_pairs.as_mut_slice(),
+ beam_size,
+ |&(_, score)| Reverse(NotNan::new(score).unwrap()),
+ );
+ label_score_pairs.truncate(beam_size);
+ label_score_pairs
  }
  _ => unreachable!(),
  })