huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 11:08:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retrieve data from the gpu.

Browse Source
This commit is contained in:
laurent
2023-06-22 11:01:49 +01:00
parent 083ced4428
commit 87a37b3bf3
2 changed files with 25 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
examples/cuda_basics.rs
View File

@ -1,12 +1,14 @@
use anyhow::Result; use anyhow::Result;
use candle::{Device, Tensor}; use candle::{DType, Device, Tensor};
fn main() -> Result<()> { fn main() -> Result<()> {
let device = Device::new_cuda(0)?; let device = Device::new_cuda(0)?;
let x = Tensor::new(&[3f32, 1., 4., 1., 5.], &device)?; let x = Tensor::new(&[3f32, 1., 4., 1., 5.], &device)?;
println!("{:?}", x.to_vec1::<f32>()?); println!("{:?}", x.to_vec1::<f32>()?);
let y = Tensor::new(&[2f32, 7., 1., 8., 2.], &device)?; let x = Tensor::new(&[2f32, 7., 1., 8., 2.], &device)?;
let z = (y * 3.)?; let y = (x * 3.)?;
println!("{:?}", z.to_vec1::<f32>()?); println!("{:?}", y.to_vec1::<f32>()?);
let x = Tensor::ones((3, 2), DType::F32, &device)?;
println!("{:?}", x.to_vec2::<f32>()?);
Ok(()) Ok(())
} }

36
src/tensor.rs
View File

@ -204,12 +204,13 @@ impl Tensor {
shape: self.shape().clone(), shape: self.shape().clone(),
}); });
} }
let from_cpu_storage = |cpu_storage: &crate::CpuStorage| {
let data = S::cpu_storage_as_slice(cpu_storage)?;
Ok::<_, Error>(data[0])
};
match &self.storage { match &self.storage {
Storage::Cpu(cpu_storage) => { Storage::Cpu(cpu_storage) => from_cpu_storage(cpu_storage),
let data = S::cpu_storage_as_slice(cpu_storage)?; Storage::Cuda(storage) => from_cpu_storage(&storage.to_cpu_storage()?),
Ok(data[0])
}
Storage::Cuda { .. } => todo!(),
} }
} }
@ -261,19 +262,20 @@ impl Tensor {
pub fn to_vec2<S: crate::WithDType>(&self) -> Result<Vec<Vec<S>>> { pub fn to_vec2<S: crate::WithDType>(&self) -> Result<Vec<Vec<S>>> {
let (dim1, dim2) = self.shape().r2()?; let (dim1, dim2) = self.shape().r2()?;
match &self.storage { let from_cpu_storage = |cpu_storage: &crate::CpuStorage| {
Storage::Cpu(cpu_storage) => { let data = S::cpu_storage_as_slice(cpu_storage)?;
let data = S::cpu_storage_as_slice(cpu_storage)?; let mut rows = vec![];
let mut rows = vec![]; let mut src_index = self.strided_index();
let mut src_index = self.strided_index(); for _idx_row in 0..dim1 {
for _idx_row in 0..dim1 { let row = (0..dim2).map(|_| data[src_index.next().unwrap()]).collect();
let row = (0..dim2).map(|_| data[src_index.next().unwrap()]).collect(); rows.push(row)
rows.push(row)
}
assert!(src_index.next().is_none());
Ok(rows)
} }
Storage::Cuda { .. } => todo!(), assert!(src_index.next().is_none());
Ok(rows)
};
match &self.storage {
Storage::Cpu(storage) => from_cpu_storage(storage),
Storage::Cuda(storage) => from_cpu_storage(&storage.to_cpu_storage()?),
} }
} }