Discretize only a certain arrrays in a tensor with TensorFlow

Issue

I have the following array:-

import numpy as np
import tensorflow as tf
input = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])
layer = tf.keras.layers.Discretization(num_bins=2, epsilon=0.01)
layer.adapt(input)
layer(input)

<tf.Tensor: shape=(2, 4), dtype=int64, numpy=
array([[0, 1, 1, 1],
       [0, 1, 1, 0]])>

This discretizes the whole tensor. I would like to know if there is a way through which I can just discretize the second array in the tensor.

Solution

We can create a mask based on the index of the array that needs to be discretized:

def get_mask(x, array_index): 
    x = tf.Variable(tf.ones_like(input, dtype=tf.float32))
    indices = tf.Variable(array_index, dtype=tf.int32)
    updates = tf.Variable(tf.zeros( (indices.shape[0], x.shape[1])), dtype=tf.float32)
    return tf.compat.v1.scatter_nd_update(x, indices, updates)

And calling

> mask = get_mask(input, np.array([[1]])) #second array
> 
> returns the mask of:
array([[1., 1., 1., 1.],
       [0., 0., 0., 0.]])

Then we can apply mask: tf.cast(layer(input), tf.float32) * (1-mask) + input*mask which returns:

array([[-1.5,  1. ,  3.4,  0.5],
       [ 0. ,  1. ,  1. ,  0. ]]

The above should work for any array and any array index to discretize.

Answered By – vijayachandran mariappan

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published