coreylowman / dfdx

Deep learning in Rust, with shape checked tensors and neural networks

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Add `OUTPUT_PADDING` to `ConvTrans2D`

swfsql opened this issue · comments

This issue is a request to add OUTPUT_PADDING to Conv2DTranspose.

It appears that dfdx ConvTrans2D behaves as if tensorflow Conv2DTranspose has output_padding=0.
Related code for dfdx and tf.

Code example for tf:

import tensorflow as tf
import numpy as np

x = np.zeros([1, 1, 2, 3], dtype=np.float32)
print(x.shape) # (1, 1, 2, 3)

a = tf.keras.layers.Conv2DTranspose(output_padding=0, filters=1, kernel_size=3, strides=2, padding='same', data_format='channels_first')
b = tf.keras.layers.Conv2DTranspose(output_padding=1, filters=1, kernel_size=3, strides=2, padding='same', data_format='channels_first')

ya = a(x).numpy().shape
yb = b(x).numpy().shape
print(ya) # (1, 1, 3, 5)
print(yb) # (1, 1, 4, 6)

Code example for dfdx:

use dfdx::prelude::*;

const IN_CHAN: usize = 1;
const OUT_CHAN: usize = 1;
const KERNEL_SIZE: usize = 3;
const STRIDE: usize = 2;
// for padding='same'
const PADDING: usize = ((KERNEL_SIZE - 1) * DILATION + 1) / 2; // = 1
const DILATION: usize = 1;
const GROUPS: usize = 1;
type Model = ConvTrans2DConstConfig<
    IN_CHAN,
    OUT_CHAN,
    KERNEL_SIZE,
    STRIDE,
    PADDING,
    DILATION,
    GROUPS,
>;

fn example() {
    let dev = Cpu::default();
    let model = dev.build_module::<f32>(Model::default());
    let x: Tensor<Rank4<1, 1, 2, 3>, _, _> = dev.zeros();
    let _prediction: Tensor<Rank4<1, 1, 3, 5>, _, _> = model.forward(x);
    // note: the shape is the same for `ya` from the tf example.
}