4. 基于ResNet的端到端驾驶模型

任务说明

本案例将构建一个基于ResNet-18的端到端驾驶模型，输入为单目前置摄像头图像，输出为方向盘转角预测值。模型采用模仿学习范式，使用NVIDIA提出的端到端驾驶架构思想，以ResNet-18替换原架构中的自定义卷积骨干，完成从图像到控制指令的直接映射。训练数据使用Udacity模拟器采集的driving_log.csv（包含中心摄像头图像路径与对应方向盘转角），代码将完成数据加载、模型定义、训练与评估全流程。

模型架构

模型由三部分组成：ResNet-18骨干网络（移除原始全连接层）、自定义回归头、以及输入预处理层。输入图像resize至224×224×3，经ImageNet均值/方差归一化后送入ResNet-18，输出512维特征图，经全局平均池化后得到512维特征向量，最后通过两层全连接（512→128→1）回归方向盘转角。骨干网络使用ImageNet预训练权重，回归头随机初始化。

代码实现

以下代码完整实现数据加载、模型定义、训练与评估。训练过程中每100个batch输出一次损失值，每个epoch结束后在验证集上计算MSE与相关系数R²作为性能指标。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
import pandas as pd
import numpy as np
from PIL import Image
import os
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score

# ---------- 1. 数据集定义 ----------
class DrivingDataset(Dataset):
    """从Udacity模拟器日志加载图像与转向角"""
    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'].strip())
        image = Image.open(img_path).convert('RGB')
        angle = float(row['steering'])
        if self.transform:
            image = self.transform(image)
        return image, torch.tensor(angle, dtype=torch.float32)

# ---------- 2. 数据预处理 ----------
IMG_SIZE = 224
train_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# ---------- 3. 模型定义 ----------
class ResNetEndToEnd(nn.Module):
    """ResNet-18骨干 + 回归头，输出方向盘转角"""
    def __init__(self, pretrained=True):
        super().__init__()
        # 加载ResNet-18，使用预训练权重
        backbone = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None)
        # 移除原始全连接层，保留特征提取部分
        self.features = nn.Sequential(*list(backbone.children())[:-1])  # 输出 512×1×1
        # 回归头
        self.regressor = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(128, 1)
        )

    def forward(self, x):
        x = self.features(x)      # [B, 512, 1, 1]
        x = self.regressor(x)     # [B, 1]
        return x.squeeze(1)       # [B]

# ---------- 4. 训练配置 ----------
def train_model(model, train_loader, val_loader, epochs=10, lr=1e-4, device='cuda'):
    model.to(device)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5)

    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        for batch_idx, (images, angles) in enumerate(train_loader):
            images, angles = images.to(device), angles.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, angles)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

            if (batch_idx + 1) % 100 == 0:
                avg_loss = running_loss / 100
                print(f'Epoch [{epoch+1}/{epochs}] Batch [{batch_idx+1}/{len(train_loader)}] Loss: {avg_loss:.4f}')
                running_loss = 0.0

        scheduler.step()
        # 每个epoch结束后验证
        val_mse, val_r2 = evaluate_model(model, val_loader, device)
        print(f'Epoch [{epoch+1}/{epochs}] Val MSE: {val_mse:.4f} | Val R²: {val_r2:.4f}')

# ---------- 5. 评估函数 ----------
def evaluate_model(model, val_loader, device='cuda'):
    model.eval()
    all_preds = []
    all_angles = []
    with torch.no_grad():
        for images, angles in val_loader:
            images = images.to(device)
            outputs = model(images).cpu().numpy()
            all_preds.extend(outputs)
            all_angles.extend(angles.numpy())
    mse = np.mean((np.array(all_preds) - np.array(all_angles)) ** 2)
    r2 = r2_score(all_angles, all_preds)
    return mse, r2

# ---------- 6. 主程序 ----------
if __name__ == '__main__':
    # 数据路径（请根据实际存放位置修改）
    CSV_PATH = 'data/driving_log.csv'
    IMG_DIR = 'data/IMG'

    # 加载数据集并划分训练/验证集
    full_dataset = DrivingDataset(CSV_PATH, IMG_DIR, transform=train_transform)
    train_idx, val_idx = train_test_split(
        range(len(full_dataset)), test_size=0.2, random_state=42
    )
    train_dataset = torch.utils.data.Subset(full_dataset, train_idx)
    val_dataset = torch.utils.data.Subset(full_dataset, val_idx)

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)

    # 初始化模型
    model = ResNetEndToEnd(pretrained=True)

    # 训练（使用GPU如果可用）
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f'Using device: {device}')
    train_model(model, train_loader, val_loader, epochs=10, lr=1e-4, device=device)

    # 保存模型
    torch.save(model.state_dict(), 'resnet_e2e_driving.pth')
    print('模型已保存至 resnet_e2e_driving.pth')

运行与验证

依赖安装：

pip install torch torchvision pandas numpy scikit-learn pillow

数据准备：从Udacity模拟器采集数据，确保driving_log.csv中center列指向图像文件名，图像存放于IMG目录。将CSV与IMG目录路径填入主程序对应变量。

执行训练：

python train_resnet_e2e.py

判定标准：训练过程中损失值应逐步下降；验证集MSE应低于0.05（归一化转角范围约-1~1），R²应大于0.7。若达到上述指标，说明模型已学习到有效的图像到转向映射关系。训练完成后，可使用model.load_state_dict(torch.load('resnet_e2e_driving.pth'))加载模型，输入单张图像即可获得预测转角。

预期输出示例

以下为训练日志的预期格式（数值取决于具体数据，仅作格式参考）：

Using device: cuda
Epoch [1/10] Batch [100/200] Loss: 0.0821
Epoch [1/10] Val MSE: 0.0612 | Val R²: 0.5234
Epoch [2/10] Batch [100/200] Loss: 0.0543
Epoch [2/10] Val MSE: 0.0487 | Val R²: 0.6312
...
Epoch [10/10] Val MSE: 0.0321 | Val R²: 0.7825
模型已保存至 resnet_e2e_driving.pth