3. 端到端模型在真实道路数据上的迁移学习

任务说明

本案例目标是在真实道路数据集上对预训练端到端驾驶模型进行迁移学习（微调），使模型适应真实场景与仿真场景之间的分布差异。预训练模型基于仿真数据训练，直接部署到真实道路数据上时，因光照、纹理、道路标记等分布偏移导致控制精度下降。你需要完成：加载预训练模型权重、在真实道路数据集上微调、评估微调前后模型在真实验证集上的表现差异。

数据与预训练模型说明

本案例使用Udacity模拟器数据集作为预训练来源，真实道路数据采用Comma.ai开源真实道路数据集（包含约10小时真实驾驶图像与转向角）。由于完整数据集较大，本案例提供数据加载接口，实际运行时需将数据路径指向本地已下载的数据目录。

预训练模型采用NVIDIA提出的端到端卷积神经网络架构（PilotNet），输入为单目RGB图像（66×200），输出为预测转向角。模型权重文件pretrained_pilotnet.pth需从Udacity预训练权重仓库下载后放置于./weights/目录。

实现代码

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 os
import json

# ---------- 1. PilotNet模型定义 ----------
class PilotNet(nn.Module):
    def __init__(self):
        super(PilotNet, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 24, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(24, 36, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(36, 48, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(48, 64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )
        self.fc_layers = nn.Sequential(
            nn.Linear(1152, 100),
            nn.ReLU(),
            nn.Linear(100, 50),
            nn.ReLU(),
            nn.Linear(50, 1)
        )
    
    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1)
        return self.fc_layers(x)

# ---------- 2. 真实道路数据集加载 ----------
class RealRoadDataset(Dataset):
    def __init__(self, data_dir, split='train'):
        self.data_dir = data_dir
        self.split = split
        self.samples = []
        # 读取数据索引文件（每行: 图像路径, 转向角）
        index_file = os.path.join(data_dir, f'{split}_index.json')
        with open(index_file, 'r') as f:
            for line in f:
                item = json.loads(line)
                self.samples.append((item['image'], item['steering']))
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        img_path, steering = self.samples[idx]
        img = cv2.imread(os.path.join(self.data_dir, img_path))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
        img = cv2.resize(img, (200, 66))
        img = img / 255.0
        img = torch.FloatTensor(img).permute(2, 0, 1)
        return img, torch.FloatTensor([steering])

# ---------- 3. 微调训练配置 ----------
def fine_tune(model, train_loader, val_loader, epochs=5, lr=1e-4):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0.0
        for batch_img, batch_steer in train_loader:
            batch_img = batch_img.to(device)
            batch_steer = batch_steer.to(device)
            
            optimizer.zero_grad()
            outputs = model(batch_img)
            loss = criterion(outputs, batch_steer)
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * batch_img.size(0)
        
        # 验证
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for batch_img, batch_steer in val_loader:
                batch_img = batch_img.to(device)
                batch_steer = batch_steer.to(device)
                outputs = model(batch_img)
                loss = criterion(outputs, batch_steer)
                val_loss += loss.item() * batch_img.size(0)
        
        print(f'Epoch {epoch+1}/{epochs} | Train Loss: {train_loss/len(train_loader.dataset):.4f} | Val Loss: {val_loss/len(val_loader.dataset):.4f}')
    
    return model

# ---------- 4. 主流程 ----------
if __name__ == '__main__':
    # 加载预训练权重
    model = PilotNet()
    model.load_state_dict(torch.load('./weights/pretrained_pilotnet.pth'))
    
    # 准备数据（需提前将真实道路数据按8:2划分并生成索引文件）
    train_dataset = RealRoadDataset('./data/real_road', split='train')
    val_dataset = RealRoadDataset('./data/real_road', split='val')
    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)
    
    # 微调前在真实验证集上评估
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()
    criterion = nn.MSELoss()
    pretrain_val_loss = 0.0
    with torch.no_grad():
        for batch_img, batch_steer in val_loader:
            batch_img = batch_img.to(device)
            batch_steer = batch_steer.to(device)
            outputs = model(batch_img)
            pretrain_val_loss += criterion(outputs, batch_steer).item() * batch_img.size(0)
    pretrain_val_loss /= len(val_dataset)
    print(f'微调前验证集MSE: {pretrain_val_loss:.4f}')
    
    # 执行微调
    model = fine_tune(model, train_loader, val_loader, epochs=5)
    
    # 微调后评估
    model.eval()
    finetune_val_loss = 0.0
    with torch.no_grad():
        for batch_img, batch_steer in val_loader:
            batch_img = batch_img.to(device)
            batch_steer = batch_steer.to(device)
            outputs = model(batch_img)
            finetune_val_loss += criterion(outputs, batch_steer).item() * batch_img.size(0)
    finetune_val_loss /= len(val_dataset)
    print(f'微调后验证集MSE: {finetune_val_loss:.4f}')
    
    # 保存微调后模型
    torch.save(model.state_dict(), './weights/finetuned_pilotnet.pth')
    print('微调完成，模型已保存至 ./weights/finetuned_pilotnet.pth')

数据准备与运行方式

运行前需完成以下数据准备步骤：

下载Comma.ai真实道路数据集，将图像文件存放于./data/real_road/images/目录；
编写脚本将数据集按8:2比例随机划分为训练集和验证集，并生成train_index.json与val_index.json两个索引文件，每行格式为{"image": "images/xxx.jpg", "steering": 0.123}；
确保./weights/pretrained_pilotnet.pth存在；
执行python finetune_pilotnet.py。

验证方法与判定标准

运行脚本后，观察终端输出的微调前后验证集MSE数值。判定标准：

微调后的验证集MSE应低于微调前，表明模型已适应真实道路数据分布；
训练过程中训练损失与验证损失均呈下降趋势，且无发散现象；
微调后模型权重文件成功保存至指定路径。

若微调后MSE未下降，可尝试增大训练轮数至10轮或调整学习率至5e-5后重新运行。

📚 课程目录