Tfovid
Tfovid

Reputation: 843

How to create add columns (i.e., features) on a tf.Dataset?

QUESTION

Very often, one wants to enrich a raw dataset with derived features. I.e., new columns need to be created from preexisting ones. How does one do that in an efficient (and preferably in-place) way with a tf.Dataset?

PS: I tried to play around tf.data.Dataset.map(), tf.data.Dataset.apply(), and tf.map() but can't find the right syntax that does what I'm illustrating below.

MINIMUM WORKING EXAMPLE

To show what I want to do, I'll use pandas' apply(). For example, I'm trying to add a feature that is the length of the embark_town feature in the Titanic dataset.

import pandas as pd
import tensorflow as tf # v. 2.0+

# Load the Titanic dataset
source = tf.keras.utils.get_file(
    "train.csv", 
    "https://storage.googleapis.com/tf-datasets/titanic/train.csv")

# Only select two features and one target for this example.
dataset = tf.data.experimental.make_csv_dataset(
    source, batch_size=5, label_name="survived", 
    select_columns=["embark_town", "age", "survived"],
    num_epochs=1, ignore_errors=True, shuffle=False)

# Add the derived feature `embark_town_len` via pandas.
batch, _ = next(iter(dataset))
batch = pd.DataFrame(batch)

print("Raw data:")
print(batch)

batch['embark_town_len'] = batch.apply(lambda x: len(x["embark_town"]), axis=1)

print("\nEnriched data:")
print(batch)

which produces

Raw data:
    age     embark_town
0  22.0  b'Southampton'
1  38.0    b'Cherbourg'
2  26.0  b'Southampton'
3  35.0  b'Southampton'
4  28.0   b'Queenstown'

Enriched data:
    age     embark_town  embark_town_len
0  22.0  b'Southampton'               11
1  38.0    b'Cherbourg'                9
2  26.0  b'Southampton'               11
3  35.0  b'Southampton'               11
4  28.0   b'Queenstown'               10

Note that although I'm using pandas' apply() here, what I'm really looking for is something that works directly on the whole tf.Dataset, not just a batch therein.

Upvotes: 0

Views: 1771

Answers (1)

Frederik Bode
Frederik Bode

Reputation: 2744

Assuming tensorflow 2.0:

import tensorflow as tf
cities_ds = tf.data.Dataset.from_tensor_slices(["Rome","Brussels"])
ages_ds = tf.data.Dataset.from_tensor_slices([5,7])
ds = tf.data.Dataset.zip((cities_ds, ages_ds)) 
ds = ds.map(lambda city, age: (city, age, tf.strings.length(city)))
for i in ds:
  print(i[0].numpy(), i[1].numpy(), i[2].numpy())

Upvotes: 1

Related Questions