huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Print the recognized categories in dino-v2. (#506)

Browse Source
This commit is contained in:
Laurent Mazare
2023-08-18 17:32:58 +01:00
committed by GitHub
parent cb069d6063
commit e5dd5fd1b3
2 changed files with 1018 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
candle-examples/examples/dinov2/main.rs
View File

@ -315,7 +315,17 @@ pub fn main() -> anyhow::Result<()> {
let model = vit_small(vb)?;
println!("model built");
let logits = model.forward(&image.unsqueeze(0)?)?;
let prs = candle_nn::ops::softmax(&logits, D::Minus1)?;
println!("{prs}");
let prs = candle_nn::ops::softmax(&logits, D::Minus1)?
.i(0)?
.to_vec1::<f32>()?;
let mut prs = prs.iter().enumerate().collect::<Vec<_>>();
prs.sort_by(|(_, p1), (_, p2)| p2.total_cmp(p1));
for &(category_idx, pr) in prs.iter().take(5) {
println!(
"{:24}: {:.2}%",
candle_examples::IMAGENET_CLASSES[category_idx],
100. * pr
);
}
Ok(())
}