使用5060ti显卡搭建强化学习环境：从零开始的避坑指南

1次阅读

没有评论

共计 2542 个字符，预计需要花费 7 分钟才能阅读完成。

5060ti 作为一款中端显卡，在强化学习领域具有不错的性价比，但新手在搭建环境时往往会遇到以下问题：

驱动兼容性问题：NVIDIA 驱动版本与 CUDA 版本不匹配导致无法识别显卡
CUDA 版本冲突：不同深度学习框架对 CUDA 版本要求不同，容易造成环境混乱
环境配置复杂：需要同时配置 CUDA、cuDNN、Python 环境等多层依赖
性能陷阱：默认配置可能无法充分发挥显卡性能

经过测试，5060ti 显卡推荐使用以下组合：

驱动版本：470.x 或更高
CUDA 版本：11.3（兼顾 PyTorch 和 TensorFlow 的兼容性）
cuDNN 版本：8.2.1
深度学习框架：
PyTorch 1.10+（对新手更友好）
TensorFlow 2.6+（企业级项目更常用）

卸载旧驱动（如果有）：
```
sudo apt-get purge nvidia*
```

添加官方驱动仓库：

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update

安装推荐驱动：
```
sudo apt install nvidia-driver-470
```
重启后验证：
```
nvidia-smi
```

下载 CUDA 11.3：

wget https://developer.download.nvidia.com/compute/cuda/11.3.1/local_installers/cuda_11.3.1_465.19.01_linux.run

安装（注意不安装驱动）：

sudo sh cuda_11.3.1_465.19.01_linux.run --toolkit --silent

添加环境变量：

echo 'export PATH=/usr/local/cuda-11.3/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

下载 cuDNN 8.2.1（需注册 NVIDIA 账号）

解压并复制文件：

tar -xzvf cudnn-11.3-linux-x64-v8.2.1.32.tgz
sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

以下是一个简单的 DQN 示例（PyTorch）：

import torch
import torch.nn as nn
import gym

# 检查 CUDA 是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 定义网络
class DQN(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(state_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

# 初始化环境
env = gym.make('CartPole-v1')
model = DQN(env.observation_space.shape[0], env.action_space.n).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 简单训练循环
for episode in range(10):
    state = env.reset()
    done = False
    while not done:
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
        action = model(state_tensor).argmax().item()
        next_state, reward, done, _ = env.step(action)
        state = next_state
    print(f"Episode {episode} completed")

在 CartPole 环境下的测试结果：