mirror of
https://github.com/huggingface/candle.git
synced 2025-06-17 11:08:52 +00:00
4845d5cc64
* More realistic training setup. * Compute the model accuracy. * Very inefficient backprop for index select. * More backprop. * Fix some backprop issues. * Backprop fix. * Another broadcasting backprop fix. * Better backprop for reducing ops. * Training again. * Add some gradient tests. * Get the training to work.
85 lines
3.2 KiB
Rust
85 lines
3.2 KiB
Rust
// This should rearch 91.5% accuracy.
|
|
#[cfg(feature = "mkl")]
|
|
extern crate intel_mkl_src;
|
|
|
|
use anyhow::Result;
|
|
use candle::{DType, Tensor, Var, D};
|
|
|
|
const IMAGE_DIM: usize = 784;
|
|
const LABELS: usize = 10;
|
|
|
|
fn log_softmax<D: candle::shape::Dim>(xs: &Tensor, d: D) -> candle::Result<Tensor> {
|
|
let d = d.to_index(xs.shape(), "log-softmax")?;
|
|
let max = xs.max_keepdim(d)?;
|
|
let diff = xs.broadcast_sub(&max)?;
|
|
let sum_exp = diff.exp()?.sum_keepdim(d)?;
|
|
let log_sm = diff.broadcast_sub(&sum_exp.log()?)?;
|
|
Ok(log_sm)
|
|
}
|
|
|
|
// TODO: Once the index_select backprop is efficient enough, switch to using this.
|
|
fn _nll_loss(inp: &Tensor, target: &Tensor) -> candle::Result<Tensor> {
|
|
let b_sz = target.shape().r1()?;
|
|
inp.index_select(target, 0)?.sum_all()? / b_sz as f64
|
|
}
|
|
|
|
pub fn main() -> Result<()> {
|
|
let dev = candle::Device::cuda_if_available(0)?;
|
|
let m = candle_nn::vision::mnist::load_dir("data")?;
|
|
println!("train-images: {:?}", m.train_images.shape());
|
|
println!("train-labels: {:?}", m.train_labels.shape());
|
|
println!("test-images: {:?}", m.test_images.shape());
|
|
println!("test-labels: {:?}", m.test_labels.shape());
|
|
let train_labels = m.train_labels;
|
|
let train_images = m.train_images;
|
|
let train_labels = train_labels.to_vec1::<u8>()?;
|
|
let train_label_mask = train_labels
|
|
.iter()
|
|
.flat_map(|l| (0..LABELS).map(|i| f32::from(i == *l as usize)))
|
|
.collect::<Vec<_>>();
|
|
let train_label_mask = Tensor::from_vec(train_label_mask, (train_labels.len(), LABELS), &dev)?;
|
|
let ws = Var::zeros((IMAGE_DIM, LABELS), DType::F32, &dev)?;
|
|
let bs = Var::zeros(LABELS, DType::F32, &dev)?;
|
|
let sgd = candle_nn::SGD::new(&[&ws, &bs], 3e-1);
|
|
let test_images = m.test_images;
|
|
let test_labels = m.test_labels.to_vec1::<u8>()?;
|
|
for epoch in 1..200 {
|
|
let logits = train_images.matmul(&ws)?.broadcast_add(&bs)?;
|
|
let log_sm = log_softmax(&logits, D::Minus1)?;
|
|
let loss = (&log_sm * &train_label_mask)?
|
|
.sum_all()?
|
|
.affine(-1f64 / train_images.dim(0)? as f64, 0f64)?;
|
|
sgd.backward_step(&loss)?;
|
|
|
|
let test_logits = test_images.matmul(&ws)?.broadcast_add(&bs)?;
|
|
/* TODO: Add argmax so that the following can be computed within candle.
|
|
let test_accuracy = test_logits
|
|
.argmax(Some(-1), false)
|
|
.eq_tensor(&test_labels)
|
|
.to_kind(Kind::Float)
|
|
.mean(Kind::Float)
|
|
.double_value(&[]);
|
|
*/
|
|
let test_logits = test_logits.to_vec2::<f32>()?;
|
|
let sum_ok = test_logits
|
|
.iter()
|
|
.zip(test_labels.iter())
|
|
.map(|(logits, label)| {
|
|
let arg_max = logits
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|(_, v1), (_, v2)| v1.total_cmp(v2))
|
|
.map(|(idx, _)| idx);
|
|
f64::from(arg_max == Some(*label as usize))
|
|
})
|
|
.sum::<f64>();
|
|
let test_accuracy = sum_ok / test_labels.len() as f64;
|
|
println!(
|
|
"{epoch:4} train loss: {:8.5} test acc: {:5.2}%",
|
|
loss.to_scalar::<f32>()?,
|
|
100. * test_accuracy
|
|
)
|
|
}
|
|
Ok(())
|
|
}
|