Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Un-scaled version of CoordConv2D
import keras.backend as K
import tensorflow as tf
from keras.layers import Layer
from keras.layers import Conv2D
"""Not tested, I'll play around with GANs soon with it."""
class CoordConv2D(Layer):
def __init__(self, channel, kernel, padding='valid', activation=None, **kwargs):
self.layer = Conv2D(channel, kernel, padding=padding, activation=activation)
self._name = 'CoordConv2D'
super(CoordConv2D, self).__init__(**kwargs)
def call(self, input):
input_shape = tf.unstack(K.shape(input))
if K.image_data_format() == 'channel_first':
bs, channel, w, h = input_shape
else:
bs, w, h, channel = input_shape
# Get indices
indices = tf.cast(tf.where(K.ones([bs, w, h])), dtype=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]), dtype=tf.float32)
canvas = (canvas * 2) - 1
# If channel_first, we swap
if K.image_data_format() == 'channel_first':
canvas = K.swap_axes(canvas, [0, 3, 1, 2])
# Concatenate channel-wise
input = K.concatenate([input, canvas], -1)
return self.layer(input)
def compute_output_shape(self, input_shape):
return self.layer.compute_output_shape(input_shape)
from keras.layers import Input, Conv2D
import numpy as np
from keras import Model
inp = Input([32, 32, 3])
layer = CoordConv2D(63, 3, padding='same')
x = layer(inp)
mod = Model(inp, x)
mod.compile('sgd', 'mse')
res = mod.predict(np.ones([3, 32, 32, 3]))
print(res.shape)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.