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

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

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

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

创建dataset类

import numpy as np
import config
from torch.utils.data import Dataset

class 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 summary

class 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 = maxpool
        self.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 x

class 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 x

if __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 Model

class Trainer:
    def __init__(self, model, train_dataset, model_state=None):
        self.model = model
        self.model_state = model_state
        self.train_dataset = train_dataset
        self.loss_list = []
        self.main()
    
    def main(self):
        torch.manual_seed(config.SEED)
        torch.cuda.manual_seed(config.SEED)

        model = self.model
        loss_list = self.loss_list
        model.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环境中进行自动驾驶测试