YOLO全称You Only Look Once,是一个端到端(end-to-end)的目标检测算法,现在已经发展到第三个版本。由于第三个版本已经比较复杂,我们选学习第一个版本。
github上有个同学实现了一个pytorch的版本:https://github.com/xiongzihua/pytorch-YOLO-v1
我基于他的源码学习,学习过程中的代码修改放在:https://git.dev.tencent.com/zzpu/yolov1.git
1 数据准备
根据github上的说明,需要下载voc2012和voc2007test数据,附上地址:
voc2012:http://host.robots.ox.ac.uk/pascal/VOC/voc2012/
voc2007:http://host.robots.ox.ac.uk/pascal/VOC/voc2007/
下载完之后,解压出数据,其中都包含一个名为JPEGImages的文件夹,是数据集的图片文件,将两个数据集的对应文件夹中的所有图片放到同一个文件夹中(如:allimgs),然后拷贝源码中的voc2012.txt,voc2007test.txt两个文件到该文件夹。
至此,数据准备工作完毕。
2 源码修改
你可以直接使用我学习过程中源码,改源码是修改好的。其实就两个地方需要修改。
2.1 修改支持pytorch1.0
因为源码是作者一年前实现的,现在pytorch1.0已经发布,有一个小地方需要修改:
if (i+1) % 5 == 0:
print ('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f'
%(epoch+1, num_epochs, i+1, len(train_loader), loss.data[0], total_loss / (i+1)))
num_iter += 1
vis.plot_train_val(loss_train=total_loss/(i+1))
data[0]改为item(),所有地方都要改:
if (i+1) % 5 == 0:
print ('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f'
%(epoch+1, num_epochs, i+1, len(train_loader), loss.data.item(), total_loss / (i+1)))
num_iter += 1
vis.plot_train_val(loss_train=total_loss/(i+1))
2.2 修改数据集路劲和训练数据集
路径改为你放图片的地方:
file_root = '/home/xzh/data/VOCdevkit/VOC2012/allimgs/'
数据集,训练只用voc2012:
train_dataset = yoloDataset(root=file_root,list_file=['voc2012.txt','voc2007.txt'],train=True,transform = [transforms.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=4)
改为:
train_dataset = yoloDataset(root=file_root,list_file=['voc2012.txt'],train=True,transform = [transforms.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=4)
3 执行训练命令
python train.py
如下图: