memory leak when using tft.vocabulary with labels argument

[AnnKatrinBecker]

Ubuntu 18.04 tensorflow 2.2.0 tfx 0.21.4

I am generating a vocabulary from my own dataset containing 28GB TFRecords with short description strings (up to 20 words) and integer labels from 1-100.

Generating the vocabulary without labels works fine. But as soon as the labels argument is provided to tft.vocabulary memory usage increases dramatically >100GB until the processes gets killed due to running out of memory.

`def preprocessing_fn(inputs): label = inputs['label'] desc = inputs['description'] desc = tf.strings.lower(desc.values)

remove all numbers and punctuations

desc = tf.strings.regex_replace(desc, "[^a-zA-Z¿]+", " ")
tokens = tf.strings.split(desc)
ngrams = tf.strings.ngrams(tokens, [1,2])
ngrams = ngrams.to_sparse()

tft.vocabulary(ngrams, top_k=100000, labels=tf.sparse.to_dense(label),
               vocab_filename='ngrams_100k_labels')
return {'description': desc, 'label': label}`

`def main():

Brings data into the pipeline

examples = external_input(
    'directory with tfrecords')
example_gen = ImportExampleGen(input=examples)
examples = example_gen.outputs['examples']

# Import schema
schema_importer = ImporterNode(
    instance_name='imported_schema',
    source_uri='pipelines/test/SchemaGen/schema/3',
    artifact_type=Schema)

### Perform transformation
transform = Transform(examples=example_gen.outputs['examples'],
                      schema=schema_importer.outputs['result'],
                      module_file='preprocessing.py')

pipe = pipeline.Pipeline( pipeline_name='test', pipeline_root='pipelines/test', components=[ example_gen, schema_importer, transform, ], metadata_connection_config=metadata.sqlite_metadata_connection_config( 'test/metadata.db'), enable_cache=True, beam_pipeline_args=['--direct_num_workers=0'])

absl.logging.set_verbosity(absl.logging.INFO)
BeamDagRunner().run(pipe)`

[mrcolo]

Any update on this? I'm getting a similar problem import a 38GB dataset using ImportExampleGen.

[arghyaganguly]

For large scale datasets it is recommended to use DataflowRunner (on GCP) or FlinkRunner or SparkRunner if it is a on-premises execution.