4. 基于CNN的转向角预测

任务目标

构建一个卷积神经网络(CNN)回归模型,输入单目前视摄像头图像(RGB,尺寸 66×200),输出连续转向角值(弧度)。在独立验证集上计算平均误差(MAE)作为性能指标,并可视化预测值与真实值的对比。

环境与依赖

需要:Python 3.9+,PyTorch 2.0+,torchvision,OpenCV(或PIL),NumPy,Matplotlib。数据集采用Udacity模拟器数据集(或任何包含center_camera.jpg与steering_angle的CSV格式数据集)。

数据准备与模型定义

以下代码实现完整闭环:数据加载、图像预处理、模型定义、训练一个epoch(演示用)、验证评估。实际使用时替换data_dir路径并增加训练轮次即可。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
import pandas as pd
from PIL import Image
import os

# ---------- 1. 数据集 ----------
class SteeringDataset(Dataset):
    def __init__(self, csv_path, img_dir, transform=None):
        self.df = pd.read_csv(csv_path)
        self.img_dir = img_dir
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img_path = os.path.join(self.img_dir, row['center_camera'])
        image = Image.open(img_path).convert('RGB')
        image = image.resize((200, 66))  # 宽200,高66
        image = np.array(image, dtype=np.float32) / 255.0
        image = image.transpose(2, 0, 1)  # CHW
        angle = np.float32(row['steering_angle'])
        return torch.tensor(image), torch.tensor(angle)

# ---------- 2. CNN模型 ----------
class SteeringCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 24, 5, stride=2), nn.ReLU(),
            nn.Conv2d(24, 36, 5, stride=2), nn.ReLU(),
            nn.Conv2d(36, 48, 5, stride=2), nn.ReLU(),
            nn.Conv2d(48, 64, 3), nn.ReLU(),
            nn.Conv2d(64, 64, 3), nn.ReLU(),
        )
        self.regressor = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 1 * 18, 100), nn.ReLU(),
            nn.Linear(100, 50), nn.ReLU(),
            nn.Linear(50, 1)
        )

    def forward(self, x):
        return self.regressor(self.features(x))

# ---------- 3. 训练与验证 ----------
def train_one_epoch(model, loader, optimizer, criterion):
    model.train()
    total_loss = 0
    for images, angles in loader:
        optimizer.zero_grad()
        outputs = model(images).squeeze()
        loss = criterion(outputs, angles)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

def evaluate(model, loader, criterion):
    model.eval()
    losses = []
    with torch.no_grad():
        for images, angles in loader:
            outputs = model(images).squeeze()
            loss = criterion(outputs, angles)
            losses.append(loss.item())
    return np.mean(losses)

# ---------- 4. 主流程 ----------
if __name__ == '__main__':
    # 假设数据目录结构:data/driving_log.csv, data/IMG/xxx.jpg
    data_dir = 'data'
    train_dataset = SteeringDataset(
        os.path.join(data_dir, 'driving_log.csv'),
        os.path.join(data_dir, 'IMG'))
    # 简单划分:80%训练,20%验证
    train_size = int(0.8 * len(train_dataset))
    val_size = len(train_dataset) - train_size
    train_ds, val_ds = torch.utils.data.random_split(
        train_dataset, [train_size, val_size])

    train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=64, shuffle=False)

    model = SteeringCNN()
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    # 演示:仅训练1个epoch(实际训练建议50-100轮)
    loss = train_one_epoch(model, train_loader, optimizer, criterion)
    val_mae = evaluate(model, val_loader, nn.L1Loss())

    print(f'训练损失(MSE): {loss:.4f}')
    print(f'验证MAE: {val_mae:.4f} 弧度')

    # 保存模型
    torch.save(model.state_dict(), 'steering_cnn.pth')

运行方式

  1. 准备数据:将Udacity模拟器数据解压,确保CSV中center_camera列为图片文件名,steering_angle列为转向角(弧度)。
  2. 修改data_dir指向你的数据目录。
  3. 执行:python steering_cnn.py

验证方法与通过标准

程序运行后输出验证集MAE。判定通过的参考标准:

  • MAE < 0.1 弧度(约5.7°)说明模型已学到基本转向规律;
  • MAE < 0.05 弧度(约2.9°)为良好性能;
  • 若MAE > 0.15,需增加训练轮次、调低学习率或检查数据预处理。

注意:上述数值为经验参考,实际阈值取决于数据集难度与分布。首次运行仅1个epoch仅用于验证代码流程正确,MAE通常较高(约0.2-0.4),属正常现象。

常见问题

  • 图像尺寸不匹配:输入必须为66×200×3,代码中已强制resize。
  • 数据路径错误:确保CSV中图片路径相对于img_dir正确。
  • GPU加速:如需使用GPU,在模型与数据后添加.cuda(),并将输入移至相同设备。
CNN转向角预测流程 输入图像 66×200×3 5×Conv2D 特征提取 stride=2 3×FC 回归头 100→50→1 转向角 标量(弧度) 损失函数:MSE(训练) / MAE(验证) 优化器:Adam,lr=1e-4 验证指标:平均误差(弧度) 数据流:CSV路径 → 图像读取 → resize(200×66) → 归一化 → CHW → 模型