forked from huggingface/candle
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add Beit model ( https://arxiv.org/abs/2106.08254 ) (huggingface#2305)
Co-authored-by: v-espitalier <>
- Loading branch information
1 parent
74e9e41
commit 7f1ba80
Showing
4 changed files
with
467 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,20 @@ | ||
# candle-beit | ||
|
||
[Beit](https://arxiv.org/abs/2106.08254) is a computer vision model. | ||
In this example, it is used as an ImageNet classifier: the model returns the | ||
probability for the image to belong to each of the 1000 ImageNet categories. | ||
|
||
## Running some example | ||
|
||
```bash | ||
cargo run --example beit --release -- --image candle-examples/examples/yolo-v8/assets/bike.jpg | ||
|
||
> mountain bike, all-terrain bike, off-roader: 56.16% | ||
> bicycle-built-for-two, tandem bicycle, tandem: 3.08% | ||
> maillot : 2.23% | ||
> alp : 0.88% | ||
> crash helmet : 0.85% | ||
|
||
``` | ||
|
||
![Leading group, Giro d'Italia 2021](../yolo-v8/assets/bike.jpg) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,79 @@ | ||
//! BEiT: BERT Pre-Training of Image Transformers | ||
//! https://github.com/microsoft/unilm/tree/master/beit | ||
|
||
#[cfg(feature = "mkl")] | ||
extern crate intel_mkl_src; | ||
|
||
#[cfg(feature = "accelerate")] | ||
extern crate accelerate_src; | ||
|
||
use clap::Parser; | ||
|
||
use candle::{DType, Device, IndexOp, Result, Tensor, D}; | ||
use candle_nn::{Module, VarBuilder}; | ||
use candle_transformers::models::beit; | ||
|
||
/// Loads an image from disk using the image crate, this returns a tensor with shape | ||
/// (3, 384, 384). Beit special normalization is applied. | ||
pub fn load_image384_beit_norm<P: AsRef<std::path::Path>>(p: P) -> Result<Tensor> { | ||
let img = image::io::Reader::open(p)? | ||
.decode() | ||
.map_err(candle::Error::wrap)? | ||
.resize_to_fill(384, 384, image::imageops::FilterType::Triangle); | ||
let img = img.to_rgb8(); | ||
let data = img.into_raw(); | ||
let data = Tensor::from_vec(data, (384, 384, 3), &Device::Cpu)?.permute((2, 0, 1))?; | ||
let mean = Tensor::new(&[0.5f32, 0.5, 0.5], &Device::Cpu)?.reshape((3, 1, 1))?; | ||
let std = Tensor::new(&[0.5f32, 0.5, 0.5], &Device::Cpu)?.reshape((3, 1, 1))?; | ||
(data.to_dtype(candle::DType::F32)? / 255.)? | ||
.broadcast_sub(&mean)? | ||
.broadcast_div(&std) | ||
} | ||
|
||
#[derive(Parser)] | ||
struct Args { | ||
#[arg(long)] | ||
model: Option<String>, | ||
|
||
#[arg(long)] | ||
image: String, | ||
|
||
/// Run on CPU rather than on GPU. | ||
#[arg(long)] | ||
cpu: bool, | ||
} | ||
|
||
pub fn main() -> anyhow::Result<()> { | ||
let args = Args::parse(); | ||
|
||
let device = candle_examples::device(args.cpu)?; | ||
|
||
let image = load_image384_beit_norm(args.image)?.to_device(&device)?; | ||
println!("loaded image {image:?}"); | ||
|
||
let model_file = match args.model { | ||
None => { | ||
let api = hf_hub::api::sync::Api::new()?; | ||
let api = api.model("vincent-espitalier/candle-beit".into()); | ||
api.get("beit_base_patch16_384.in22k_ft_in22k_in1k_adapted.safetensors")? | ||
} | ||
Some(model) => model.into(), | ||
}; | ||
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, &device)? }; | ||
let model = beit::vit_base(vb)?; | ||
println!("model built"); | ||
let logits = model.forward(&image.unsqueeze(0)?)?; | ||
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(()) | ||
} |
Oops, something went wrong.