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

Abstract the gradient storage.

Browse Source
This commit is contained in:
laurent
2023-06-21 14:29:48 +01:00
parent 68f525f321
commit 7adffafeda
5 changed files with 87 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
Cargo.toml
View File

@ -13,9 +13,14 @@ readme = "README.md"
[dependencies]
safetensors = "0.3.1"
thiserror = "1"
cudarc = { version = "0.9.9", optional = true }
[dev-dependencies]
anyhow = "1"
clap = { version = "4.2.4", features = ["derive"] }
rand = "0.8.5"
tokenizers = "0.13.3"
[features]
default = []
cuda = ["dep:cudarc"]

21
src/device.rs
View File

@ -54,27 +54,36 @@ impl<S: crate::WithDType, const N: usize, const M: usize> NdArray for &[[S; N];
}
impl Device {
pub(crate) fn ones(&self, shape: &Shape, dtype: DType) -> Storage {
pub(crate) fn ones(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
match self {
Device::Cpu => Storage::Cpu(CpuStorage::ones_impl(shape, dtype)),
Device::Cpu => {
let storage = Storage::Cpu(CpuStorage::ones_impl(shape, dtype));
Ok(storage)
}
Device::Cuda { gpu_id: _ } => {
todo!()
}
}
}
pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Storage {
pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
match self {
Device::Cpu => Storage::Cpu(CpuStorage::zeros_impl(shape, dtype)),
Device::Cpu => {
let storage = Storage::Cpu(CpuStorage::zeros_impl(shape, dtype));
Ok(storage)
}
Device::Cuda { gpu_id: _ } => {
todo!()
}
}
}
pub(crate) fn tensor<A: NdArray>(&self, array: A) -> Storage {
pub(crate) fn tensor<A: NdArray>(&self, array: A) -> Result<Storage> {
match self {
Device::Cpu => Storage::Cpu(array.to_cpu_storage()),
Device::Cpu => {
let storage = Storage::Cpu(array.to_cpu_storage());
Ok(storage)
}
Device::Cuda { gpu_id: _ } => {
todo!()
}

94
src/tensor.rs
View File

@ -86,9 +86,9 @@ impl Tensor {
dtype: DType,
device: Device,
is_variable: bool,
) -> Self {
) -> Result<Self> {
let shape = shape.into();
let storage = device.ones(&shape, dtype);
let storage = device.ones(&shape, dtype)?;
let stride = shape.stride_contiguous();
let tensor_ = Tensor_ {
id: TensorId::new(),
@ -98,18 +98,18 @@ impl Tensor {
op: None,
is_variable,
};
Self(Arc::new(tensor_))
Ok(Self(Arc::new(tensor_)))
}
pub fn ones<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Self {
pub fn ones<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
Self::ones_impl(shape, dtype, device, false)
}
pub fn ones_var<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Self {
pub fn ones_var<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
Self::ones_impl(shape, dtype, device, true)
}
pub fn ones_like(&self) -> Self {
pub fn ones_like(&self) -> Result<Self> {
Tensor::ones(self.shape(), self.dtype(), self.device())
}
@ -118,9 +118,9 @@ impl Tensor {
dtype: DType,
device: Device,
is_variable: bool,
) -> Self {
) -> Result<Self> {
let shape = shape.into();
let storage = device.zeros(&shape, dtype);
let storage = device.zeros(&shape, dtype)?;
let stride = shape.stride_contiguous();
let tensor_ = Tensor_ {
id: TensorId::new(),
@ -130,18 +130,18 @@ impl Tensor {
op: None,
is_variable,
};
Self(Arc::new(tensor_))
Ok(Self(Arc::new(tensor_)))
}
pub fn zeros<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Self {
pub fn zeros<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
Self::zeros_impl(shape, dtype, device, false)
}
pub fn zeros_var<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Self {
pub fn zeros_var<S: Into<Shape>>(shape: S, dtype: DType, device: Device) -> Result<Self> {
Self::zeros_impl(shape, dtype, device, true)
}
pub fn zeros_like(&self) -> Self {
pub fn zeros_like(&self) -> Result<Self> {
Tensor::zeros(self.shape(), self.dtype(), self.device())
}
@ -151,7 +151,7 @@ impl Tensor {
is_variable: bool,
) -> Result<Self> {
let shape = array.shape()?;
let storage = device.tensor(array);
let storage = device.tensor(array)?;
let stride = shape.stride_contiguous();
let tensor_ = Tensor_ {
id: TensorId::new(),
@ -376,16 +376,16 @@ impl Tensor {
nodes
}
pub fn backward(&self) -> Result<HashMap<TensorId, Tensor>> {
pub fn backward(&self) -> Result<GradStore> {
let sorted_nodes = self.sorted_nodes();
println!("{}", sorted_nodes.len());
let mut grads = HashMap::new();
grads.insert(self.id, self.ones_like());
let mut grads = GradStore::new();
grads.insert(self, self.ones_like()?);
for node in sorted_nodes.iter() {
if node.is_variable {
continue;
}
let grad = grads.remove(&node.id).unwrap();
let grad = grads.remove(node).unwrap();
// TODO: We should perform all these operations in place (or at least not track the
// whole graph).
// The only drawback would be if we wanted to support grad of grad but this is out of
@ -393,51 +393,51 @@ impl Tensor {
if let Some(op) = &node.op {
match op {
Op::Add(lhs, rhs) => {
let lhs_sum_grad = grads.entry(lhs.id).or_insert_with(|| lhs.zeros_like());
let lhs_sum_grad = grads.or_insert(lhs)?;
*lhs_sum_grad = lhs_sum_grad.add(&grad)?;
let rhs_sum_grad = grads.entry(rhs.id).or_insert_with(|| rhs.zeros_like());
let rhs_sum_grad = grads.or_insert(rhs)?;
*rhs_sum_grad = rhs_sum_grad.add(&grad)?;
}
Op::Sub(lhs, rhs) => {
let lhs_sum_grad = grads.entry(lhs.id).or_insert_with(|| lhs.zeros_like());
let lhs_sum_grad = grads.or_insert(lhs)?;
*lhs_sum_grad = lhs_sum_grad.add(&grad)?;
let rhs_sum_grad = grads.entry(rhs.id).or_insert_with(|| rhs.zeros_like());
let rhs_sum_grad = grads.or_insert(rhs)?;
*rhs_sum_grad = rhs_sum_grad.add(&grad.neg()?)?;
}
Op::Mul(lhs, rhs) => {
let lhs_grad = grad.mul(rhs)?;
let lhs_sum_grad = grads.entry(lhs.id).or_insert_with(|| lhs.zeros_like());
let lhs_sum_grad = grads.or_insert(lhs)?;
*lhs_sum_grad = lhs_sum_grad.add(&lhs_grad)?;
let rhs_grad = grad.mul(lhs)?;
let rhs_sum_grad = grads.entry(rhs.id).or_insert_with(|| rhs.zeros_like());
let rhs_sum_grad = grads.or_insert(rhs)?;
*rhs_sum_grad = rhs_sum_grad.add(&rhs_grad)?;
}
Op::Div(lhs, rhs) => {
let lhs_grad = grad.div(rhs)?;
let lhs_sum_grad = grads.entry(lhs.id).or_insert_with(|| lhs.zeros_like());
let lhs_sum_grad = grads.or_insert(lhs)?;
*lhs_sum_grad = lhs_sum_grad.add(&lhs_grad)?;
let rhs_grad = grad.mul(lhs)?.div(&rhs.sqr()?)?;
let rhs_sum_grad = grads.entry(rhs.id).or_insert_with(|| rhs.zeros_like());
let rhs_sum_grad = grads.or_insert(rhs)?;
*rhs_sum_grad = rhs_sum_grad.add(&rhs_grad)?;
}
Op::Affine { arg, mul, .. } => {
let arg_grad = grad.affine(*mul, 0.)?;
let sum_grad = grads.entry(arg.id).or_insert_with(|| arg.zeros_like());
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?
}
Op::Neg(arg) => {
let arg_grad = grad.neg()?;
let sum_grad = grads.entry(arg.id).or_insert_with(|| arg.zeros_like());
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?
}
Op::Sqr(arg) => {
let arg_grad = arg.mul(&grad)?.affine(2., 0.)?;
let sum_grad = grads.entry(arg.id).or_insert_with(|| arg.zeros_like());
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?
}
Op::Sqrt(arg) => {
let arg_grad = grad.div(arg)?.affine(0.5, 0.)?;
let sum_grad = grads.entry(arg.id).or_insert_with(|| arg.zeros_like());
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?
}
};
@ -503,3 +503,39 @@ bin_trait!(Add, add, |_| 1., |v| v);
bin_trait!(Sub, sub, |_| 1., |v: f64| -v);
bin_trait!(Mul, mul, |v| v, |_| 0.);
bin_trait!(Div, div, |v| 1. / v, |_| 0.);
pub struct GradStore(HashMap<TensorId, Tensor>);
impl GradStore {
fn new() -> Self {
GradStore(HashMap::new())
}
pub fn get_id(&self, id: TensorId) -> Option<&Tensor> {
self.0.get(&id)
}
pub fn get(&self, tensor: &Tensor) -> Option<&Tensor> {
self.0.get(&tensor.id)
}
pub fn remove(&mut self, tensor: &Tensor) -> Option<Tensor> {
self.0.remove(&tensor.id)
}
pub fn insert(&mut self, tensor: &Tensor, grad: Tensor) -> Option<Tensor> {
self.0.insert(tensor.id, grad)
}
fn or_insert(&mut self, tensor: &Tensor) -> Result<&mut Tensor> {
use std::collections::hash_map::Entry;
let grad = match self.0.entry(tensor.id) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => {
let grad = tensor.zeros_like()?;
entry.insert(grad)
}
};
Ok(grad)
}
}

2
tests/grad_tests.rs
View File

@ -6,7 +6,7 @@ fn simple_grad() -> Result<()> {
let x = Tensor::var(&[3f32, 1., 4.], Device::Cpu)?;
let y = (((&x * &x)? + &x * 5f64)? + 4f64)?;
let grads = y.backward()?;
let grad_x = grads.get(&x.id()).context("no grad for x")?;
let grad_x = grads.get(&x).context("no grad for x")?;
assert_eq!(x.to_vec1::<f32>()?, [3., 1., 4.]);
// y = x^2 + 5.x + 4
assert_eq!(y.to_vec1::<f32>()?, [28., 10., 40.]);

2
tests/tensor_tests.rs
View File

@ -2,7 +2,7 @@ use candle::{DType, Device, Result, Tensor};
#[test]
fn zeros() -> Result<()> {
let tensor = Tensor::zeros((5, 2), DType::F32, Device::Cpu);
let tensor = Tensor::zeros((5, 2), DType::F32, Device::Cpu)?;
let (dim1, dim2) = tensor.shape().r2()?;
assert_eq!(dim1, 5);
assert_eq!(dim2, 2);