将 pytorch 2d 填充转换为 tensorflow keras

Question

What would be the equivalent of this:

nn.ReflectionPad2d(1)

in TensorFlow 2? The above line is from PyTorch.

In TensorFlow 2 Keras, I'm currently looking into using tf.pad() as a TF version of that, but it seems PyTorch is able to handle varying dimensions with that single integer 1. For instance, if it gets an input of shape [batch size, 1, 1, 100], nn.ReflectionPad2D will handle that well, but in TensorFlow, I get an error if I try to use

tf.pad(t, tf.constant([0,0], [1,1], [1,1], [0,0]]), 'REFLECT')

Any suggestions on how to replicate nn.ReflectinPad2d in TensorFlow 2 keras? Thanks!

Answer 1

When I was training CycleGan on TF2, I created this custom layer for myself:

class ReflectionPad2D(tf.keras.layers.Layer):
  def __init__(self, paddings=(1,1,1,1)):
    super(ReflectionPad2D, self).__init__()
    self.paddings = paddings

  def call(self, input):
    l, r, t, b = self.paddings

    return tf.pad(input, paddings=[[0,0], [t,b], [l,r], [0,0]], mode='REFLECT')

You can just use it as is by putting it in a Model/Sequential, eg:

model = Sequential([
    ReflectionPad2D((3, 3, 3, 3)),
    Conv2D(64, kernel_size=7, strides=1, padding='valid', use_bias=False),
    BatchNormalization()
])

model.build(input_shape=(8, 128, 128, 3))
model.summary()

Example output:

Model: "sequential_13"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
reflection_pad2d_198 (Reflec (8, 134, 134, 3)          0         
_________________________________________________________________
conv2d_474 (Conv2D)          (8, 128, 128, 64)         9408      
_________________________________________________________________
batch_normalization_41 (Batc (8, 128, 128, 64)         256       
=================================================================
Total params: 9,664
Trainable params: 9,536
Non-trainable params: 128
_________________________________________________________________