Support dilation in conv-transpose2d. (#671)

This commit is contained in:
Laurent Mazare
2023-08-30 09:22:00 +01:00
committed by GitHub
parent 9b25113393
commit 393690387f
4 changed files with 40 additions and 17 deletions

View File

@ -1186,12 +1186,6 @@ impl<'a> Map2 for ConvTranspose2D<'a> {
const OP: &'static str = "conv_transpose2d";
fn f<T: WithDType>(&self, inp: &[T], inp_l: &Layout, k: &[T], k_l: &Layout) -> Result<Vec<T>> {
let p = self.0;
if p.dilation != 1 {
crate::bail!(
"dilation {} is not supported for conv-transpose2d",
p.dilation
)
}
let inp = &inp[inp_l.start_offset()..];
let (inp_s0, inp_s1, inp_s2, inp_s3) = crate::shape::dims4(inp_l.stride())?;
let k = &k[k_l.start_offset()..];
@ -1235,8 +1229,8 @@ impl<'a> Map2 for ConvTranspose2D<'a> {
for b_idx in 0..p.b_size {
for inp_y in 0..p.i_h {
for inp_x in 0..p.i_w {
let out_x = inp_x * p.stride + k_x;
let out_y = inp_y * p.stride + k_y;
let out_x = inp_x * p.stride + k_x * p.dilation;
let out_y = inp_y * p.stride + k_y * p.dilation;
if out_x < p.padding || out_y < p.padding {
continue;
}