原版翻译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开始训练吧。

避障 – 训练模型

训练的方法很简单，我们将使用图像分类器来训练两个类free和blocked，我们用这个训练完的模型来避免碰撞。为此，我们将使用一个流行的深度学习库 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独自运行，并且能避开障碍。