Use all the CPU available on a single node for `generate_statistics_from_tfrecord`

[zippeurfou]

Hi, I've been trying to use all the CPU available for my single node to compute the statistics. The default behavior of tfdv.generate_statistics_from_tfrecord only use a single CPU which is really slow. One solution I tried which does use all the CPU is this one:

train_stats=tfdv.generate_statistics_from_tfrecord(train_paths,stats_options=tfdv.StatsOptions(enable_semantic_domain_stats=True),pipeline_options=PipelineOptions(['--runner','DirectRunner','--direct_num_workers', '0', '--direct_running_mode', 'multi_processing']))

However, it somehow silently ignore a subset of my data. For example I have 12 M records and it will display 8M. Meanwhile without the pipeline options it does display the 12M. I do not have access to gcp as we are on aws. Is there anything I am missing here? PS: I also tried to use StatisticsGen component but get similar issues.

[singhniraj08]

@zippeurfou,

If you can load the whole data into memory, you can use tfdv.generate_statistics_from_dataframe with n_jobs argument set to -1 to utilize all the CPU cores.

Can you share us an example graphs where it ignores the subset of data? I am unable to replicate the issue with below code. Also, make sure you are using Apache beam >= 2.22.0 as direct_num_workers = 0 support was introduced in 2.22.0.

stats = tfdv.generate_statistics_from_tfrecord(data_location=train_tf_file,
                                              #  stats_options=tfdv.StatsOptions(enable_semantic_domain_stats=True),
                                               pipeline_options=PipelineOptions(['--runner','DirectRunner','--direct_num_workers', '0', '--direct_running_mode', 'multi_processing']),
                                               output_path=new_stats_location,
                                               )

Thank you!

[zippeurfou]

Thank you for the reply. I would like to avoid using the data frame one. I will share an example when I can. I think the main difference most likely is that I am working with a folders that has multiple tfrecord files compressed.

[singhniraj08]

@zippeurfou,

generate_statistics_from_tfrecord by default will pick multiple tfrecord files from given data folder(location). Make sure the given path have all the tfrecord files and avoid sub folders. Please share us the example code to replicate this issue on our end. Thanks.

[singhniraj08]

Closing this due to inactivity. Please take a look into the answers provided above, feel free to reopen and post your comments(if you still have queries on this). Thank you!