Custom TF installation on top of NGC container

AI & Data Science Deep Learning (Training & Inference) Frameworks
1

Hi, I’m having a problem running this notebook with TF NGC container version 20.01-tf2-py3:

https://github.com/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb

I’m getting a segfault when doing inference:

output_dict = model(input_tensor)

But seems to be a TF problem, already adjusted at 2.0 nightly (https://devtalk.nvidia.com/default/topic/1068486/container-tensorflow/segmentation-fault-when-running-tf-detection-tutorial/).

Point is, we do not have a 2.0 + version ready. What would be the most adequate way to them update the framework within this container but also keeping the GPU support?

I’ve tried a naive approach to update

pip install -U --pre tensorflow=="2.*")

But seems that inference is not running on GPU at all, only CPU.

Thoughts?

2

The 20.02 NGC container should be released in the next few days and will include TF 2.1 with the needed fix.