27 lines
922 B
Bash
27 lines
922 B
Bash
|
#!/bin/bash
|
|||
|
|
|
|||
|
|
echo "安装PyTorch GPU版本（通过官方源）"
|
|||
|
|
echo "==================================="
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
echo "请选择CUDA版本:"
|
|||
|
|
echo "1. CUDA 12.1 (适用于较新的NVIDIA GPU)"
|
|||
|
|
echo "2. CUDA 11.8 (适用于较旧的NVIDIA GPU)"
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
read -p "请输入选项 (1/2): " choice
|
|||
|
|
|
|||
|
|
if [ "$choice" = "1" ]; then
|
|||
|
|
echo "正在安装PyTorch CUDA 12.1版本..."
|
|||
|
|
pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cu121
|
|||
|
|
elif [ "$choice" = "2" ]; then
|
|||
|
|
echo "正在安装PyTorch CUDA 11.8版本..."
|
|||
|
|
pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cu118
|
|||
|
|
else
|
|||
|
|
echo "无效的选项！"
|
|||
|
|
exit 1
|
|||
|
|
fi
|
|||
|
|
|
|||
|
|
echo ""
|
|||
|
|
echo "验证PyTorch GPU支持状态..."
|
|||
|
|
python -c "import torch; print('CUDA是否可用:', torch.cuda.is_available()); print('PyTorch版本:', torch.__version__); print('CUDA版本:', torch.version.cuda if torch.cuda.is_available() else '无')"
|