使用PyTorch+CUDA的Docker容器异常：CUDA初始化与传参问题

方案1：让bash -c方式正确传递命令行参数 #

bash -c执行模式下，默认会将后续参数作为命令的位置参数，但$@的作用范围有限，需通过特殊写法捕获Vertex AI传入的参数。修改Dockerfile的ENTRYPOINT为：

ENTRYPOINT ["/bin/bash", "-c", "python train_pytorch_model.py \"$@\"", "--"]

说明："--"是核心，它会让bash将所有后续传入的参数（即Vertex AI传递的命令行参数）全部映射到$@中；转义引号包裹$@可避免参数中的空格被错误拆分。

方案2：修复直接执行脚本时的CUDA初始化警告 #

直接用ENTRYPOINT ["python", "train_pytorch_model.py"]时，容器未加载完整的NVIDIA库路径环境变量，导致PyTorch检测CUDA时出现驱动版本误判，以下是三种修复方式：

方法A：显式设置LD_LIBRARY_PATH #

在Dockerfile中添加ENV指令，确保NVIDIA库路径被优先加载：

# 补充NVIDIA库路径到环境变量
ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/nvidia/lib:$LD_LIBRARY_PATH
ENV PATH=/usr/local/nvidia/bin:$PATH

ENTRYPOINT ["python", "train_pytorch_model.py"]

方法B：使用登录shell执行脚本 #

通过bash登录模式加载基础镜像预设的环境变量配置：

ENTRYPOINT ["bash", "-l", "-c", "python train_pytorch_model.py"]

说明：-l参数让bash以登录shell启动，会自动读取/etc/profile或~/.bash_profile中的环境变量配置，其中包含NVIDIA容器运行时设置的正确库路径。

方法C：在脚本开头手动加载环境变量 #

若不想修改Dockerfile，可在train_pytorch_model.py开头添加以下代码，强制设置正确的CUDA库路径：

import os
os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib64:/usr/local/nvidia/lib:' + os.environ.get('LD_LIBRARY_PATH', '')
os.environ['PATH'] = '/usr/local/nvidia/bin:' + os.environ.get('PATH', '')

# 后续PyTorch初始化代码
import torch
print(torch.cuda.is_available())

内容的提问来源于stack exchange，提问作者Michael