玩转JetBot自动驾驶 (六)训练JetBot自己认识危险

原版翻译notebook可在Github下载:
https://github.com/ling3ye/jetbot/blob/master/notebooks/collision_avoidance/train_model.ipynb

在训练之前,建议你把JetBot关机,使用一个5V,3A的充电头或专用电源给JetBot供电。不建议使用移动电源,因为在训练的过程中会消耗很多资源很费电,如果嫌麻烦的请把移动电源充满电,以免在训练在训练过程中没电了,白费劲了。

开机,通过浏览器输入如下网址:
http://<你JetBot的IP地址>:8888

会看到Jupyter Lab的界面提示登陆:

默认的账户密码都是:jetbot

在Jupyter Lab的文件浏览器中,找到collision_avoidance文件夹,点开train_model.ipynb开始训练吧。

避障 – 训练模型

训练的方法很简单,我们将使用图像分类器来训练两个类freeblocked,我们用这个训练完的模型来避免碰撞。为此,我们将使用一个流行的深度学习库 PyTorch

import torch
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

上传并提取数据集

上次,我们上一次收集并打包成.zip的数据集,然后通过调用shell (命令行)命令来提取(解压缩)此数据集。如果你直接就在jetson nano上训练,可以跳过这一步。

!unzip -q dataset.zip

你应该见到一个名为dataset的文件夹出现在文件浏览器上。

创建数据集实例

现在我们使用torchvision.datasets 库中的ImageFolder数据集类。里面有个附加torchvision.transforms库用于转换数据,为训练模型做好准备。

dataset = datasets.ImageFolder(
    'dataset',
    transforms.Compose([
        transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
)

将数据集拆分为训练集和测试集

接下来,我们将数据集拆分为 训练集 和 测试集。测试集将用于验证我们训练完的模型准确性。

train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 50, 50])

创建数据加载器以批量加载数据

我们将创建两个DataLoader实例,它们为洗牌数据提供实用程序,生成批次图像,并与多个任务并行加载样本。

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4
)

test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4
)

定义神经网络

现在,我们定义我们将要训练的神经网络。 torchvision 库提供了一系列我们可以使用的预训练模型。

在一个称为迁移学习的过程中,我们可以重新利用预先训练的模型(在数百万图像上进行训练),以获得可能少的数据,准确完成尽量多的任务。

在预训练模型的原始训练中学到的重要特征可重复用于新任务。 我们将使用alexnet模型。

model = models.alexnet(pretrained=True)

alexnet模型最初是针对具有1000个类标签的数据集进行训练的,但我们的数据集只有两个类标签! 我们将把最好的层替换为最新的,未经训练的层只有两个输出。

model.classifier[6] = torch.nn.Linear(model.classifier[6].in_features, 2)

最后,我们将模型转移到GPU上执行

device = torch.device('cuda')
model = model.to(device)

训练神经网络

使用下面的代码,将开始训练我们的神经网络,在运行完每个世代后,保存表现最佳的模型。

一个世代是所有数据运行一遍

NUM_EPOCHS = 30
BEST_MODEL_PATH = 'best_model.pth'
best_accuracy = 0.0

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

for epoch in range(NUM_EPOCHS):
    
    for images, labels in iter(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = F.cross_entropy(outputs, labels)
        loss.backward()
        optimizer.step()
    
    test_error_count = 0.0
    for images, labels in iter(test_loader):
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        test_error_count += float(torch.sum(torch.abs(labels - outputs.argmax(1))))
    
    test_accuracy = 1.0 - float(test_error_count) / float(len(test_dataset))
    print('%d: %f' % (epoch, test_accuracy))
    if test_accuracy > best_accuracy:
        torch.save(model.state_dict(), BEST_MODEL_PATH)
        best_accuracy = test_accuracy

当这完成后,你应该会见到一个文件best_model.pth在Jupyter Lab的文件浏览器上,鼠标右键可以选择下载,可以保存这个模型在你的系统平台上。

可以说,这就是JetBot自己学习的知识了。在下一篇,我们将使用这个知识,让JetBot独自运行,并且能避开障碍。

发表评论

电子邮件地址不会被公开。 必填项已用*标注

*
*

BACK TO TOP