Running JetBot Collision Avoidance Train避障训练新模型

为何要重新训练DNN模型？

1、默认的训练模型采集的数据有限，并不适用所有场景，用户需根据自己场景自行采集数据和训练模型
2、摄像头安装、车身可能存在的差异，也需要重新进行数据的采集和模型的训练

如何训练DNN 模型？

JetBot提供了采集的方法及代码：

Notebooks\collision_avoidance\data_collection.ipynb

第一步，采集数据

• 采集图片数量300张(free无障碍物)+300张(blocked有障碍物)，当然越多越好
• 小车在桌边时，采集的数据要多一些包括free和blocked
• 采集224x224分辨率的图片作为神经网络的输入
• 采集图片时会自动生成Notebooks\collision_avoidance\dataset\free 和 blocked 文件夹，并将添加的free 和 blockded图片存放到对应目录中

free&blocked.png

第二步，保存数据集

采集完图片后，运行并保存dataset.zip ：!zip -r -q dataset.zip dataset

save

第三步，训练采集的数据集

运行train_model_collision_avoidance.py（见附件，此代码参考：train_model.ipynb），训练完成后会生成best_model.pth！

PS:
训练采用Alexnet （若需下载，请打开连接alexnet model ）

第四步，训练完成后导入DNN模型并运行demo.py

复制best_model.pth，到collision_avoidance 目录并运行demo.py，查看训练效果

cp best_model.pth /home/jetbot/Notebooks/collision_avoidance/

---

附件:
train_model_collision_avoidance.ipynb （代码）

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

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])

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
)

model = models.alexnet(pretrained=True)

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

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