基于5060ti显卡的强化学习环境搭建指南：从硬件配置到性能优化

1次阅读

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

5060ti 显卡作为中高端 GPU，在强化学习任务中具有不错的性价比。然而在实际使用中，开发者常遇到以下问题：

驱动兼容性问题：新显卡需要特定版本的驱动支持
CUDA 版本冲突：不同深度学习框架对 CUDA 版本要求不同
计算性能瓶颈：未优化时无法充分发挥显卡潜力

主流深度学习框架对 5060ti 的支持情况：

框架	推荐版本	CUDA 要求	优势
PyTorch	1.12+	CUDA 11.6	动态图，调试方便
TensorFlow	2.10+	CUDA 11.2	静态图，部署方便

推荐使用 PyTorch，因其对 5060ti 的兼容性更好，且强化学习开发更灵活。

显卡：NVIDIA 5060ti (8GB 显存以上)
CPU：Intel i7 或同等 AMD 处理器
内存：32GB 以上
存储：建议 SSD 1TB

卸载旧驱动（如存在）
```
sudo apt-get purge nvidia*
```

添加官方驱动仓库

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

安装推荐驱动（以 515 版本为例）
```
sudo apt install nvidia-driver-515
```

下载 CUDA 11.6

wget https://developer.download.nvidia.com/compute/cuda/11.6.0/local_installers/cuda_11.6.0_510.39.01_linux.run
sudo sh cuda_11.6.0_510.39.01_linux.run

配置环境变量

export PATH=/usr/local/cuda-11.6/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.6/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

安装 cuDNN 8.4

tar -xzvf cudnn-11.6-linux-x64-v8.4.1.50.tgz
sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64

conda create -n rl python=3.8
conda activate rl
pip install torch==1.12.1+cu116 -f https://download.pytorch.org/whl/torch_stable.html
pip install gym[atari] tensorboard

import torch
import torch.nn as nn
import gym
from collections import deque
import random

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')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

# 初始化网络和优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DQN(state_dim, action_dim).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练循环
for episode in range(100):
    state = env.reset()
    total_reward = 0

    while True:
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
        action = model(state_tensor).argmax().item()

        next_state, reward, done, _ = env.step(action)
        total_reward += reward

        # 这里简化为在线学习，实际应使用经验回放
        loss = ... # 计算损失
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        state = next_state
        if done:
            break

    print(f'Episode {episode}, Reward: {total_reward}')

watch -n 0.5 nvidia-smi

从较小的 batch size 开始（如 32）
逐步增加直到显存占用达到 80%
使用 torch.cuda.empty_cache() 定期清理缓存

使用混合精度训练

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

使用梯度累积

for i, (inputs, labels) in enumerate(train_loader):
    loss = model(inputs, labels)
    loss = loss / accumulation_steps
    loss.backward()

    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()