Module tf_nndct not found in vitis-ai-tensorflow2-cpu

[dguenzel]

Hello,

I am trying to follow the tutorial TF2-Vitis-AI-Optimizer for Vitis AI 3.5. My current development machine does not have a GPU, so I use the CPU-only Docker image. I have tried with the image I built myself according to the tutorial and with the one on Docker Hub.

In step 4 (section 3.4) Pruning and Fine-Tuning I get an error:

(vitis-ai-tensorflow2) vitis-ai-user@fpgadev:/workspace/tutorials/TF2-Vitis-AI-Optimizer/files/dogs-vs-cats_mobilenetv2$ python -u implement.py --mode prune --build_dir ${BUILD} 2>&1 | tee ${LOG}/prune.log
tf_nndct               not found

-----------------------------------------
Keras version      :  2.12.0
TensorFlow version :  2.12.0
3.8.6 | packaged by conda-forge | (default, Oct  7 2020, 19:08:05) 
[GCC 7.5.0]
-----------------------------------------
Please install GPU version of TF2

[DB INFO] IMPLEMENT PRUNING ...

Loading weights from ./build_pr/float_model/f_model.h5

[DB INFO] Original model accuracy: 93.4200 %
-----------------------------------------
Pruning iteration  1  of 6  Pruning ratio:  0.1
Target accuracy for this iteration:  0.9061739891767502
Traceback (most recent call last):
  File "implement.py", line 342, in <module>
    run_main()
  File "implement.py", line 339, in run_main
    implement(args.build_dir, args.mode, args.target)
  File "implement.py", line 219, in implement
    pruned_model = prune(model,prune_ratio,test_dataset)
  File "implement.py", line 151, in prune
    runner = IterativePruningRunner(model, input_spec)
NameError: name 'IterativePruningRunner' is not defined

Indeed, there is no Python module "tf_nndct" installed anywhere in the Docker image. If I look inside the image for ROCm it is in: /opt/vitis_ai/conda/envs/vitis-ai-tensorflow2/lib/python3.8/site-packages/tf_nndct

It appears that there has been a similar issue #1283 with this module in the past, but the offered solution does not change anything for me.

I could not find any indication that the iterative pruning is a GPU-only feature, so is this a bug during creation of the CPU-only Docker image?

Thank you and best regards!

[janifer112x]

Hi @dguenzel, Thanks for trying with VitisAI. this maybe some mis-leading information in the documentation. optimizer is only a feature in GPU docker(ROCM and Nivida GPU), for CPU, the feature is not supported.

[dguenzel]

Thank you for the quick reply. Is a GPU a general requirement for the optimizer or just a limitation of the specific tutorial?

[janifer112x]

Hi @dguenzel , it is general requirement. if you want to use optimizer, have to use GPU docker instead.:)

[dguenzel]

Thanks for the clarification. I am a bit surprised as I don't recall this being mentioned on either github.io or in UG1414.

I used the ROCm TF2 Docker image now and was able to proceed with the tutorial even without a GPU. I will close this issue, perhaps the documentation can be updated to mention that the GPU containers have to be used if optimization is desired.

[janifer112x]

Thanks @dguenzel , I will inform our Document team to update accordingly in the github.io

[haminhson1st]

Hi @janifer112x, I'm working on the same tutorial of the post but using GPU, however, I cannot generate the .h5 file in the first step when running implement.py --mode training. It kept stopping in the middle of the program without error logging anything. Please help me, the setup is pretty much the same.

[adrivg10]

Hi @haminhson1st I faced the same issue. You are probably running out of memory (the default batch size is 150). You should open the config.py and modify it, I tried batchsize=64 and it works for me (I`m using NVIDIA GeForce RTX 3060)