玩转JetBot自动驾驶（六）训练JetBot自己认识危险¶

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

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

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

默认的账户密码都是：jetbot

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

避障 – 训练模型¶

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

import torch.optim as optim

import torch.nn.functional as F

import torchvision.datasets as datasets

import torchvision.models as models

import torchvision.transforms as transforms

上传并提取数据集¶

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

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

创建数据集实例¶

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

dataset = datasets.ImageFolder(

transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),

transforms.Resize((224, 224)),

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(

test\_loader = torch.utils.data.DataLoader(

定义神经网络¶

现在，我们定义我们将要训练的神经网络。 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')

训练神经网络¶

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

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

BEST\_MODEL\_PATH = 'best\_model.pth'

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)

loss = F.cross\_entropy(outputs, labels)

for images, labels in iter(test\_loader):

images = images.to(device)

labels = labels.to(device)

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独自运行，并且能避开障碍。

玩转JetBot自动驾驶 （六）训练JetBot自己认识危险¶