About initialization of Nearest Conv.

[dslisleedh]

Hi @Algolzw, Your works really inspired me. But I think there are points for improvement in initialization of Nearest Conv.

Current Nearest Conv's initialization is like this.

w = np.transpose(np.array([[
[[1, 0, 0], 
 [0, 1, 0], 
 [0, 0, 1]]*scale_factor**2
]]), (0, 1, 3, 2))

res_conv = Conv2D(out_c*(scale_factor**2), kernel_size=1, padding='same', use_bias=False)
res_conv.trainable = False
res_conv.build(x_in.shape)
res_conv.set_weights([w])
x_res = res_conv(x_in)

But this way can't initialize input channel other than 3. So I implemented initializer for Nearest Conv.

class NearestConvInit(tf.keras.initializers.Initializer):
    def __call__(self, shape, dtype=None):
        assert len(shape) == 4, f"NearestConv only support 2D conv, but got kernel shape of {len(shape) - 2}D Conv"
        assert shape[0] == shape[1] == 1, f'only support 1x1 kernel, but got kernel shape of {shape[0]}x{shape[1]}'
        assert shape[-1] % shape[-2] == 0, f'out_channel*scale**2 must be divisible by input channel, but got {shape[-1]} and {shape[-2]}'

        kernel = tf.transpose(
            tf.reshape(
                tf.repeat(tf.eye(shape[-2], dtype=dtype)[tf.newaxis, :], shape[-1]//shape[-2], axis=0), (1, 1, -1, shape[-2])
            ), perm=(0, 1, 3, 2)
        )
        return kernel

    def get_config(self):
        return super().get_config()

You can initialize Nearest Conv. simply by using the kernel_initialzier argument. Here's exmaple.

input_shape = (1, 256, 256, 3)
input_shape_build = (None, None, None, 3)

inputs = tf.nn.relu(tf.random.normal(input_shape))
scale_factor = 3

conv_init = tf.keras.layers.Conv2D(input_shape[-1]*(scale_factor**2), 1, padding='SAME', kernel_initializer=NearestConvInit(), use_bias=False)
conv_init.build(input_shape_build)
print('Conv_init kernel:', conv_init.kernel.shape)
output_init = conv_init(inputs)
print('Conv_init output shape: ', output_init.shape)

conv = tf.keras.layers.Conv2D(input_shape[-1]*(scale_factor**2), 1, padding='SAME', use_bias=False)
w = np.transpose(np.array([[
    [[1, 0, 0], 
     [0, 1, 0], 
     [0, 0, 1]]*scale_factor**2
    ]]), (0, 1, 3, 2))
conv.build(input_shape_build)
conv.set_weights([w])
print('\nConv kernel:', conv.kernel.shape)
output = conv(inputs)
print('Conv output shape: ', output.shape)

print('\nDifference between init and non-init: ', tf.reduce_sum(tf.abs(output_init - output)))
----------------------------------------------------------------------------------------------------
Conv_init kernel: (1, 1, 3, 27)
Conv_init output shape:  (1, 256, 256, 27)

Conv kernel: (1, 1, 3, 27)
Conv output shape:  (1, 256, 256, 27)

Difference between init and non-init:  tf.Tensor(0.0, shape=(), dtype=float32)

You can see there is no difference in the results.
I tested this initialzier with edited basenet func and variable input size.

def basenet_init(n_feat=64, inp_c=3, out_c=3, scale_factor=3):
    x_in = Input(shape=(None, None, inp_c))

    res_conv = Conv2D(out_c*(scale_factor**2), kernel_size=1, padding='same', use_bias=False, kernel_initializer=NearestConvInit())
    res_conv.trainable = False
    x_res = res_conv(x_in)
    
    x = Conv2D(n_feat, kernel_size=3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x_in)
    x = Conv2D(n_feat, kernel_size=3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    x = Conv2D(n_feat, kernel_size=3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    x = Conv2D(n_feat, kernel_size=3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    x = Conv2D(n_feat, kernel_size=3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    x = Conv2D(out_c*(scale_factor**2), kernel_size=3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    x = Conv2D(out_c*(scale_factor**2), kernel_size=3, padding='same', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    x = Add()([x_res, x])
    depth_to_space = Lambda(lambda x: tf.nn.depth_to_space(x, scale_factor))
    x = depth_to_space(x)

    clip_func = Lambda(lambda x: tf.clip_by_value(x, 0., 255.))
    x = clip_func(x)

    return Model(x_in, x)


for i in range(1, 4):
    inp_c = i
    inputs = tf.nn.relu(tf.random.normal((1, 256, 256, inp_c)))
    try:
        print('\ninput_shape:', inputs.shape)
        init_out = basenet_init(inp_c=inp_c,out_c=3)(inputs)
        print('Output shape: ', init_out.shape)
    except Exception as e:
        print('Error: ', e)
-----------------------------------------------------------
input_shape: (1, 256, 256, 1)
Output shape:  (1, 768, 768, 3)

input_shape: (1, 256, 256, 2)
Error:  out_channel*scale**2 must be divisible by input channel, but got 27 and 2

input_shape: (1, 256, 256, 3)
Output shape:  (1, 768, 768, 3)

And it works perpectly when input_channel == output_channel

for i in range(1, 4):
    inp_c = i
    inputs = tf.nn.relu(tf.random.normal((1, 256, 256, inp_c)))
    try:
        print('\ninput_shape:', inputs.shape)
        init_out = basenet_init(inp_c=inp_c,out_c=inp_c)(inputs)
        print('Output shape: ', init_out.shape)
    except Exception as e:
        print('Error: ', e)
----------------------------------------------
input_shape: (1, 256, 256, 1)
Output shape:  (1, 768, 768, 1)

input_shape: (1, 256, 256, 2)
Output shape:  (1, 768, 768, 2)

input_shape: (1, 256, 256, 3)
Output shape:  (1, 768, 768, 3)

I hope this helps.
Thank you.

[Algolzw]

Hi @dslisleedh, Thank you and I agree that your implementation seems more flexible!
I will add a link to this issue on the readme page and thanks again for your efforts in this work.

Best

[dslisleedh]

Thank you :)