Carla 使用神经网络训练自动驾驶车辆—模型搭建及训练

上一节已经搭建好了一个carla仿真环境，并且进行了数据采集

现在需要使用采集到的摄像头图片和转角数据进行模型训练

创建dataset类

import numpy as np
import config
from torch.utils.data import Datasetclass TrainDataSet(Dataset):def __init__(self):self.x=np.load(config.DATASET_PATH+"x_train.npy")[:config.TRAIN_DATASET_SIZES]self.y=np.load(config.DATASET_PATH+"y_train.npy")[:config.TRAIN_DATASET_SIZES]def __len__(self):return len(self.x)def __getitem__(self,index):return (self.x[index],self.y[index])class TestDataSet(Dataset):def __init__(self):self.x = np.load(config.DATASET_PATH+'x_train.npy')[config.TRAIN_DATASET_SIZES:]self.y = np.load(config.DATASET_PATH+'y_train.npy')[config.TRAIN_DATASET_SIZES:]def __len__(self):return len(self.x)def __getitem__(self, index):return (self.x[index], self.y[index])

模型部分

import torch 
import torch.nn as nn
from torchsummary import summaryclass Block(nn.Module):def __init__(self, in_channels, out_channels, maxpool=False):super().__init__()self.conv = nn.Conv2d(in_channels, out_channels, 3, 1, 1)self.relu = nn.ReLU()self.maxpool = maxpoolself.max_pool = nn.MaxPool2d(2, 2)def forward(self, x):x = self.conv(x)x = self.relu(x)if self.maxpool:x = self.max_pool(x)return xclass Model(nn.Module):def __init__(self, in_channels=3, out_size=1):super().__init__()self.conv_block = nn.Sequential(Block(in_channels, 16),Block(16, 16, True),Block(16, 32),Block(32, 32, True),Block(32, 64),Block(64, 64, True),Block(64, 128),Block(128, 128, True),Block(128, 128),Block(128, 128, True),)self.fc_block = nn.Sequential(nn.Linear(1536, 100),nn.ReLU(),nn.Linear(100, 20),nn.ReLU(),nn.Linear(20, out_size))def forward(self, x):x = self.conv_block(x)x = x.view(x.size(0), -1)x = self.fc_block(x)return xif __name__ == '__main__':model = Model()summary(model, input_size=(3, 66, 200), device='cpu')data = torch.ones(1, 3, 66, 200)out = model(data)print(out.shape)

训练部分：

import torch
import torch.nn as nn
import keyboard
import matplotlib.pyplot as plt
import config
from torch.utils.data import DataLoader
from dataset import TrainDataSet
from model import Modelclass Trainer:def __init__(self, model, train_dataset, model_state=None):self.model = modelself.model_state = model_stateself.train_dataset = train_datasetself.loss_list = []self.main()def main(self):torch.manual_seed(config.SEED)torch.cuda.manual_seed(config.SEED)model = self.modelloss_list = self.loss_listmodel.to(config.DEVICE)if self.model_state:state = torch.load(self.model_state)model.load_state_dict(state)model.train()dataloader = DataLoader(self.train_dataset, batch_size=config.BATCH_SIZE, shuffle=True)optimizer = torch.optim.Adam(model.parameters(), lr=config.LEARNING_RATE)criterion = nn.MSELoss()for epoch in range(config.NUM_EPOCH):for x, y in dataloader:x, y = x.to(config.DEVICE).float(), y.to(config.DEVICE).float()x = x.reshape(x.size(0), 3, 66, 200)optimizer.zero_grad()out = model(x)loss = criterion(out, y)loss.backward()optimizer.step()loss_list.append(loss.item())self._show_loss(loss_list, '3')print('epoc[%i/%i] loss=%.5f' % (epoch, config.NUM_EPOCH, loss.item()))torch.save(model.state_dict(), config.MODEL_PATH+'model_state.pth')def _show_loss(self, loss_list, key='3'):if keyboard.is_pressed(key):plt.plot(loss_list)plt.ylim(0, 0.1)plt.show()if __name__ == '__main__':# Trainer(Model(), TrainDataSet(), model_state='model_state.pth')Trainer(Model(), TrainDataSet())

这样就会得到神经网络的权重model_state.pth，下一节使用训练好的神经网络在carla环境中进行自动驾驶测试