Implementation:LaurentMazare Tch rs Imagenet Top

Knowledge Sources	tch-rs
Domains	Computer_Vision, Model_Inference
Last Updated	2026-02-08 14:00 GMT

Overview

Concrete tool for extracting top-K ImageNet class predictions from a probability tensor provided by the tch vision module.

Description

imagenet::top takes a probability tensor and returns the top K (probability, class_name) pairs. It handles input shapes of [1000], [1, 1000], or [1, 1, 1000] by flattening to [1000]. The function uses Tensor::topk for efficient selection and maps indices to the 1000 ImageNet class names stored in the CLASSES constant array.

Usage

Use after running a pretrained ImageNet model and applying softmax to the output logits. Typically called with k=5 for top-5 predictions.

Code Reference

Source Location

• Repository: tch-rs
• File: src/vision/imagenet.rs
• Lines: 1155-1170

Signature

pub fn top(tensor: &Tensor, k: i64) -> Vec<(f64, String)>

Import

use tch::vision::imagenet;

I/O Contract

Inputs

Name	Type	Required	Description
tensor	&Tensor	Yes	Probability tensor of shape [1000], [1, 1000], or [1, 1, 1000]
k	i64	Yes	Number of top predictions to return

Outputs

Name	Type	Description
Vec<(f64, String)>	Vec	Top-k (probability, class_name) pairs, sorted descending

Usage Examples

use tch::{vision::imagenet, Kind};

let probs = output.softmax(-1, Kind::Float);
let top5 = imagenet::top(&probs, 5);

for (prob, class) in &top5 {
    println!("{:.2}% {}", prob * 100.0, class);
}
// Example output:
// 89.34% golden retriever
// 4.21% Labrador retriever
// ...

Related Pages

Implements Principle

• Principle:LaurentMazare_Tch_rs_Top_K_Prediction