YOLOV1详解——Pytorch版

  • 1 YOLOV1
    • 1 数据处理
      • 1.1 数据集划分
      • 1.2 读入xml文件
      • 1.3 数据增强
    • 2 训练
      • 2.1 Backbone
      • 2.2 Loss
      • 2.3 train
    • 3 预测
      • 3.1 图片输入模型得到预测结果

1 YOLOV1

  1. YOLOV1的算法流程。YOLOV1把图像从448x448x3下采样6次得到特征图7x7x30，特征图上的每个网格生成两个框预测物体和物体存在的概率以及20个物体的类别，30=Bx(loc+conf)+cls = 2x(4+1)+20。B的数量为2，是预测框的数量。loc的数量为4，是预测框中心店的偏移和宽高。conf的数量为2，代表每个框与物体的概率。cls代表预测物体的数量。如果真实框落在某个网格内，网格内置信度高得一个框预测真实框。
  2. YOLOV1的优缺点。YOLOV1是end-to-end的模型，模型简单；速度快，预测精度逊色与Faster_RCNN，但是预测速度大幅提升；在不同的数据集上验证证明模型泛化能力强。不需要滑动窗口得到区域框，预测精度相对于其他的模型有所提高。

由于YOLOV1是end没有设置先验框，因此预测精度受的影响且迁移能力差；只有一个特征层，对小物体和群体的小物体预测能力差，这样也意味着有更多的修改空间。

参考代码：https://github.com/abeardear/pytorch-YOLO-v1

'''
数据集：VOC2007VOC2007/Annotations  00001.xml00002.xml.....VOC2007/JPEGImages   00001.jpg00002.jpg.....VOCdevkit/VOC2007/ImageSets/Maintrain.txttest.txtval.txt
代码流程：1. 数据处理1.1 data_split_11.py 把数据集划分成训练集train.txt、验证集、测试集1.2 data_jpgxlm_12.py 根据生成的数据集取出数据集每个样本的的图片地址和框、类别信息1.3 data_label_13.py 数据增强和编码生成标签2.训练2.1 主干网络、迁移学习 resnet_yolo.py2.2 损失函数 yoloLoss.py2.3 优化器   2.4 导入数据，生成器2.5 训练3.预测3.1 图片输入模型得到预测结果。result = predict_gpu(model,image_name)3.2 画图1.1 data_split_11.py功能：获取数据集索引。inputs  : xmlpath = r'./VOCdevkit/VOC2007/Annotations'   # 所有数据集的.xml.outputs : ./VOCdevkit/VOC2007/ImageSets/Main/train.txt   # 数据集索引./VOCdevkit/VOC2007/ImageSets/Main/test.txt./VOCdevkit/VOC2007/ImageSets/Main/val.txt./VOCdevkit/VOC2007/ImageSets/Main/trainval.txt./VOCdevkit/VOC2007/Annotations:00001.xml00002.xml....../VOCdevkit/VOC2007/ImageSets/Maintrain.txttest.txttrainval.txtval.txttrain.txt:000025000018000022.....1.1.1 获取'./VOCdevkit/VOC2007/Annotations' 下以'.xml'结尾的文件名。1.1.2 从总的样本中按照数据集比例抽取样本，得到每个数据集的索引1.1.3 不同数据集的存储地址1.1.4 遍历样本，根据抽取的样本索引放入不同的数据集中1.1.5 关闭文件1.2 data_jpgxlm_12.py 根据生成的数据集取出数据集每个样本的的图片地址和框、类别信息inputs  : sets、classesoutputs : 2007_train.txt、2007_test.txt、2007_val.txt2007_train.txt:VOCdevkit/VOC2007/JPEGImages/000022.jpg 68,103,368,283,12 186,44,255,230,14VOCdevkit/VOC2007/JPEGImages/000028.jpg 63,18,374,500,7......1.2.1 遍历每个数据集获取数据集图片名。生成解析文件，用来存储图片地址和框的信息。1.2.2 遍历每张图片，解析文件写入图片地址，打开xml文件，读入类别信息和框的信息。(1) getroot()-->'object'-->'difficult','name','bndbox'-->cls_id + box1.2.3 关闭解析文件1.3 data_label_13.py  inputs  : list_f = r'2007_train.txt'outputs :1.3.1 读入文件信息，存储图片、框、类别fenb放入fnames、boxes、labels中。1.3.2 根据索引遍历每张图片，如果是train模式，进行数据增强。（1）随机左右反转。先翻转图片，计算框与左右边的距离，框左右边数值调换位置。（2）随机变换尺度。选择随机数放缩宽，框的x轴数值按同比例放缩。（3）随机模糊。（4）随机改变色度、饱和度、明暗。bgr->hsv->放缩->bgr（5）随机平移。随机选平移数值，平移图片。判断框框平移后是否在图片上，如果框平移后存在，平移框。筛选平移后留下的目标。（6）随机剪切。生成随机剪切起点和宽高，选取剪切图片上的框和目标，剪切图片。1.3.3 图像归一化1.3.4 编码（1）计算框的中心点和宽高（2）计算框再特征图上的位置，计算中心点偏移（3）再特征图上真实框对应的位置写入框、置信度、类别信息1.3.5 根据需要转换通道2.1 主干网络、迁移学习 resnet_yolo.pyinputs  : torch.Size([2, 3, 224, 224])outputs : torch.Size([2, 7, 7, 30])2.1.1 input --> cnnv(s=2)+bn+re+maxpool(s=2) --> layers*5((s=2) * 3) --> cnnv+bn+sigmoid2.2 损失函数 yoloLoss.pyinputs  :pred.shape:  [n, 7, 7, 30]target.shape:[n, 7, 7, 30]outputs :scalar2.2.1 有物体的掩码coo_mask[n, 7, 7, 30]，没有物体的掩码noo_mask[n, 7, 7, 30]。2.2.2 根据noo_mask得到没有物体的预测负样本noo_pred = pred_tensor[noo_mask].view(-1,30)和物体的真实负样本noo_target = target_tensor[noo_mask].view(-1,30)，提取置信度,计算负样本的损失.noo_mask.shape: torch.Size([2, 14, 14, 30])pred_tensor.shape: torch.Size([2, 14, 14, 30])noo_pred.shape: torch.Size([0, 30])noo_target.shape: torch.Size([0, 30])2.2.3 根据coo_mask提取预测框、预测类别、真实框、真实类别。根据预测类别和真实类别计算类别损失。2.2.4 遍历真实框，根据IoU标记与真实框最匹配的预测框和最优的IoU.提取最匹配的预测框，分别计算框和正样本的损失。2.2.5 对各个损失加权，计算总损失。2.3 训练2.3.1 加载模型，设置损失函数、优化器、加载数据、前向传播、向后传播跟新参数。3.1 图片输入模型得到预测结果inputs  : imgoutputs : boxes,cls_indexs,probs 3.1.1 数据预处理+前向传 --> 预测结果pred7x7x303.1.2 预测结果解码 + 筛选(1) 先根据置信度筛选满足条件的mask.(2) 遍历特征图上每一个网格点和mask，对满足条件的框解码。(3) 根据置信度和预测物体的条件概率得到预测概率，根据预测概率对框筛选。(4) 对上一步筛选的框进行NMS。3.1.3 把预测框映射到原图3.2 画框3.2.1 画框3.2.2 在框上标记预测物体类别和概率'''

1 数据处理

数据处理主要包括数据集划分、读入xml文件、数据增强三大部分。

首先把数据划分为训练集、验证集和测试集，每种类型的数据集里存储的是图片的名称，比如2002001.jpg,2002002.jpg,2002003.jpg图片，在数据集中是2002001，2002002、2002003。第二步根据数据集读入图片的地址和和真实框的信息。第三部就是根据图片和真实框的信息进行数据增强及编码得到标签。

三部分的代码依次存放在 data_split_11.py 、data_jpgxlm_12.py、 data_label_13.py，点击运行就会得到对应结果。然后点击 train.py 就可以训练模型，根据训练模型的参数就可以点击 predict.py 预测了。

1.1 数据集划分

数据集划分在data_split_11.py文件中。

已知信息：在VOCdevkit/VOC2007/Annotations存储的是每个图片对应的xml信息、训练集、验证集、测试集比例。
输出的是各个数据集的图片名称。

数据集划分流程：

  1. 获取’./VOCdevkit/VOC2007/Annotations’ 下以’.xml’结尾的文件名。
  2. 从总的样本中按照数据集比例抽取样本，得到每个数据集的索引。
  3. 不同数据集的存储地址。
  4. 遍历样本，根据抽取的样本索引放入不同的数据集中。
  5. 关闭文件。
'''
dataSplit
'''
import os
import randomxml_path = r'./VOCdevkit/VOC2007/Annotations'
base_path = r'./VOCdevkit/VOC2007/ImageSets/Main'# 1 样本名字
tmp = []
img_names = os.listdir(xml_path)
for i in img_names:if i.endswith('.xml'):tmp.append(i[:-4])# 2 数据集划分
trainval_ratio = 0.9
train_ratio = 0.9
N = len(tmp)
trainval_num = int(trainval_ratio*N)
train_num = int(train_ratio*trainval_num)
trainval_idx = random.sample(range(N),trainval_num)
train_idx = random.sample(trainval_idx,train_num)# 3 数据集的存储地址
ftrainval = open(os.path.join(base_path,'LS_trainval.txt'),'w')
ftrain = open(os.path.join(base_path,'LS_train.txt'),'w')
fval = open(os.path.join(base_path,'LS_val.txt'),'w')
ftest = open(os.path.join(base_path,'LS_test.txt'),'w')# 4 写入数据
for i in range(N):name = tmp[i]+'\n'if i in trainval_idx:ftrainval.write(name)if i in train_idx:ftrain.write(name)else:fval.write(name)else:ftest.write(name)# 5 关闭文件
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

1.2 读入xml文件

data_jpgxlm_12.py把图片信息和xml信息放在一起。

  1. 遍历每个数据集获取数据集图片名。生成解析文件，用来存储图片地址和框的信息。
  2. 遍历每张图片，解析文件写入图片地址，打开xml文件，读入类别信息和框的信息。 getroot()–>‘object’–>‘difficult’,‘name’,‘bndbox’–>cls_id + box
  3. 关闭解析文件
import xml.etree.ElementTree as ETsets = [('2007','train'),('2007','val'),('2007','test')]
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle","bus", "car", "cat", "chair", "cow", "diningtable","dog", "horse", "motorbike", "person", "pottedplant","sheep", "sofa", "train", "tvmonitor"]       # len(classes) = 20def convert_annotation(year,img_id,list_file):list_file.write("VOCdevkit/VOC%s/JPEGImages/%s.jpg"%(year,img_id))in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year,img_id))root = ET.parse(in_file).getroot()for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult)==1:continuecls_id = classes.index(cls)xml_box = obj.find('bndbox')b = (int(xml_box.find('xmin').text),int(xml_box.find('ymin').text),int(xml_box.find('xmax').text),int(xml_box.find('ymax').text))list_file.write(" "+','.join([str(i) for i in b])+','+str(cls_id))for year ,img_set in sets:img_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year ,img_set)).read().strip().split()list_file = open('%s_%s.txt'%(year,img_set),'w')for img_id in img_ids:convert_annotation(year,img_id,list_file)list_file.write('\n')list_file.close()

1.3 数据增强

data_label_13.py对图片数据增强，增加样本，提高模型的泛化能力。

  1. 首先把图片名称、框、类别信息分别存储。
  2. 生成迭代器，对每个图片进行数据增强。翻转、缩放、模糊、随机变换亮度、随机变换色度、随机变换饱和度、随机平移、随机剪切。
  3. 编码。对于增强后的图片，根据图片宽高，获取框在图片的相对位置、去均值、统一图片尺寸、编码。编码的关键在于找到真实框在特征图上的相对位置。先把真实框左上角和右下角坐标转换为中心点和宽高，中心点坐标x特征图宽高后，向下取整在-1就得到真实框在特征图上的位置ij。中心点坐标x特征图宽高-ij得到真实偏移。根据ij输入偏移、宽高、类别信息。

（1）数据增强代码


''' 数据增强 '''import os
import sys
import cv2
import torch
import random
import os.path
import numpy as np
import matplotlib.pyplot as plt
import torch.utils.data as data
import torchvision.transforms as transformsclass yoloDataset(data.Dataset):image_size = 448def __init__(self,root,list_file,train,transform):print('data init')self.root = root    ## img.jpg_pathself.train = train  ## bool:True，如果是训练模式，就进行数据增强。self.transform = transform  # 转置self.fnames = []   # 存储图片名  eg: 00014.jpgself.boxes = []    # 存放真实框信息self.labels = []   # 存放标签self.mean = (123,117,104)   # 图片各个通道均值，用于归一化数据，加速训练。with open(list_file) as f:lines = f.readlines()   # 读取数据 # line： /Users/ls/PycharmProjects/YOLOV1_LS/VOCdevkit/VOC2007/JPEGImages/000022.jpg 68,103,368,283,12 186,44,255,230,14for line in lines:   splited = line.strip().split()self.fnames.append(splited[0])  # 保存图片名num_boxes = (len(splited)-1)    # 一张图片中真实框的个数box = []     # 存储框label = []   # 存储标签for i in range(1,num_boxes+1): # 遍历每个框tmp = [float(j) for j in splited[i].split(',')] # 把真实框油字符变为float类型，并用‘,’隔开。box.append(tmp[:4]) label.append(int(tmp[4])+1)self.boxes.append(torch.Tensor(box))self.labels.append(torch.LongTensor(label))self.num_samples = len(self.boxes)def __getitem__(self,idx):fname = self.fnames[idx] # 用迭代器遍历每张图片img = cv2.imread(fname)  # 读取图片 cv2.imread(os.path.join(self.root+fname))boxes = self.boxes[idx].clone()labels = self.labels[idx].clone()if self.train:img,boxes = self.random_flip(img,boxes)img,boxes = self.randomScale(img,boxes)img = self.randomBlur(img)img = self.RandomBrightness(img)img = self.RandomHue(img)img = self.RandomSaturation(img)img,boxes,labels = self.randomShift(img,boxes,labels)img,boxes,labels = self.randomCrop(img,boxes,labels)h,w,_ = img.shapeboxes /= torch.Tensor([w,h,w,h]).expand_as(boxes)img = self.BGR2RGB(img)img = self.subMean(img,self.mean)img = cv2.resize(img,(self.image_size,self.image_size))target = self.encoder(boxes,labels)for t in self.transform:img = t(img)return  img,targetdef __len__(self):return self.num_samplesdef encoder(self,boxes,labels):grid_num = 7target = torch.zeros((grid_num,grid_num,30))wh = boxes[:,2:] - boxes[:,:2]cxcy = (boxes[:,2:] + boxes[:,:2])/2for i in range(cxcy.size()[0]):cxcy_sample = cxcy[i]ij = (cxcy_sample*grid_num).ceil()-1dxy = cxcy_sample*grid_num-ijtarget[int(ij[1]),int(ij[0]),:2] = target[int(ij[1]),int(ij[0]),5:7] = dxytarget[int(ij[1]),int(ij[0]),2:4] = target[int(ij[1]),int(ij[0]),7:9] = wh[i]target[int(ij[1]),int(ij[0]),4] = target[int(ij[1]),int(ij[0]),9] = 1target[int(ij[1]),int(ij[0]),int(labels[i])+9] = 1return targetdef BGR2RGB(self,img):return cv2.cvtColor(img,cv2.COLOR_BGR2RGB)def BGR2HSV(self,img):return cv2.cvtColor(img,cv2.COLOR_BGR2HSV)def HVS2BGR(self,img):return cv2.cvtColor(img,cv2.COLOR_HSV2BGR)def RandomBrightness(self,bgr):if random.random()<0.5:hsv = self.BGR2HSV(bgr)h,s,v = cv2.split(hsv)v = v.astype(float)v *= random.choice([0.5,1.5])v = np.clip(v,0,255).astype(hsv.dtype)hsv = cv2.merge((h,s,v))bgr = self.HVS2BGR(hsv)return  bgrdef RandomSaturation(self,bgr):if random.random()<0.5:hsv = self.BGR2HSV(bgr)h,s,v = cv2.split(hsv)s = s.astype(float)s *= random.choice([0.5,1.5])s = np.clip(s,0,255).astype(hsv.dtype)hsv = cv2.merge((h,s,v))bgr = self.HVS2BGR(hsv)return bgrdef RandomHue(self,bgr):if random.random() < 0.5:hsv = self.BGR2HSV(bgr)h,s,v = cv2.split(hsv)h = h.astype(float)h *= random.choice([0.5,1.5])h = np.clip(h,0,255).astype(hsv.dtype)hsv=cv2.merge((h,s,v))bgr = self.HVS2BGR(hsv)return bgrdef randomBlur(self,bgr):if random.random() < 0.5:bgr = cv2.blur(bgr,(5,5))return bgrdef randomShift(self,bgr,boxes,labels):center = (boxes[:,2:]+boxes[:,:2])/2if random.random()<0.5:height,width,c = bgr.shapeafter_shift_imge = np.zeros((height,width,c),dtype=bgr.dtype)after_shift_imge[:,:,:] = (104,117,123)shift_x = random.uniform(-width*0.2,width*0.2)shift_y = random.uniform(-height*0.2,height*0.2)if shift_x>=0 and shift_y>=0:after_shift_imge[int(shift_y):,int(shift_x):,:] = bgr[:height-int(shift_y),:width-int(shift_x),:]elif shift_x>=0 and shift_y<0:after_shift_imge[:height+int(shift_y),int(shift_x):,:] = bgr[-int(shift_y):,:width-int(shift_x),:]elif shift_x <0 and shift_y >=0:after_shift_imge[int(shift_y):,:width+int(shift_x),:] = bgr[:height-int(shift_y),-int(shift_x):,:]elif shift_x<0 and shift_y<0:after_shift_imge[:height+int(shift_y),:width+int(shift_x),:] = bgr[-int(shift_y):,-int(shift_x):,:]shift_xy = torch.FloatTensor([[int(shift_x),int(shift_y)]]).expand_as(center)center = center + shift_xymask1 = (center[:,0]>0)& (center[:,0]>height)mask2 = (center[:,1]>0)& (center[:,1]>width)mask = (mask1 & mask2).view(-1,1)boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)if len(boxes_in) == 0:return bgr,boxes,labelsbox_shift = torch.FloatTensor([[int(shift_x),int(shift_y),int(shift_x),int(shift_y)]]).expand_as(boxes_in)boxes_in = boxes_in+box_shiftlabels_in = labels[mask.view(-1)]return after_shift_imge,boxes_in,labels_inreturn bgr,boxes,labelsdef randomScale(self,bgr,boxes):if random.random() < 0.5:scale = random.uniform(0.8,1.2)h,w,c = bgr.shapebgr = cv2.resize(bgr,(int(w*scale),h))scale_tensor = torch.FloatTensor([[scale,1,scale,1]]).expand_as(boxes)boxes = boxes*scale_tensorreturn bgr,boxesreturn bgr,boxesdef randomCrop(self,bgr,boxes,labels):if random.random() < 0.5:center = (boxes[:,:2]+boxes[:,2:])/2height,width,c = bgr.shapeh = random.uniform(0.6*height,height)w = random.uniform(0.6*width,width)x = random.uniform(0,width-w)y = random.uniform(0,height-h)x,y,h,w = int(x),int(y),int(h),int(w)center = center - torch.FloatTensor([[x,y]]).expand_as(center)mask1 = (center[:,0]>0) & (center[:,0]<w)mask2 = (center[:,1]>0) & (center[:,0]<h)mask = (mask1 & mask2).view(-1,1)boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)if(len(boxes_in)==0):return bgr,boxes,labelsbox_shift = torch.FloatTensor([[x,y,x,y]]).expand_as(boxes_in)boxes_in = boxes_in-box_shiftboxes_in[:,0]=boxes_in[:,0].clamp(0,w)boxes_in[:,2]=boxes_in[:,2].clamp(0,w)boxes_in[:,1]=boxes_in[:,1].clamp(0,h)boxes_in[:,3]=boxes_in[:,3].clamp(0,h)labels_in = labels[mask.view(-1)]img_croped = bgr[y:y+h,x:x+w,:]return img_croped,boxes_in,labels_inreturn bgr,boxes,labelsdef subMean(self,bgr,mean):mean = np.array(mean,dtype=np.float32)bgr = bgr - meanreturn bgrdef random_flip(self,im,boxes):if random.random() < 0.5:im_lr = np.fliplr(im).copy()h,w,_ = im.shapexmin = w - boxes[:,2]xmax = w - boxes[:,0]boxes[:,0] = xminboxes[:,2] = xmaxreturn im_lr,boxesreturn im,boxesdef random_bright(self,im,delta=16):alpha = random.random()if alpha > 0.3:im = im * alpha + random.randrange(-delta,delta)im = im.clip(min=0,max=255).astype(np.uint8)return imif __name__=='__main__':from torch.utils.data import DataLoaderimport torchvision.transforms as transformsfile_root ='/Users/ls/PycharmProjects/YOLOV1_LS'  ## xx.jpglist_f = r'/Users/ls/PycharmProjects/YOLOV1_LS/2007_train.txt'train_dataset = yoloDataset(root=file_root,list_file=list_f,train=True,transform = [transforms.ToTensor()])train_loader = DataLoader(train_dataset,batch_size=1,shuffle=False,num_workers=0)train_iter = iter(train_loader)# for i in range(5):#     img,target = next(train_iter)#     print(target[target[...,0]>0])for i,(images,target) in enumerate(train_loader):print(1111111111111111111111)print(target)print(images)

（2）以下是数据增强中主要函数的讲解。数据增强中，在对图片操作的同时，关键是也要对真实框作相应的操作。例如对图像进行翻转，目标物体位置发生变化，真实框也要进行同样的翻转，利于保障训练预测结果的准确性。

  1. 编码要点

I.编码中输入的框是框在图像中的相对位置，即真实框的左上角和右下角坐标除以图像宽高后的值。因为要在找到特征图上真实框的位置，虽然原图和特征图的尺寸不一样，但是真实框在原图和特征图的相对位置一样，通过相对位置把真实框映射在特征图上。

img.shape:[400,500,3]         # 原图 h:400;w:500
box:[100,120,200,250]         # 真实框的坐标 [x1,y1,x2,y2]
box_img: [100/500，120/400,200/500,250/400]            # 真实框在原图上的相对位置 [x1/w,y1/h,x2/w,y2/h]
feature.shape:[7,7,30]        # 特征图
box_feature: [100/500*7，120/400*7,200/500*7,250/400*7] # 真实框在特征图上的位置

II. dxy

cxcy_sample = cxcy[i]          # 真实框相对原图的中心点
ij = (cxcy_sample*grid_num).ceil()-1  # 真实框在特征图上的中心点，为了防止中心点越界，因此中心点坐标向下取整并减一。
dxy = cxcy_sample*grid_num-ij  # 偏移

III. target.shape[7,7,30]，前10个数据是两个框和置信度数据，后面20个是类别的one_hot形式。

dim = 2
label = 5
0   1  2  3    4    5   6   7  8    9     10  11  12  13  14 ...29
dx  dy w  h  conf   dx  dy  w  h  conf     0   0   0   0   1 ... 0

编码流程：

I. 真实框左上角、右下角[x1,y1,x2,y2]坐标转换为中心点和宽高[cx,cy,w,h]。

II. 遍历图像中真实框，计算真实框在特征图上的中心点坐标、中心点坐标偏移。把框、置信度、类别信息填写在标签target的对应位置上。


def encoder(self,boxes,labels):grid_num = 7     # 特征图边长，feature.shape:[7,7,30]target = torch.zeros((grid_num,grid_num,30))   # 标签wh = boxes[:,2:] - boxes[:,:2]         # 真实框相对原图的宽高cxcy = (boxes[:,2:] + boxes[:,:2])/2   # 真实框相对原图的中心点for i in range(cxcy.size()[0]):  # 遍历每个框cxcy_sample = cxcy[i]         ij = (cxcy_sample*grid_num).ceil()-1  # 真实框在特征图上的中心点dxy = cxcy_sample*grid_num-ij     # 中心偏移target[int(ij[1]),int(ij[0]),:2] = target[int(ij[1]),int(ij[0]),5:7] = dxy   target[int(ij[1]),int(ij[0]),2:4] = target[int(ij[1]),int(ij[0]),7:9] = wh[i]target[int(ij[1]),int(ij[0]),4] = target[int(ij[1]),int(ij[0]),9] = 1  # 置信度target[int(ij[1]),int(ij[0]),int(labels[i])+9] = 1   # 类别return target
  1. 随机调节亮度
def RandomBrightness(self,bgr):if random.random()<0.5:        # 随机值hsv = self.BGR2HSV(bgr)    # bgr-->hsvh,s,v = cv2.split(hsv)     # 通道分割v = v.astype(float)        # 取出亮度通道，转换数据类型v *= random.choice([0.5,1.5])  # 随机变换亮度值v = np.clip(v,0,255).astype(hsv.dtype)  # 限制数据范围，转换数据类型hsv = cv2.merge((h,s,v))   # 通道合并bgr = self.HVS2BGR(hsv)    # hsv-->bgrreturn  bgr
  1. 随机平移

I. 找到平移距离shift_x，shift_y，平移图片。

II. 判断框平移后是否合理。

III.对框平移。

bgr.shape[100,100,3]
1: shift_x= 20，shift_y= 30,after_shift_imge[30:,20:,:]=bgr[:70,:80,:]向右下角移动
2: shift_x= 20，shift_y=-30,after_shift_imge[:70,20:,:]=bgr[30:,:80,:]向上平移
3: shift_x=-20，shift_y= 30,after_shift_imge[30:,:80,:]=bgr[:70,20:,:]向左平移
4: shift_x=-20，shift_y=-30,after_shift_imge[:70,:80,:]=bgr[30:,20:,:]向左上角移动
def randomShift(self,bgr,boxes,labels):center = (boxes[:,2:]+boxes[:,:2])/2if random.random()<0.5:height,width,c = bgr.shapeafter_shift_imge = np.zeros((height,width,c),dtype=bgr.dtype)after_shift_imge[:,:,:] = (104,117,123)shift_x = random.uniform(-width*0.2,width*0.2)shift_y = random.uniform(-height*0.2,height*0.2)if shift_x>=0 and shift_y>=0:after_shift_imge[int(shift_y):,int(shift_x):,:] = bgr[:height-int(shift_y),:width-int(shift_x),:]elif shift_x>=0 and shift_y<0:after_shift_imge[:height+int(shift_y),int(shift_x):,:] = bgr[-int(shift_y):,:width-int(shift_x),:]elif shift_x <0 and shift_y >=0:after_shift_imge[int(shift_y):,:width+int(shift_x),:] = bgr[:height-int(shift_y),-int(shift_x):,:]elif shift_x<0 and shift_y<0:after_shift_imge[:height+int(shift_y),:width+int(shift_x),:] = bgr[-int(shift_y):,-int(shift_x):,:]shift_xy = torch.FloatTensor([[int(shift_x),int(shift_y)]]).expand_as(center)center = center + shift_xy # 框中心点移动后的位置mask1 = (center[:,0]>0)& (center[:,0]>height)  # mask2 = (center[:,1]>0)& (center[:,1]>width)mask = (mask1 & mask2).view(-1,1) boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)  # 筛选移动后仍存在的框if len(boxes_in) == 0:return bgr,boxes,labelsbox_shift = torch.FloatTensor([[int(shift_x),int(shift_y),int(shift_x),int(shift_y)]]).expand_as(boxes_in)boxes_in = boxes_in+box_shiftlabels_in = labels[mask.view(-1)]   # 筛选移动后仍存在的标签return after_shift_imge,boxes_in,labels_inreturn bgr,boxes,labels
  1. 随机缩放
def randomScale(self,bgr,boxes):if random.random() < 0.5:scale = random.uniform(0.8,1.2)  # 缩放因子h,w,c = bgr.shapebgr = cv2.resize(bgr,(int(w*scale),h)) # 缩放图片宽，高度不变scale_tensor = torch.FloatTensor([[scale,1,scale,1]]).expand_as(boxes)boxes = boxes*scale_tensor    # 缩放框的宽度return bgr,boxesreturn bgr,boxes
  1. 随机剪切

I.确定随机剪切的宽高hw，根据宽高确定剪切起始点xy.

II.根据框的中心点判断框是否在剪切区间。

III.对在剪切区间的框移动。返回剪切后的图片、框、标签。

def randomCrop(self,bgr,boxes,labels):if random.random() < 0.5:center = (boxes[:,:2]+boxes[:,2:])/2height,width,c = bgr.shapeh = random.uniform(0.6*height,height)w = random.uniform(0.6*width,width)x = random.uniform(0,width-w)y = random.uniform(0,height-h)x,y,h,w = int(x),int(y),int(h),int(w)center = center - torch.FloatTensor([[x,y]]).expand_as(center)mask1 = (center[:,0]>0) & (center[:,0]<w)mask2 = (center[:,1]>0) & (center[:,0]<h)mask = (mask1 & mask2).view(-1,1)boxes_in = boxes[mask.expand_as(boxes)].view(-1,4)if(len(boxes_in)==0):return bgr,boxes,labelsbox_shift = torch.FloatTensor([[x,y,x,y]]).expand_as(boxes_in)boxes_in = boxes_in-box_shiftboxes_in[:,0]=boxes_in[:,0].clamp(0,w)boxes_in[:,2]=boxes_in[:,2].clamp(0,w)boxes_in[:,1]=boxes_in[:,1].clamp(0,h)boxes_in[:,3]=boxes_in[:,3].clamp(0,h)labels_in = labels[mask.view(-1)]img_croped = bgr[y:y+h,x:x+w,:]return img_croped,boxes_in,labels_inreturn bgr,boxes,labels
  1. 随机左右翻转

图片左右翻转后，框距左右边界翻转。

def random_flip(self,im,boxes):if random.random() < 0.5:im_lr = np.fliplr(im).copy()h,w,_ = im.shapexmin = w - boxes[:,2]   # 框距右边边界的距离xmax = w - boxes[:,0]   # 框距左边边界的距离# 调换左右边界距离就是翻转后的框boxes[:,0] = xmin       boxes[:,2] = xmaxreturn im_lr,boxesreturn im,boxes

2 训练

论文中的主干模型是由24层卷积和两个全联接组成。代码中
训练主要包括backbone(ResNet、VGG)、LOSS、代入数据训练模型。作者在ImageNet数据集上训练模型，达到top5上达到 88%的准确率。

2.1 Backbone

（1）ResNet
ResNet主要分为三部分。首先通过卷积和池化进行两次步长为2的下采样，然后通过残差模块layer1～layer5扩展通道数和三次步长为2的下采样，最后一次卷积、批归一化、激活得到特征图[7,7,30]。

resnet_yolo.py

import math
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
import torch.nn.functional as F__all__ = ['ResNet','resnet18','resnet34','resnet50','resnet101','resnet152']model_urls = {'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth','resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth','resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth','resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth','resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}def conv3x3(in_planes,out_planes,stride=1):return nn.Conv2d(in_planes,out_planes,kernel_size=3,stride=stride,padding=1,bias=False)class BasicBlock(nn.Module):expansion = 1def __init__(self,inplanes,planes,stride=1,downsample=None):super(BasicBlock, self).__init__()self.conv1 = conv3x3(inplanes, planes, stride)self.bn1 = nn.BatchNorm2d(planes)self.relu = nn.ReLU(inplace=True)self.conv2 = conv3x3(planes, planes)self.bn2 = nn.BatchNorm2d(planes)self.downsample = downsampleself.stride = stridedef forward(self,x):residual = xout = self.conv1(x)out = self.bn1(out)out = self.relu(out)out = self.conv2(out)out = self.bn2(out)if self.downsample is not None:residual = self.downsample(x)out += residualout = self.relu(out)return outclass Bottleneck(nn.Module):expansion = 4def __init__(self, inplanes, planes, stride=1, downsample=None):super(Bottleneck, self).__init__()self.conv1 = nn.Conv2d(inplanes,planes,kernel_size=1,bias=False)self.bn1 = nn.BatchNorm2d(planes)self.conv2 = nn.Conv2d(planes,planes,kernel_size=3,stride=stride,padding=1,bias=False)self.bn2 = nn.BatchNorm2d(planes)self.conv3 = nn.Conv2d(inplanes, planes * self.expansion, kernel_size=1,bias=False)self.bn3 = nn.BatchNorm2d(planes* self.expansion)self.relu = nn.ReLU(inplace=True)self.downsample = downsampleself.stride = stridedef forward(self,x):residual = xout = self.conv1(x)out = self.bn1(out)out = self.relu(out)out = self.conv2(out)out = self.bn2(out)out = self.relu(out)out = self.conv3(out)out = self.bn3(out)if self.downsample is not None:residual = self.downsample(x)out += residualout = self.relu(out)return outclass detnet_bottleneck(nn.Module):# no expansion# dilation = 2# type B use 1x1 convexpansion = 1def __init__(self, in_planes, planes, stride=1, block_type='A'):super(detnet_bottleneck, self).__init__()self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)self.bn1 = nn.BatchNorm2d(planes)self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=2, bias=False,dilation=2)self.bn2 = nn.BatchNorm2d(planes)self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)self.bn3 = nn.BatchNorm2d(self.expansion*planes)self.downsample = nn.Sequential()if stride != 1 or in_planes != self.expansion*planes or block_type == 'B':self.downsample = nn.Sequential(nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),nn.BatchNorm2d(self.expansion*planes))def forward(self, x):out = F.relu(self.bn1(self.conv1(x)))out = F.relu(self.bn2(self.conv2(out)))out = self.bn3(self.conv3(out))out += self.downsample(x)out = F.relu(out)return outclass ResNet(nn.Module): # model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)def __init__(self, block, layers, num_classes=1470):self.inplanes = 64super(ResNet, self).__init__()self.conv1 = nn.Conv2d(3,64,kernel_size=7,stride=2,padding=3,bias=False)self.bn1 = nn.BatchNorm2d(64)self.relu = nn.ReLU(inplace=True)self.maxpool = nn.MaxPool2d(kernel_size=3,stride=2,padding=1)  # torch.Size([2, 64, 112, 112])self.layer1 = self._make_layer(block,64,layers[0])self.layer2 = self._make_layer(block,128,layers[1],stride=2)self.layer3 = self._make_layer(block,256,layers[2],stride=2)self.layer4 = self._make_layer(block,512,layers[3],stride=2)self.layer5 = self._make_detnet_layer(in_channels=512)   # (in_channels=2048)self.conv_end = nn.Conv2d(256,30,kernel_size=3,stride=1,padding=1,bias=False)self.bn_end = nn.BatchNorm2d(30)for m in self.modules():if isinstance(m,nn.Conv2d):n = m.kernel_size[0]*m.kernel_size[1]*m.out_channelsm.weight.data.normal_(0,math.sqrt(2./n))elif isinstance(m,nn.BatchNorm2d):m.weight.data.fill_(1)m.bias.data.zero_()def _make_layer(self,block, planes, blocks, stride=1): # 64,3downsample = Noneif stride != 1 or self.inplanes != planes * block.expansion:downsample = nn.Sequential(nn.Conv2d(self.inplanes,planes*block.expansion,kernel_size=1,stride=stride,bias=False),nn.BatchNorm2d(planes*block.expansion),)layers = []layers.append(block(self.inplanes,planes,stride,downsample))self.inplanes = planes*block.expansionfor i in range(1,blocks):layers.append(block(self.inplanes,planes))return nn.Sequential(*layers)def _make_detnet_layer(self,in_channels):layers = []layers.append(detnet_bottleneck(in_planes=in_channels, planes=256, block_type='B'))layers.append(detnet_bottleneck(in_planes=256, planes=256, block_type='A'))layers.append(detnet_bottleneck(in_planes=256, planes=256, block_type='A'))return nn.Sequential(*layers)def forward(self,x):x = self.conv1(x)x = self.bn1(x)x = self.relu(x)x = self.maxpool(x)x = self.layer1(x)x = self.layer2(x)x = self.layer3(x)x = self.layer4(x)# print(x.shape)x = self.layer5(x)# print(x.shape)x = self.conv_end(x)x = self.bn_end(x)x = F.sigmoid(x)x = x.permute(0,2,3,1)return xdef resnet18(pretrained=False,**kwargs):model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))return modeldef resnet34(pretrained=False,**kwargs):model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))return modeldef resnet50(pretrained=False,**kwargs):model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))return modeldef resnet101(pretrained=False,**kwargs):model = ResNet(BasicBlock, [3, 4, 23, 3], **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))return modeldef resnet152(pretrained=False,**kwargs):model = ResNet(BasicBlock, [3, 8, 36, 3], **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))return modeldef test():import torchfrom torch.autograd import Variablemodel = resnet18()x = torch.rand(2, 64, 112, 112)x = Variable(x)out = model(x)print(out.shape)if __name__ == '__main__':test()

（2）VGG
VGG 模型分为两大部分，一部分用卷积、池化提取特征，然后通过两次全连接得到特征。
net.py

#encoding:utf-8
import math
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo__all__ = ['VGG','vgg11','vgg11_bn','vgg13','vgg13_bn', 'vgg16','vgg16_bn','vgg19','vgg19_bn']model_urls = {'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth','vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth','vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth','vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth','vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth','vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth','vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth','vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
}class VGG(nn.Module):def __init__(self,features,num_classes=1000,image_size=448):super(VGG,self).__init__()self.features = featuresself.image_size = image_sizeself.classifier = nn.Sequential(nn.Linear(512*7*7,4096),nn.ReLU(True),nn.Dropout(),nn.Linear(4096, 4096),nn.ReLU(True),nn.Dropout(),nn.Linear(4096,1470))self._initialize_weights()def forward(self,x):x = self.features(x)x = x.view(x.size(0),-1)x = self.classifier(x)x = F.sigmoid(x)x = x.view(-1,7,7,30)return xdef _initialize_weights(self):for m in self.modules():if isinstance(m,nn.Conv2d):n = m.kernel_size[0] * m.kernel_size[1] * m.out_channelsm.weight.data.normal_(0,math.sqrt(2./n))if m.bias is not None:m.bias.data.zero_()elif isinstance(m, nn.BatchNorm2d):m.weight.data.fill_(1)m.bias.data.zero_()elif isinstance(m, nn.Linear):m.weight.data.normal_(0, 0.01)m.bias.data.zero_()def make_layers(cfg,batch_norm=False):layers = []in_channels = 3first_flag = Truefor v in cfg:s = 1if (v == 64 and first_flag):s = 2first_flag = Falseif v == 'M':layers += [nn.MaxPool2d(kernel_size=2,stride=2)]else:conv2d = nn.Conv2d(in_channels,v,kernel_size=3,stride=s,padding=1)if batch_norm:layers += [conv2d,nn.BatchNorm2d(v),nn.ReLU(inplace=True)]else:layers += [conv2d,nn.ReLU(inplace=True)]in_channels = vreturn nn.Sequential(*layers)cfg = { 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']}def vgg11(pretrained=False,**kwargs):model = VGG(make_layers(cfg['A']),**kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg11']))return modeldef vgg11_bn(pretrained=False,**kwargs):model  = VGG(make_layers(cfg['A'],batch_norm=True),**kwargs)if pretrained:model.load_state_dict(model_zoo.load_url([vgg11_bn]))return modeldef vgg13(pretrained=False,**kwargs):model = VGG(make_layers(cfg['B']),**kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))return modeldef vgg13_bn(pretrained=False, **kwargs):model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn']))return modeldef vgg16(pretrained=False, **kwargs):model = VGG(make_layers(cfg['D']), **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))return modeldef vgg16_bn(pretrained=False, **kwargs):model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))return modeldef vgg19(pretrained=False, **kwargs):model = VGG(make_layers(cfg['E']), **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))return modeldef vgg19_bn(pretrained=False, **kwargs):model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs)if pretrained:model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn']))return modeldef test():# import torch# from torch.autograd import Variable# model = vgg16()# img = torch.rand(2,3,448,448)# img = Variable(img)# output = model(img)# print(output.size())import torchfrom torch.autograd import Variablemodel = vgg16()img = torch.rand(2,3,448,448)img = Variable(img)output = model(img)print(output.size())if __name__ == '__main__':test()

2.2 Loss

YOLOV1计算损失的特殊性：用MSE计算损失，对框回归的宽高先开方在进行MSE，解决大小物体损失差异过大的问题。对回归和前景分类赋予不同的权重，解决正负样本不均衡问题。
yoloLoss.py
计算损失时，输入的是真实框target_tensor、和解码后的预测框pred_tensor[batch_size,7,7,30].

（1）计算损失流程：

  1. 根据真实框的置信度对target_tensor和pred_tensor取出没有真实框的样本sample_nobj[-1,30],在取出样本的第5列和第10列，用mse计算负样本的损失noobj_loss。
  2. 根据真实框的置信度对target_tensor和pred_tensor取出没有真实框的样本sample_obj[-1,30]。对取出的样本分别在提取target_tensor和pred_tensor的框和物体类别，计算类别损失。
  3. 根据sample_obj，计算预测框和真实框的IuO,根据IuO选出与真实框匹配的样本，计算框的回归损失和正样本的损失。预测正样本的真实值用IuO计算。
  4. 对各种损失加权以平衡正负样本不平衡。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variableclass yoloLoss(nn.Module):def __init__(self,S,B,l_coord,l_noobj):super(yoloLoss, self).__init__()self.S = Sself.B = Bself.l_coord = l_coordself.l_noobj = l_noobjdef compute_iou(self,box1,box2):'''Args:box1[N,4],box2[M,4]Return:iou, sized [N,M].'''N = box1.size()[0]M = box2.size()[0]lt = torch.max(box1[:,:2].unsqueeze(1).expand(N,M,2),box2[:,:2].unsqueeze(0).expand(N,M,2))rd = torch.min(box1[:,2:].unsqueeze(1).expand(N,M,2),box2[:,2:].unsqueeze(0).expand(N,M,2))wh = rd-ltwh[wh<0] = 0inter = wh[...,0] * wh[...,1]area1 = ((box1[:,2]-box1[:,0])*(box1[:,3]-box1[:,1])).unsqueeze(1).expand_as(inter)area2 = ((box2[:,2]-box2[:,0])*(box2[:,3]-box2[:,1])).unsqueeze(0).expand_as(inter)iou = inter/(area1+area2-inter)return ioudef forward(self,pred_tensor,target_tensor):''' pred_tensor[b,S,S,B*5+20] ; target_tensor[b,S,S,30]'''# 1 mask_obj_nobjN = pred_tensor.size(0)coo_mask = target_tensor[...,4] > 0   # 存在物体的mask [batch_size,7,7]noo_mask = target_tensor[...,4] == 0  # 不存在物体的maskcoo_mask = coo_mask.unsqueeze(-1).expand_as(target_tensor) # [b,7,7,30]noo_mask = noo_mask.unsqueeze(-1).expand_as(target_tensor)# 2 nobj lossnoo_pred = pred_tensor[noo_mask].view(-1,30) # 没有物体的预测值# print('noo_mask.shape:',noo_mask.shape)# print('pred_tensor.shape:',pred_tensor.shape)# print('noo_pred.shape:',noo_pred.shape)noo_target = target_tensor[noo_mask].view(-1,30) # 存在物体的预测值noo_pred_c = noo_pred[:,[4,9]].flatten() # 取出预测值中的负样本的置信度noo_target_c = noo_target[:,[4,9]].flatten()  # 取出标签中负样本的置信度noobj_loss = F.mse_loss(noo_pred_c,noo_target_c,size_average=False)  # 计算负样本损失# 3  obj: box , classcoo_pred = pred_tensor[coo_mask].view(-1,30)  # 存在物体的预测值box_pred = coo_pred[:,:10].contiguous().view(-1,5) # 预测框class_pred = coo_pred[:,10:]  # 预测类别coo_target = target_tensor[coo_mask].view(-1,30)  # 存在物体的标签box_target = coo_target[:,:10].contiguous().view(-1,5)  # 真实框class_target = coo_target[:,10:] # 真实类别# 3.1  class lossclass_loss = F.mse_loss(class_pred,class_target,size_average=False) # 类别损失# 4  obj_iou（每个网格上有两个预测框，根据IoU选出与真实框最匹配的预测框计算回归损失和正样本损失）coo_response_mask = torch.ByteTensor(box_target.size()).zero_()# coo_response_mask = torch.tensor(coo_response_mask,dtype=torch.bool)box_target_iou = torch.zeros(box_target.size())for i in range(0,box_target.size(0),2):  # 遍历存在物体的框box1 = box_pred[i:i+2]   # 存在物体的两个预测框box1_xy = Variable(torch.FloatTensor(box1.size()))box1_xy[:,:2] = box1[:,:2] / 14. - 0.5*box1[:,2:4]box1_xy[:,2:4] = box1[:,:2] / 14. + 0.5*box1[:,2:4]box2 = box_target[i].view(-1,5)  # 存在物体的一个真实框box2_xy = Variable(torch.FloatTensor(box2.size()))box2_xy[:,:2] = box2[:,:2] / 14. - 0.5*box2[:,2:4]box2_xy[:,2:4] = box2[:,:2] / 14. + 0.5*box2[:,2:4]iou = self.compute_iou(box1_xy[:,:4],box2_xy[:,:4])max_iou,max_index = iou.max(0)  # 计算预测框和真实框的IoU，并返回最有的IoU和预测框的下标coo_response_mask[i+max_index] = 1box_target_iou[i+max_index,4] = max_ioubox_target_iou = Variable(box_target_iou)# 4.1 obj_lossbox_pred_response = box_pred[coo_response_mask].view(-1,5) # 与真实框最匹配的预测框box_target_response = box_target[coo_response_mask].view(-1,5)  # 真是框，这一步多余。box_target_response_iou = box_target_iou[coo_response_mask].view(-1,5)  # 正样本的概率# 4.1.1 contain_losscontain_loss = F.mse_loss(box_pred_response[:,4],box_target_response_iou[:,4],size_average = False)   # 正样本损失# 4.1.2 loc_lossloc_loss = F.mse_loss(box_pred_response[:,:2],box_target_response[:,:2],size_average = False)+ \F.mse_loss(torch.sqrt(box_pred_response[:,2:]),torch.sqrt(box_target_response[:,2:]),size_average = False)  # 框的回归损失return (self.l_noobj*noobj_loss + class_loss + 2*contain_loss + self.l_coord*loc_loss)/N   加权平均损失if __name__ == '__main__':pred_tensor = torch.randn(2,14,14,30)target_tensor =  pred_tensor+0.01yolo_loss = yoloLoss(14,8,5,0.5)loss = yolo_loss(pred_tensor,target_tensor)print(loss)

（2）IoU
I. 计算框相交部分的左上角和右下角坐标lt,rd。
II. 计算交集面积inter和相交框的各自面积area1、area2。
III.根据以上步骤计算交并比iou。

def compute_iou(self,box1,box2):'''Args:box1[N,4],box2[M,4]Return:iou, sized [N,M].'''N = box1.size()[0]M = box2.size()[0]lt = torch.max(box1[:,:2].unsqueeze(1).expand(N,M,2),  # box1.shape[N,4]-->box1[:,:2].shape[N,2]-->box1[:,:2].unsqueeze(1).shape[N,1,2]-->lt.shape[N,M,2]box2[:,:2].unsqueeze(0).expand(N,M,2))rd = torch.min(box1[:,2:].unsqueeze(1).expand(N,M,2),box2[:,2:].unsqueeze(0).expand(N,M,2))wh = rd-lt     # wh.shape(N,M,2)wh[wh<0] = 0inter = wh[...,0] * wh[...,1]  # [N,M]area1 = ((box1[:,2]-box1[:,0])*(box1[:,3]-box1[:,1])).unsqueeze(1).expand_as(inter)  # area1.shape[N,M]area2 = ((box2[:,2]-box2[:,0])*(box2[:,3]-box2[:,1])).unsqueeze(0).expand_as(inter)iou = inter/(area1+area2-inter)   # iou.shape[N,M]return iou

2.3 train

train.py
训练流程:

  1. 导入库
  2. 设置超参数
  3. 模型
  4. 导入模型参数
  5. 损失函数
  6. 设置优化器
  7. 导入数据
  8. 训练
# 1 导入库
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision import models
from torch.autograd import Variablefrom net import vgg16_bn
from resnet_yolo import resnet50
from yoloLoss import yoloLoss
from data_label_13 import yoloDataset
from visualize import Visualizer
import numpy as np# 2 设置参数
# use_gpu = torch.cuda.is_available()
file_root = r'/Users/liushuang/PycharmProjects/YOLOV1_LS'
learning_rate = 0.001
num_epochs = 2
batch_size = 1
use_resnet = False# 3 backbone
if use_resnet:net = resnet50()
else:net = vgg16_bn()
# print(net)# 3.1 导入预训练参数
if use_resnet:resnet = models.resnet50(pretrained=False)  # Truenew_state_dict = resnet.state_dict()dd = net.state_dict()for k in new_state_dict.keys():if k in dd.keys() and not k.startswith('fc'):dd[k] = new_state_dict[k]net.load_state_dict(dd)
else:vgg = models.vgg16_bn(pretrained=False)new_state_dict = vgg.state_dict()dd = net.state_dict()for k in new_state_dict.keys():if k in dd.keys() and k.startswith('features'):dd[k] = new_state_dict[k]net.load_state_dict(dd)
if False:net.load_state_dict(torch.load('best.pth'))# 4 Loss
criterion = yoloLoss(7,2,5,0.5)# if use_gpu:
#     net.cuda()
# 模型训练
net.train()# 5 参数
params = []
params_dict = dict(net.named_parameters())for k,v in params_dict.items():if k.startswith('features'):params += [{'params':[v],'lr':learning_rate*1}]else:params += [{'params':[v],'lr':learning_rate*1}]# 6 优化器
optimizer = torch.optim.SGD(params,lr=learning_rate,momentum=0.9,weight_decay=5e-4)# 7 导入数据train_dataset = yoloDataset(root=file_root,list_file='2007_train.txt',train=True,transform = [transforms.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=0)test_dataset = yoloDataset(root=file_root,list_file='2007_test.txt',train=False,transform = [transforms.ToTensor()] )
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=4)print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
logfile = open('log.txt', 'w')num_iter = 0
vis = Visualizer(env='LS')
best_test_loss = np.inf# 8 训练
for epoch in range(num_epochs):net.train()if epoch == 30:learning_rate = 0.0001if epoch == 40:learning_rate = 0.00001for params_group in optimizer.param_groups:params_group['lr'] = learning_rateprint('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))print('Learning Rate for this epoch: {}'.format(learning_rate))total_loss = 0.for i,(images,target) in enumerate(train_loader):images = Variable(images)target = Variable(target)# if use_gpu:#     images,target = images.cuda(),target.cuda()pred = net(images)# print('pred.shape:',pred.shape)# print('target.shape:',target.shape)loss = criterion(pred,target)total_loss += loss.data.item()optimizer.zero_grad()loss.backward()optimizer.step()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 += 1vis.plot_train_val(loss_train = total_loss/(i+1))validation_loss = 0.0net.eval()for i,(images,target) in enumerate(test_loader):images = Variable(images,volatile=True)target = Variable(target,volatile=True)if use_gpu:images,target = images.cuda(),target.cuda()pred = net(images)loss = criterion(pred,target)validation_loss += loss.data[0]validation_loss /= len(test_loader)vis.plot_train_val(loss_val=validation_loss)if best_test_loss > validation_loss:best_test_loss = validation_lossprint('get best test loss %.5f' % best_test_loss)torch.save(net.state_dict(),'best.pth')logfile.writelines(str(epoch) + '\t' + str(validation_loss) + '\n')logfile.flush()torch.save(net.state_dict(),'yolo.pth')

3 预测

3.1 图片输入模型得到预测结果

  1. 图片输入模型得到预测结果

（1） 预测
predict.py

import torch
from torch.autograd import Variable
from resnet_yolo import resnet50
import torchvision.transforms as transforms
import cv2
import numpy as npVOC_CLASSES = (    # always index 0'aeroplane', 'bicycle', 'bird', 'boat','bottle', 'bus', 'car', 'cat', 'chair','cow', 'diningtable', 'dog', 'horse','motorbike', 'person', 'pottedplant','sheep', 'sofa', 'train', 'tvmonitor')Color = [[0, 0, 0],[128, 0, 0],[0, 128, 0],[128, 128, 0],[0, 0, 128],[128, 0, 128],[0, 128, 128],[128, 128, 128],[64, 0, 0],[192, 0, 0],[64, 128, 0],[192, 128, 0],[64, 0, 128],[192, 0,128],[64, 128, 128],[192, 128, 128],[0, 64, 0],[128, 64, 0],[0, 192, 0],[128, 192, 0],[0, 64, 128]]def nms(bboxes,scores,threshold=0.5):x1 = bboxes[:,0]y1 = bboxes[:,1]x2 = bboxes[:,2]y2 = bboxes[:,3]areas = (x2-x1)*(y2-y1)_,order = scores.sort(0,descending=True)keep = []while order.numel() > 0:if order.numel()>1:i = order[0]else:i = orderkeep.append(i)if order.numel() == 1:breakxx1 = x1[order[1:]].clamp(min=x1[i])yy1 = y1[order[1:]].clamp(min=y1[i])xx2 = x2[order[1:]].clamp(max=x1[i])yy2 = y2[order[1:]].clamp(max=y1[i])w = (xx2-xx1).clamp(min=0)h = (yy2-yy1).clamp(min=0)inter = w*hove = inter/(areas[i]+areas[order[1:]]-inter)#ids = (ove <= threshold).nonzero().squeeze()ids = torch.nonzero(ove <= threshold).squeeze()if ids.numel() == 0:breakorder = order[ids+1]return torch.LongTensor(keep)def decoder(pred):grid_num = 7boxes = []cls_indexs = []probs = []cell_size = 1./grid_numpred = pred.datapred = pred.squeeze(0)  # 7x7x30contain1 = pred[:,:,4].unsqueeze(2)  # [7, 7, 1]contain2 = pred[:,:,9].unsqueeze(2)  # [7, 7, 1]contain = torch.cat((contain1,contain2),2)  # [7, 7, 2]mask1 = contain > 0.1  # [7, 7, 2]mask2 = (contain==contain.max())  # [7, 7, 2]mask = (mask1+mask2).gt(0)  # [7, 7, 2]for i in range(grid_num):for j in range(grid_num):for b in range(2):if mask[i,j,b] == 1:box = pred[i,j,b*5:b*5+4]contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])xy = torch.FloatTensor([j,i])*cell_size #cell左上角  up left of cellbox[:2] = box[:2]*cell_size + xy # return cxcy relative to imagebox_xy = torch.FloatTensor(box.size())#转换成xy形式    convert[cx,cy,w,h] to [x1,y1,x2,y2]box_xy[:2] = box[:2] - 0.5*box[2:]box_xy[2:] = box[:2] + 0.5*box[2:]max_prob,cls_index = torch.max(pred[i,j,10:],0)if float((contain_prob*max_prob)[0]) > 0.1:boxes.append(box_xy.view(1,4))cls_indexs.append(torch.LongTensor(cls_index,0))probs.append(contain_prob*max_prob)if len(boxes) == 0:boxes = torch.zeros((1,4))probs = torch.zeros(1)cls_indexs = torch.zeros(1)else:boxes = torch.cat(boxes,0) #(n,4)probs = torch.cat(probs,0) #(n,)cls_indexs = torch.cat(cls_indexs,0) #(n,)keep = nms(boxes,probs)return boxes[keep],cls_indexs[keep],probs[keep]def predict_gpu(model,image_name,root_path='/Users/ls/PycharmProjects/YOLOV1_LS/VOCdevkit/VOC2007/JPEGImages/'):result = []image = cv2.imread(root_path+image_name)# 1 图片预处理h,w,_ = image.shapeimg = cv2.resize(image,(448,448))  # 统一输入模型的图片尺寸img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)  # 色彩空间转换mean = (123,117,104)  #RGBimg = img - np.array(mean,dtype=np.float32)  # 去均值transform = transforms.Compose([transforms.ToTensor(),])img = transform(img)  # 转置img = Variable(img[None,:,:,:],volatile=True)# img = img.cuda()# 2 预测pred = model(img) #1x7x7x30   pred = pred.cpu()# 3 解码boxes,cls_indexs,probs =  decoder(pred)for i,box in enumerate(boxes):x1 = int(box[0]*w)x2 = int(box[2]*w)y1 = int(box[1]*h)y2 = int(box[3]*h)cls_index = cls_indexs[i]if cls_index.numel()==0:returncls_index = int(cls_index) # convert LongTensor to intprob = probs[i]prob = float(prob)result.append([(x1,y1),(x2,y2),VOC_CLASSES[cls_index],image_name,prob])return resultif __name__ == '__main__':model = resnet50()print('load model...')# model.load_state_dict(torch.load('best.pth'))model.eval()#model.cuda()image_name = '000015.jpg'image = cv2.imread(image_name)print('predicting...')result = predict_gpu(model,image_name)#  4 画框for left_up,right_bottom,class_name, _ ,prob in result:color = Color[VOC_CLASSES.index(class_name)]cv2.rectangle(image,left_up,right_bottom,color,2)label = class_name+str(round(prob,2))text_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)p1 = (left_up[0], left_up[1]- text_size[1])cv2.rectangle(image, (p1[0] - 2//2, p1[1] - 2 - baseline), (p1[0] + text_size[0], p1[1] + text_size[1]), color, -1)cv2.putText(image, label, (p1[0], p1[1] + baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255,255,255), 1, 8)cv2.imwrite('result.jpg',image)

（2）NMS

NMS的目的是根据IoU删除重复的预测框，原理是同一物体的预测框IoU 大，不同物体的预测框IoU小。

I. 计算预测框面积。

II. 对置信度降序排列，返回排序下标。把排序好的第一个概率大的预测框取出，计算剩余框与第一个框的IoU。

III. 根据阈值筛选保留的框，并对保留的框做同样的操作，直到剩余一个框，则停止操作。


def nms(bboxes,scores,threshold=0.5):x1 = bboxes[:,0]y1 = bboxes[:,1]x2 = bboxes[:,2]y2 = bboxes[:,3]# 1 计算所有预测框的面积areas = (x2-x1)*(y2-y1)# 2 按照预测概率排序_,order = scores.sort(0,descending=True)keep = []while order.numel() > 0:# 3 取出第一个框if order.numel()>1:i = order[0]else:i = orderkeep.append(i)if order.numel() == 1:break# 4 计算剩余框与第一个框的IoUxx1 = x1[order[1:]].clamp(min=x1[i])yy1 = y1[order[1:]].clamp(min=y1[i])xx2 = x2[order[1:]].clamp(max=x1[i])yy2 = y2[order[1:]].clamp(max=y1[i])w = (xx2-xx1).clamp(min=0)h = (yy2-yy1).clamp(min=0)inter = w*hove = inter/(areas[i]+areas[order[1:]]-inter)#ids = (ove <= threshold).nonzero().squeeze()# 5 根据IoU剔除重合的框ids = torch.nonzero(ove <= threshold).squeeze()if ids.numel() == 0:break# 6 取出与第一个框重叠小或不重叠的框作为下一轮筛选的对象order = order[ids+1]return torch.LongTensor(keep)

（3）解码

I. 取出预测值中的置信度，根据置信度阈值初筛预测框。

II.遍历输出特征图的行、列、每个网格的框，取出对应的预测框、类别概率。根据预测偏移计算预测框的中心点。预测类别概率与置信度乘积作为最终的预测概率，再根据最终的预测概率设置阈值帅选一遍框。

III.根据预测框和物体类别概率进行非极大值抑制，输出符合条件的预测值。

def decoder(pred):grid_num = 7boxes = []cls_indexs = []probs = []cell_size = 1./grid_numpred = pred.datapred = pred.squeeze(0)  # 7x7x30contain1 = pred[:,:,4].unsqueeze(2)  # [7, 7, 1]contain2 = pred[:,:,9].unsqueeze(2)  # [7, 7, 1]# 1 根据置信度筛选框contain = torch.cat((contain1,contain2),2)  # [7, 7, 2]mask1 = contain > 0.1  # [7, 7, 2]mask2 = (contain==contain.max())  # [7, 7, 2]mask = (mask1+mask2).gt(0)  # [7, 7, 2]for i in range(grid_num):for j in range(grid_num):for b in range(2):if mask[i,j,b] == 1:box = pred[i,j,b*5:b*5+4]contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])xy = torch.FloatTensor([j,i])*cell_size #cell左上角  up left of cell# 2 解码box[:2] = box[:2]*cell_size + xy # return cxcy relative to imagebox_xy = torch.FloatTensor(box.size())#转换成xy形式    convert[cx,cy,w,h] to [x1,y1,x2,y2]box_xy[:2] = box[:2] - 0.5*box[2:]box_xy[2:] = box[:2] + 0.5*box[2:]max_prob,cls_index = torch.max(pred[i,j,10:],0)# 3 根据最终预测概率筛选框if float((contain_prob*max_prob)[0]) > 0.1:boxes.append(box_xy.view(1,4))cls_indexs.append(torch.LongTensor(cls_index,0))probs.append(contain_prob*max_prob)if len(boxes) == 0:boxes = torch.zeros((1,4))probs = torch.zeros(1)cls_indexs = torch.zeros(1)else:boxes = torch.cat(boxes,0) #(n,4)probs = torch.cat(probs,0) #(n,)cls_indexs = torch.cat(cls_indexs,0) #(n,)# 4 非极大值抑制keep = nms(boxes,probs)return boxes[keep],cls_indexs[keep],probs[keep]

YOLOV1详解——Pytorch版相关推荐

  1. iOS核心动画详解swift版----基础动画

    2019独角兽企业重金招聘Python工程师标准>>> iOS核心动画详解swift版---基础动画 创建工程,添加2个ViewController,通过rootViewContro ...

  2. python中squeeze函数_详解pytorch中squeeze()和unsqueeze()函数介绍

    squeeze的用法主要就是对数据的维度进行压缩或者解压. 先看torch.squeeze() 这个函数主要对数据的维度进行压缩,去掉维数为1的的维度,比如是一行或者一列这种,一个一行三列(1,3)的 ...

  3. Web IDE优势在哪？详解Web版数据库管理工具SQL Studio

    Web IDE优势在哪?详解Web版数据库管理工具SQL Studio 去年年末,GitPod在A轮融资中获得2500万美元,用来兑现云端开发环境(Cloud Development Environm ...

  4. matlab guide图像处理实例,现代数字图像处理技术提高及应用案例详解 MATLAB版-教科书.pdf...

    数字,图像处理,技术,提高,应用,案例,详解,教科书 [General Information] 书名=现代数字图像处理技术提高及应用案例详解 MATLAB版 作者=赵小川编著 页数=384 SS号 ...

  5. 超详解pytorch实战Kaggle比赛：房价预测

    详解pytorch实战Kaggle比赛:房价预测 教程名称 教程地址 机器学习/深度学习 [李宏毅]机器学习/深度学习国语教程(双语字幕) 生成对抗网络 [李宏毅]生成对抗网络国语教程(双语字幕) 目 ...

  6. python 动态执行 内存变化_详解Pytorch显存动态分配规律探索

    下面通过实验来探索Pytorch分配显存的方式. 实验显存到主存 我使用VSCode的jupyter来进行实验,首先只导入pytorch,代码如下: import torch 打开任务管理器查看主存与 ...

  7. 详解pytorch实现猫狗识别98%附代码

    详解pytorch实现猫狗识别98%附代码 前言 一.为什么选用pytorch这个框架? 二.实现效果 三.神经网络从头到尾 1.来源:仿照人为处理图片的流程,模拟人们的神经元处理信息的方式 2.总览 ...

  8. unity导出面板详解(pc版)

    unity导出面板详解(pc版) 目录 Resolution and Presentation(分辨率和显示) Fullscreen Mode(全屏模式) 全屏模式 参数 功能 Fullscreen ...

  9. （！详解 Pytorch实战：①）kaggle猫狗数据集二分类：加载（集成/自定义）数据集

    这系列的文章是我对Pytorch入门之后的一个总结,特别是对数据集生成加载这一块加强学习 另外,这里有一些比较常用的数据集,大家可以进行下载: 需要注意的是,本篇文章使用的PyTorch的版本是v0. ...

最新文章

  1. impala的substr从第几位截到最后一位_【题解循环嵌套】书的最后一页（粉丝求助）...
  2. LeetCode 443. 压缩字符串（双指针）
  3. 【数据结构与算法】哈夫曼树的Java实现
  4. 如何在 Intellij IDEA 更高效地将应用部署到容器服务 Kubernetes
  5. JAVA编码（5）——JAVA输入流输出流
  6. springboot项目启动成功后执行一段代码的两种方式
  7. Linux时间子系统之七：定时器的应用--msleep()，hrtimer_nanosleep()
  8. 菜鸟html代码在线,html菜鸟教程,HTML新手如何快速入门
  9. python第六周项目答案_Python语言程序设计-课后练习-第6周.pdf
  10. 图形图像处理案例4——彩色勾线画生成器，彩色绕线画生成器，绕线画生成器
  11. 关于微信无法登陆网页版的问题
  12. Golang interface 接口详解
  13. 迅雷极速版服务器响应,如何阻止迅雷极速版强制更新？这个方法很简单
  14. 无聊吗？写个【飞机大战】来玩吧（下篇）
  15. javaSE (十六）Randon、System类的方法(生成随机数、主动垃圾回收、终止jvm、对程序进行记时、System自带的复制数组的方法）
  16. 大学计算机word图文混排,Word 2003从入门到精通第五讲（图文混排）讲稿
  17. js 去除最后一个逗号
  18. 上班族致富五大途径解密
  19. Beta冲刺-第四天
  20. python flask 微信小程序_python-flask微信小程序搭建

热门文章

  1. 共享单车神经网络预测（pytorch )每行代码详细解释
  2. 有些腹肌不必撕裂，是吃出来的
  3. 苹果手机之间怎么传照片_相机与手机之间传送RAW格式照片问题！
  4. 李嘉诚在长江商学院毕业典礼上的致词《强者的有为》
  5. eclipse字体和保护色的设置
  6. 电脑开机时嘀嘀嘀声音代表含义
  7. 前后端分离--整套解决方案
  8. 逃离同质化，OPPO折叠屏正在笃定远一点的未来
  9. win10卸载java8,win10系统彻底卸载java jdk的操作方法
  10. 满足70%项目需要，葛兰岱尔推出免费商用版3D/BIM/GIS引擎