1060 AI算力入门指南：从零搭建你的第一个深度学习推理环境

1次阅读

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

NVIDIA GTX 1060（6GB 显存版）作为 Pascal 架构的中端显卡，在 AI 计算中有以下特性：

FP16 半精度：支持基础运算但无 Tensor Core 加速
INT8 量化：需通过 TensorRT 等工具链实现
显存限制：6GB 容量建议模型参数量控制在 1 亿以下
对比参考：FP32 性能约 4.4 TFLOPS，相当于 RTX 3060 的 1 /5

在 Ubuntu 20.04 环境下测试 ResNet50 推理速度：

CUDA 版本	平均推理时延(ms)	显存占用(MB)
10.2	28.5	3420
11.7	26.1	3380

推荐选择：
– 新项目建议 CUDA 11.7 + cuDNN 8.5
– 需兼容旧框架时选择 CUDA 10.2

# 添加官方驱动 PPA
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update

# 安装 470 版驱动（兼容性最佳）sudo apt install nvidia-driver-470

# 下载 CUDA 11.7 本地安装包
wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.65.01_linux.run

# 运行安装程序（注意取消驱动选项）sudo sh cuda_11.7.1_515.65.01_linux.run --toolkit --silent --override

# 解压下载的 tar 包
sudo tar -xzvf cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz

# 复制库文件
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include/
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64/

创建 conda 环境：

conda create -n pytorch_env python=3.8
conda activate pytorch_env

# 安装 PyTorch 1.12（匹配 CUDA 11.7）conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch

验证脚本gpu_check.py：

import torch

# 检查 CUDA 可用性
if torch.cuda.is_available():
    device = torch.device("cuda")
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"CUDA 版本: {torch.version.cuda}")
else:
    raise RuntimeError("CUDA 设备不可用，请检查驱动安装")

from torchvision import datasets, transforms

# 使用 GPU 加速的数据预处理
transform = transforms.Compose([transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_set = datasets.FashionMNIST(
    'data', 
    download=True,
    train=True,
    transform=transform
)

# 使用 Pin Memory 加速数据加载
train_loader = torch.utils.data.DataLoader(
    train_set,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True  # 关键参数
)

model = SimpleCNN()  # 自定义网络

# 错误写法：先 to(device)再定义优化器
model = model.to('cuda')
optimizer = torch.optim.Adam(model.parameters())  # 必须在同一设备

# 正确流程：model = SimpleCNN()
optimizer = torch.optim.Adam(model.parameters())
model = model.to('cuda')

当出现 CUDA out of memory 错误时：

逐步降低 batch size（从 64→32→16）

启用梯度检查点：

from torch.utils.checkpoint import checkpoint

# 在前向传播中分段计算
def forward(self, x):
    x = checkpoint(self.block1, x)
    x = checkpoint(self.block2, x)
    return x

在模型输入尺寸固定时启用：

torch.backends.cudnn.benchmark = True  # 自动寻找最优卷积算法

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()

# 实时监控命令
watch -n 1 nvidia-smi

输出关键字段解读：
– Volatile GPU-Util：瞬时利用率
– Memory-Usage：已用 / 总显存
– Processes：占用显存的进程

CUDA 版本	最低驱动版本
10.2	440.33
11.7	515.43.04

运行 nvidia-smi 确认显存占用
减少 batch size 或模型规模
使用torch.cuda.empty_cache()

考虑使用梯度累积：

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

    if (i+1) % 4 == 0:  # 每 4 个 batch 更新一次
        optimizer.step()
        optimizer.zero_grad()