How to implement a set lookup in Tensorflow?

user_206253 0
Viewed 104

During the preprocessing of a tensorflow dataset I need to check whether a certain value is contained in an unmutable set. If it isn't I need to replace it with a default value. Essentially it is about censoring/replacing certain outliers

In python I would do something like this:

def map_id (value):
  s = frozenset([1,2,3])
  if value in s:
    return value
  else:
    return 0 # default for all outliers

This map_id function will be called like this

def preprocess(item):
  return (map_id(item["investment_id"]), item["features"]), item["target"]

The preprocess function will be called like this

def make_dataset(file_paths, batch_size=4096, mode="train"):
  ds = tf.data.TFRecordDataset(file_paths)
  ds = ds.map(decode_function)
  ds = ds.map(preprocess)
  if mode == "train":
    ds = ds.shuffle(batch_size * 4)
  ds = ds.batch(batch_size).cache().prefetch(tf.data.AUTOTUNE)
  return ds

How to write this map_id function in Tensorflow 2.x?

tensorflow-datasets tensorflow2.0 tensorflow set python
1 Answers

I am not sure what your data looks like, but you should be able to use a simple StaticHashTable as a Set alternative for your use case, since it will run in graph mode:

import tensorflow as tf

data = {
    "investment_id": [1, 2, 3, 4, 5], 
    "features": [12, 912, 28, 90, 17],
    "target": [1, 0, 1, 1, 1]
}

keys_tensor = tf.constant([1, 2, 3])
vals_tensor = tf.constant([1, 2, 3])
table = tf.lookup.StaticHashTable(
    tf.lookup.KeyValueTensorInitializer(keys_tensor, vals_tensor),
    default_value=0)


ds = tf.data.Dataset.from_tensor_slices(data)
ds = ds.map(lambda item: (table.lookup(item['investment_id']), item['features'], item['target']))

for d in ds:
  print(d)

(<tf.Tensor: shape=(), dtype=int32, numpy=1>, <tf.Tensor: shape=(), dtype=int32, numpy=12>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=int32, numpy=2>, <tf.Tensor: shape=(), dtype=int32, numpy=912>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)
(<tf.Tensor: shape=(), dtype=int32, numpy=3>, <tf.Tensor: shape=(), dtype=int32, numpy=28>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=int32, numpy=0>, <tf.Tensor: shape=(), dtype=int32, numpy=90>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=int32, numpy=0>, <tf.Tensor: shape=(), dtype=int32, numpy=17>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
user_96578610
Related