TensorFlow 2.2.0 update for the tensorflow estimator for Azure ML, or disable horovod?

Question

I would like to use the tensorflow hub to retrain existing models, however tensorflow supports the hub library only on their 2.2 version. And The Estimator azure presents supports tf 2.0.

When I list tensorflow 2.2 as a required dependency as a pip package, during docker image creation the system fails - it seems like horovod is responsible, - that it cannot find the correct libraries.

Is this possible to be fixed? as in either an Estimator with tf 2.2 support, or an esitmator without the horovod - as I do not need a distributed system for my solution.

Answer 1

Following the pointers from @romungi-MSFT, defining estimator with gpubase image; "mcr.microsoft.com/azureml/base-gpu:openmpi3.1.2-cuda10.1-cudnn7-ubuntu18.04" solves the problem, and Tensorflow 2.2 can be included. Tensorflow uses GPU by default when available.

 estimator = Estimator(source_directory=experiment_folder,
                       compute_target=compute_target,
                       script_params=script_params,
                       entry_script='rps_efn_b0.py',
                       node_count=1,        
                       conda_packages=['ipykernel'],
                       pip_packages = ['azureml-sdk',
                                       'pyarrow',
                                       'pyspark',
                                       'azureml-mlflow',
                                       'joblib',
                                       'matplotlib',
                                       'Pillow',
                                       'tensorflow==2.2',
                                       'tensorflow-datasets',
                                       'tensorflow-hub',
                                       'azureml-defaults',
                                       'azureml-dataprep[fuse,pandas]'],
                       custom_docker_image='mcr.microsoft.com/azureml/base-gpu:openmpi3.1.2-cuda10.1-cudnn7-ubuntu18.04')