Un-scaled version of CoordConv2D

coordconv2d.py

import keras.backend as K
import tensorflow as tf
from tensorflow.keras.layers import Layer

"""Not tested, I'll play around with GANs soon with it."""
from tensorflow.keras.layers import Conv2D
import numpy as np


class CoordConv2D(Layer):
    def __init__(self, channel, kernel, padding='valid', **kwargs):
        self.layer = Conv2D(channel, kernel, padding=padding)
        super(CoordConv2D, self).__init__(**kwargs)

    def call(self, inputs, **kwargs):
        indices = tf.ones_like(inputs)
        if K.image_data_format() == 'channel_first':
            # bs, channel, w, h
            indices = indices[:, 0, ...]
        else:
            #bs, w, h, channel = input_shape
            indices = indices[..., 0]

        # Get indices
        bs, w, h = [tf.shape(indices)[k] for k in range(3)]
        indices = K.cast(tf.where(indices), tf.float32)
        canvas = K.reshape(indices, [bs, w, h, 3])[..., 1:]
        # Normalize the canvas
        canvas = canvas / tf.cast(K.reshape([w, h], [1, 1, 1, 2]), tf.float32)
        canvas = (canvas * 2) - 1

        # If channel_first, we swap
        if K.image_data_format() == 'channel_first':
            canvas = K.permute_dimensions(canvas, [0, 3, 1, 2])

        # Concatenate channel-wise
        input = K.concatenate([inputs, canvas], -1)
        return self.layer(input)

    def compute_output_shape(self, input_shape):
        return self.layer.compute_output_shape(input_shape)

class CustomModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.l = CoordConv2D(63, 3, padding='same')

    def call(self, inputs):
        x = self.l(inputs)
        return x

def main():
    mod = CustomModel()

    mod.compile('sgd', 'mse')
    mod.run_eagerly = True
    res = mod.predict(np.ones([3, 32, 32, 3]))
    print(res.shape)

if __name__ == '__main__':
    main()

@Dref360, I've tried modifying the above for Tensorflow 2.1. (see my fork). Any chance you understand what is causing the error that results when one runs it?

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: model/coord_conv2d/unstack:0

Hi Thomas,
I updated the gist to TF2.
Thank you!

Thanks!