feat: 2.16.1的bug修复

script/2_text_test.sh

@@ -9,8 +9,17 @@ fi
 if [[ "$1" == *"pytorch"* ]]; then 
   docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/pytorch/gpt2 $1 python infer.py; fi
 
-if [[ "$1" == *"tensorflow"* ]]; then 
-  docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/bert $1 python infer.py; fi
+if [[ "$1" == *"tensorflow"* ]]; then
+  tensorflow_version=$(echo "$1" | cut -d: -f2 | cut -d- -f1)
+  # 当tensorflow版本为2.16.1时，不添加环境变量找不到cuda，所以需要这样执行验证。在正常交互式启动容器时，会默认激活/etc/bash.bashrc，可以正常找到cuda
+  if [[ "$tensorflow_version" == "2.16.1" ]]; then
+    python_version=$(echo $1 | awk -F'[-:]' '{for(i=3;i<=NF;i++) if($i ~ /^py[0-9]+\.[0-9]+$/) {gsub(/^py/,"",$i); print $i; exit}}')
+    docker run --rm --platform=linux/amd64 --gpus all \
+    -e CUDNN_PATH="/opt/conda/lib/python$python_version/site-packages/nvidia/cudnn" \
+    -e LD_LIBRARY_PATH="/opt/conda/lib/python$python_version/site-packages/nvidia/cudnn/lib:/usr/local/cuda/lib64" \
+    -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/bert $1 python infer.py
+  else
+    docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/bert $1 python infer.py; fi; fi
 if [[ "$1" == *"paddle"* ]]; then 
    TARGET_DIR=gpu-base-image-test/paddletest 
    docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace $1 python text.py; fi

script/3_image_test.sh

@@ -9,8 +9,17 @@ fi
 if [[ "$1" == *"pytorch"* ]]; then 
   docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/pytorch/stable-diffusion-v1-4 $1 python infer.py; fi
 
-if [[ "$1" == *"tensorflow"* ]]; then 
-  docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/mnist $1 python train.py; fi
+if [[ "$1" == *"tensorflow"* ]]; then
+  tensorflow_version=$(echo "$1" | cut -d: -f2 | cut -d- -f1)
+  # 当tensorflow版本为2.16.1时，不添加环境变量找不到cuda，所以需要这样执行验证。在正常交互式启动容器时，会默认激活/etc/bash.bashrc，可以正常找到cuda
+  if [[ "$tensorflow_version" == "2.16.1" ]]; then
+    python_version=$(echo $1 | awk -F'[-:]' '{for(i=3;i<=NF;i++) if($i ~ /^py[0-9]+\.[0-9]+$/) {gsub(/^py/,"",$i); print $i; exit}}')
+    docker run --rm --platform=linux/amd64 --gpus all \
+    -e CUDNN_PATH="/opt/conda/lib/python$python_version/site-packages/nvidia/cudnn" \
+    -e LD_LIBRARY_PATH="/opt/conda/lib/python$python_version/site-packages/nvidia/cudnn/lib:/usr/local/cuda/lib64" \
+    -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/mnist $1 python train.py
+  else
+    docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/mnist $1 python train.py; fi; fi
 
 if [[ "$1" == *"paddle"* ]]; then 
   TARGET_DIR=gpu-base-image-test/paddletest