INNER CODE UNIT · Rust

predict

nogibjj/rust-mlops-template · actixtorch/src/lib.rs:13

pub fn predict() -> Result<()> {
    let weights = std::path::Path::new(MODEL);
    let image = IMAGE.to_owned();
    let image = imagenet::load_image_and_resize224(image)?;

    // Create the model and load the weights from the file.
    let mut vs = tch::nn::VarStore::new(tch::Device::Cpu);
    let net: Box<dyn ModuleT> = match weights.file_name().unwrap().to_str().unwrap() {
        "resnet18.ot" => Box::new(resnet::resnet18(&vs.root(), imagenet::CLASS_COUNT)),

        _ => bail!("unknown model, use a weight file named e.g. resnet18.ot"),
    };
    vs.load(weights)?;

    // Apply the forward pass of the model to get the logits.
    let output = net
        .forward_t(&image.unsqueeze(0), /* train= */ false)
        .softmax(-1, tch::Kind::Float); // Convert to probability.

View source record →