Transform won't yield memory in tfx after transform and it takes up total memory

[axelning]

If the bug is related to a specific library below, please raise an issue in the
respective repo directly:

TensorFlow Data Validation Repo

TensorFlow Model Analysis Repo

TensorFlow Transform Repo

TensorFlow Serving Repo

System information

• Have I specified the code to reproduce the issue
  (Yes/No): yes
• Environment in which the code is executed (e.g., Local
  (Linux/MacOS/Windows), Interactive Notebook, Google Cloud, etc): - TensorFlow
  version (you are using): 2.3.2- TFX Version: 0.26.1- Python version:3.6.7

Describe the current behavior
In tfx transform module it calls tensorflow_transform> beam >impl.py:1058

schema = schema_inference.infer_feature_schema_v2(
      structured_outputs,
      metadata_fn.get_concrete_function(),
      evaluate_schema_overrides=False)

this will call infer_feature_schma_v2 in schema_inference.py :163

in this function, tf2_utils.supply_missing_inputs(structured_inputs, batch_size=1) in line 195 will tries to convert inputs to tensor and will not release the gpu memory when finished. By default this operation takes 7715 MB on my singlee Tesla p40

When I run into OOM because the following training starts to apply for the GPU, and after I stop the whole process and continue, cause the transform has been saved and the trainning goes successful, which means this part does not need to keep in the GPU from when it ends.

[arghyaganguly]

duplicate tfx#3343

[arghyaganguly]

@zoyahav , shall we track this issue here or in tfx#3343 ?

[zoyahav]

Let's keep it here for now.

@axelning are you able to check if the issue occurs with CPU as well?

[axelning]

Let's keep it here for now.

@axelning are you able to check if the issue occurs with CPU as well?

by setting the growth limitation and worker_num limitation， this issue can be circumvented
and in cpu just bcz i got 32GB memory， so this issue is not reproduced during running。

still，the memory management of gpu is keeping emerging，may be some architect optimization is needed

[pritamdodeja]

@axelning have you considered overriding the executor in the Transform component with a canonical tft based implementation?