Harrison Taylor
Harrison Taylor

Reputation: 23

How to apply element-wise rounding function to Keras Tensor?

I am trying to write a Lambda-style layer in Keras that quantizes each weight in the previous 1D Dense(output_len) layer to the nearest step of 1/128.

I have tried to use the map_tf function in the Keras backend, but I am having no luck so far.

Essentially what I'm trying to do is apply the following function element-wise to the 1D input tensor:

def quantize(x):
    'Squashes x (0->1) to steps of 1/128'
    precision = 3
    base = 0.0078125 # 1/128
    if x < 0:
        x = 0
    elif x > 1:
        x = 1

    return round(base * round(float(x)/base) - 1/256, precision)

So, for example, this would be the result of a certain prediction:

input (after going through the CNN):
  [0.21940812, 0.7998919 , 0.5420448 , 0.33850232 ]
output (after leaving the quantization layer):
  [0.215, 0.793, 0.535, 0.332 ]

Is what I'm trying to achieve possible?

Thanks.

Upvotes: 2

Views: 1512

Answers (1)

Peter Szoldan
Peter Szoldan

Reputation: 4868

This is how I would do it:

import tensorflow as tf
from keras.models import Model
from keras.layers import Input, Lambda
import keras.backend as K
import numpy as np

def quantize(x):
    'Squashes x (0->1) to steps of 1/128'
    precision = 3
    base = 0.0078125 # 1/128
    x = K.clip( x, min_value = 0.0, max_value = 1.0 )
    return K.round( 1000 * ( base * K.round( x / base ) - 1.0 / 256 ) ) / 1000

a = Input( shape = ( 4, ) )
b = Lambda( quantize )( a )
model = Model( inputs = a, outputs = b )
print ( model.predict( np.array( [ [0.21940812, 0.7998919 , 0.5420448 , 0.33850232 ] ] ) ) )

Outputs:

[[0.215 0.79300004 0.535 0.33200002]]

If you can live with the rounding errors...

Upvotes: 1

Related Questions