From a424d95473ea9268ffb1dde4d73ce0cff9904845 Mon Sep 17 00:00:00 2001 From: laurent Date: Tue, 4 Jul 2023 11:15:45 +0100 Subject: [PATCH] Add more of the conv1d op. --- candle-core/src/conv.rs | 24 +++++++++++++++++++++ candle-core/src/cpu_backend.rs | 3 +-- candle-core/src/cuda_backend.rs | 3 +-- candle-core/src/dummy_cuda_backend.rs | 3 +-- candle-core/src/lib.rs | 1 + candle-core/src/storage.rs | 7 +++--- candle-core/src/tensor.rs | 27 ++++++++++++++++++------ candle-examples/examples/whisper/main.rs | 3 +-- 8 files changed, 52 insertions(+), 19 deletions(-) create mode 100644 candle-core/src/conv.rs diff --git a/candle-core/src/conv.rs b/candle-core/src/conv.rs new file mode 100644 index 00000000..90bb5229 --- /dev/null +++ b/candle-core/src/conv.rs @@ -0,0 +1,24 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ParamsConv1D { + pub(crate) b_size: Option, + pub(crate) c_out: usize, + pub(crate) c_in: usize, + pub(crate) k_size: usize, + pub(crate) padding: usize, + pub(crate) stride: usize, +} + +impl ParamsConv1D { + pub(crate) fn l_out(&self, l_in: usize) -> usize { + let dilation = 1; + (l_in + 2 * self.padding - dilation * (self.k_size - 1) - 1) / self.stride + 1 + } + + pub(crate) fn out_dims(&self, l_in: usize) -> Vec { + let l_out = self.l_out(l_in); + match self.b_size { + None => vec![self.c_out, l_out], + Some(n) => vec![n, self.c_out, l_out], + } + } +} diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs index ed3a5998..54002184 100644 --- a/candle-core/src/cpu_backend.rs +++ b/candle-core/src/cpu_backend.rs @@ -632,8 +632,7 @@ impl CpuStorage { _l: &Layout, _kernel: &Self, _kernel_l: &Layout, - _padding: usize, - _stride: usize, + _params: &crate::conv::ParamsConv1D, ) -> Result { todo!() } diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs index ec69688c..917655fc 100644 --- a/candle-core/src/cuda_backend.rs +++ b/candle-core/src/cuda_backend.rs @@ -806,8 +806,7 @@ impl CudaStorage { _l: &Layout, _kernel: &Self, _kernel_l: &Layout, - _padding: usize, - _stride: usize, + _params: &crate::conv::ParamsConv1D, ) -> Result { todo!() } diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs index eca5961b..0dbd8d54 100644 --- a/candle-core/src/dummy_cuda_backend.rs +++ b/candle-core/src/dummy_cuda_backend.rs @@ -105,8 +105,7 @@ impl CudaStorage { _l: &Layout, _kernel: &Self, _kernel_l: &Layout, - _padding: usize, - _stride: usize, + _params: &crate::conv::ParamsConv1D, ) -> Result { Err(Error::NotCompiledWithCudaSupport) } diff --git a/candle-core/src/lib.rs b/candle-core/src/lib.rs index 0d4c2a8d..2365a34d 100644 --- a/candle-core/src/lib.rs +++ b/candle-core/src/lib.rs @@ -1,4 +1,5 @@ mod backprop; +mod conv; mod cpu_backend; #[cfg(feature = "cuda")] mod cuda_backend; diff --git a/candle-core/src/storage.rs b/candle-core/src/storage.rs index 235080c0..53ea1544 100644 --- a/candle-core/src/storage.rs +++ b/candle-core/src/storage.rs @@ -149,18 +149,17 @@ impl Storage { l: &Layout, kernel: &Self, kernel_l: &Layout, - padding: usize, - stride: usize, + params: &crate::conv::ParamsConv1D, ) -> Result { self.same_device(kernel, "conv1d")?; self.same_dtype(kernel, "conv1d")?; match (self, &kernel) { (Storage::Cpu(inp), Storage::Cpu(kernel)) => { - let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?; + let s = inp.conv1d(l, kernel, kernel_l, params)?; Ok(Self::Cpu(s)) } (Storage::Cuda(inp), Storage::Cuda(kernel)) => { - let s = inp.conv1d(l, kernel, kernel_l, padding, stride)?; + let s = inp.conv1d(l, kernel, kernel_l, params)?; Ok(Self::Cuda(s)) } (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp { diff --git a/candle-core/src/tensor.rs b/candle-core/src/tensor.rs index 26d44718..590b81c4 100644 --- a/candle-core/src/tensor.rs +++ b/candle-core/src/tensor.rs @@ -433,13 +433,26 @@ impl Tensor { } pub fn conv1d(&self, kernel: &Self, padding: usize, stride: usize) -> Result { - let storage = self.storage.conv1d( - self.layout(), - &kernel.storage, - kernel.layout(), + let (c_out, c_in_k, k_size) = kernel.shape().r3()?; + let (b_size, c_in, l_in) = match *self.dims() { + [b_size, c_in, l_in] => (Some(b_size), c_in, l_in), + [c_in, l_in] => (None, c_in, l_in), + _ => todo!("proper error message"), + }; + if c_in != c_in_k { + todo!("proper error message") + } + let params = crate::conv::ParamsConv1D { + b_size, + c_out, + c_in, + k_size, padding, stride, - )?; + }; + let storage = + self.storage + .conv1d(self.layout(), &kernel.storage, kernel.layout(), ¶ms)?; let op = if self.track_op() || kernel.track_op() { Some(Op::Conv1D { arg: self.clone(), @@ -450,8 +463,8 @@ impl Tensor { } else { None }; - let dims = self.dims(); - Ok(from_storage(storage, dims, op, false)) + let out_dims = params.out_dims(l_in); + Ok(from_storage(storage, out_dims, op, false)) } pub fn matmul(&self, rhs: &Self) -> Result { diff --git a/candle-examples/examples/whisper/main.rs b/candle-examples/examples/whisper/main.rs index 75ab2189..a380d30e 100644 --- a/candle-examples/examples/whisper/main.rs +++ b/candle-examples/examples/whisper/main.rs @@ -236,8 +236,7 @@ impl Conv1D { fn forward(&self, x: &Tensor) -> candle::Result { let (bsize, _, _) = x.shape().r3()?; let w = self.weight.broadcast_left(bsize)?.t()?; - // TODO: Add the conv1d operation - let x = x.matmul(&w)?; + let x = x.conv1d(&w, self.config.padding, self.config.stride)?; match &self.bias { None => Ok(x), Some(bias) => x.broadcast_add(bias),