代码收藏家 YOLOV1（pytorch） 代码讲解
引言
本文通过代码实现了改进的YOLOV1算法，将主干网络替换为了ResNet50其余不变，使用的是pytorch框架，版本为1.7.1。YOLOV1的算法原理本文不做过多阐述。本文针对小白对代码以及相关函数进行讲解，建议配合代码进行阅读，代码中我进行了详细的备注，因此读者可以更加容易理解代码的含义，本文只展示了部分代码，全部代码可以通过GitHub下载。
更新：
2022.8.9 predict.py 修复了一些bug
使用的数据集为VOC格式数据集，全部代码可以分为以下几个部分:
- 读取XML文件信息（write_txt.py）
- 数据集预处理（yoloData.py）
- YOLOV1网络结构定义（resnet50.py）
- 损失函数定义（yoloLoss.py）
- 训练网络（train.py）
- 进行预测（predict.py）
write_txt.py
此文件的主要任务是，读取文件夹内所有的xml文件及其信息，然后将其信息（name,bbox,class）写入一个txt文件，再此阶段训练集以及测试集被划分开来这里是按照7：3的比例进行划分，后续的数据集处理需要用到这些信息。
python中有专门的库用于提取xml文件信息名为"ElementTree"，另外还需导入os库用于系统的文件操作。按照个人的理解xml文件分为多层目录，其基本的思路为逐级读取目录，然后从目录内提取信息，以下为常用操作
tree = ET.parse(xml_filename) # 生成一个总目录名为tree
src = tree.findall('source') # 在总目录下找到'source'目录 名为 src
data = obj.find('database') # 若此目录没有下一层目录则用find而不用findall
print(data) # data = The VOC2007 Databasedef parse_rec(filename):
函数的功能:读取xml文件信息，在"object"目录下查看"difficult"值是否为1，若不为1则在名为"obj_struct"的字典中存入"bbox"和"name"的信息，再将这个字典作为名为"objects"的列表的元素，最终输出这个列表。所以这个名为"objects"的列表中的每一个元素都是一个字典。
def parse_rec(filename): # 输入xml文件名
tree = ET.parse(filename)
objects = []
for obj in tree.findall('object'):
obj_struct = {}
difficult = int(obj.find('difficult').text)
if difficult == 1: # 若为1则跳过本次循环
continue
obj_struct['name'] = obj.find('name').text
bbox = obj.find('bndbox')
obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
int(float(bbox.find('ymin').text)),
int(float(bbox.find('xmax').text)),
int(float(bbox.find('ymax').text))]
objects.append(obj_struct)
return objects
def write_txt():
在划分好训练集以及测试集后，生成voctrain.txt以及voctest.txt，并且输入以下信息: 图片文件名（带后缀），bbox的四个角坐标，类名代号（数字0~19），这四个信息，图片内的信息都在一行内输入完成，所以有多少张图片文件就有多少行。
def write_txt():
count = 0
for train_list in train_lists: # 生成训练集txt
count += 1
image_name = train_list.split('.')[0] + '.jpg' # 图片文件名
results = parse_rec(Annotations + train_list)
if len(results) == 0:
print(train_list)
continue
train_set.write(image_name)
for result in results:
class_name = result['name']
bbox = result['bbox']
class_name = VOC_CLASSES.index(class_name)
train_set.write(' ' + str(bbox[0]) +
' ' + str(bbox[1]) +
' ' + str(bbox[2]) +
' ' + str(bbox[3]) +
' ' + str(class_name))
train_set.write('\n')
train_set.close()
for test_list in test_lists: # 生成测试集txt
count += 1
image_name = test_list.split('.')[0] + '.jpg' # 图片文件名
results = parse_rec(Annotations + test_list)
if len(results) == 0:
print(test_list)
continue
test_set.write(image_name)
for result in results:
class_name = result['name']
bbox = result['bbox']
class_name = VOC_CLASSES.index(class_name)
test_set.write(' ' + str(bbox[0]) +
' ' + str(bbox[1]) +
' ' + str(bbox[2]) +
' ' + str(bbox[3]) +
' ' + str(class_name))
test_set.write('\n')
test_set.close()
最终的部分结果如上图所示，一张图片可能含有多个框所以按照顺序依次排列，红色框内的信息上面已经解释过了，以000001.jpg为例48，240，195，371 为bbox信息，11为类别信息按照定义的元组11代表dog，而14代表person,从下图可以证实代码的正确性。
VOC_CLASSES = ( # 定义所有的类名
'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant',
'sheep', 'sofa', 'train', 'tvmonitor')
yoloData.py
此文件主要任务就是根据txt文件内的信息制作ground truth，并且还会进行一定的数据增强。最终输出一个7*7*30的张量。在（train.py中会用到）
以下为文件的结构
class yoloDataset(Dataset):
def __init__(self, img_root, list_file, train, transform)：
def __getitem__(self, idx)：
def __len__(self)：
def encoder(self, boxes, labels): # 输出ground truth （一个7*7*30的张量）
# 以下都是数据增强操作
def random_flip(self, im, boxes): # 随机翻转
def randomScale(self, bgr, boxes): # 随机伸缩变换
def randomBlur(self, bgr): # 随机模糊处理
def RandomBrightness(self, bgr): # 随机调整图片亮度
def randomShift(self, bgr, boxes, labels): # 平移变换首先是 __init__ 部分，大致的操作为：逐行读取生成的文本文件的内容，然后对其进行分类，将信息保存在fnames，boxes，labels三个列表中，以下为需要用到的函数
file_txt = open(list_file) # os模块函数，用于打开文件
lines = file_txt.readlines() # 读取txt文件每一行
splited = line.strip().split()
'''
strip() # 移除首位的换行符号
split() # 以空格为分界线，将所有元素组成一个列表
'''
fnames = splited[0]
x_max = splited[1 + 5*i] # i 表示bbox的个数
y_max = splited[2 + 5*i]
x_min = splited[3 + 5*i]
y_min = splited[4 + 5*i] def __init__(self, img_root, list_file, train, transform): # list_file为txt文件 img_root为图片路径
self.root = img_root
self.train = train
self.transform = transform
# 后续要提取txt文件信息，分类后装入以下三个列表
self.fnames = []
self.boxes = []
self.labels = []
self.S = 7 # YOLOV1
self.B = 2 # 相关
self.C = 20 # 参数
self.mean = (123, 117, 104) # RGB
file_txt = open(list_file)
lines = file_txt.readlines() # 读取txt文件每一行
for line in lines: # 逐行开始操作
splited = line.strip().split() # 移除首位的换行符号再生成一张列表
self.fnames.append(splited[0]) # 存储图片的名字
num_boxes = (len(splited) - 1) // 5 # 每一幅图片里面有多少个bbox
box = []
label = []
for i in range(num_boxes): # bbox四个角的坐标
x = float(splited[1 + 5 * i])
y = float(splited[2 + 5 * i])
x2 = float(splited[3 + 5 * i])
y2 = float(splited[4 + 5 * i])
c = splited[5 + 5 * i] # 代表物体的类别，即是20种物体里面的哪一种
box.append([x, y, x2, y2])
label.append(int(c))
self.boxes.append(torch.Tensor(box))
self.labels.append(torch.LongTensor(label))
self.num_samples = len(self.boxes)
其次是 __getitem__ 部分，此为python类中专有的方法，其功能是按照索引获取值，大致的操作为：先对图片进行一些数据增强，再对坐标进行归一化处理，最后通过 encoder 方法得到ground truth。最后对每一张图片都返回预处理后的 图片以及他的ground truth。
def __getitem__(self, idx):
fname = self.fnames[idx]
img = cv2.imread(os.path.join(self.root + fname))
boxes = self.boxes[idx].clone()
labels = self.labels[idx].clone()
if self.train: # 数据增强里面的各种变换用torch自带的transform是做不到的，因为对图片进行旋转、随即裁剪等会造成bbox的坐标也会发生变化，所以需要自己来定义数据增强
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.shape
boxes /= torch.Tensor([w, h, w, h]).expand_as(boxes) # 坐标归一化处理，为了方便训练
img = self.BGR2RGB(img) # because pytorch pretrained model use RGB
img = self.subMean(img, self.mean) # 减去均值
img = cv2.resize(img, (self.image_size, self.image_size)) # 将所有图片都resize到指定大小
target = self.encoder(boxes, labels) # 将图片标签编码到7x7*30的向量
for t in self.transform:
img = t(img)
return img, target进行数据增强比如对图片进行拉长或者旋转其bbox信息也会随之改变，但是用pytorch自带的数据增强方法并不会改变bbox信息，所以这里都是自己定义的增强方法，在对图片进行变换的过程后也会输出bbox更改后的信息。
boxes /= torch.Tensor([w, h, w, h]).expand_as(boxes)
'''
此处为归一化处理
expand_as(box) # 扩张张量至与box的形状一致
'''最后是encoder方法，它需要输入bbox以及labels信息最终输出一个7*7*30的tensor作为ground truth。值得注意的是bbox的坐标信息中的 x 与 y 含义是左上角顶点坐标相对与中心点坐标（xc,yc）的偏移是一个小于1的数，因此在进行预测的时候还需要进行解码。
def encoder(self, boxes, labels): # 输入的box为归一化形式(X1,X2,Y1,Y2) , 输出ground truth (7*7)
grid_num = 7
target = torch.zeros((grid_num, grid_num, 30))
cell_size = 1. / grid_num # 1/7
wh = boxes[:, 2:] - boxes[:, :2] # wh = [w, h] 1*1
# 物体中心坐标集合
cxcy = (boxes[:, 2:] + boxes[:, :2]) / 2 # 归一化含小数的中心坐标
for i in range(cxcy.size()[0]):
cxcy_sample = cxcy[i] # 中心坐标 1*1
ij = (cxcy_sample / cell_size).ceil() - 1 # 左上角坐标 （7*7)为整数
# 第一个框的置信度
target[int(ij[1]), int(ij[0]), 4] = 1
# 第二个框的置信度
target[int(ij[1]), int(ij[0]), 9] = 1
target[int(ij[1]), int(ij[0]), int(labels[i]) + 9] = 1 # 对应类别的概率设置为1
xy = ij * cell_size # 归一化左上坐标 （1*1）
delta_xy = (cxcy_sample - xy) / cell_size # 中心与左上坐标差值 （7*7）
# 坐标w,h代表了预测的bounding box的width、height相对于整幅图像width,height的比例
target[int(ij[1]), int(ij[0]), 2:4] = wh[i] # w1,h1
target[int(ij[1]), int(ij[0]), :2] = delta_xy # x1,y1
# 每一个网格有两个边框
target[int(ij[1]), int(ij[0]), 7:9] = wh[i] # w2,h2
# 由此可得其实返回的中心坐标其实是相对左上角顶点的偏移，因此在进行预测的时候还需要进行解码
target[int(ij[1]), int(ij[0]), 5:7] = delta_xy # [5,7) 表示x2,y2
return target # (xc,yc) = 7*7 (w,h) = 1*1在此需要注意单位的不同，根据论文描述（xc，yc）是基于7*7网格的，而（w，h）是基于1*1的网格。那为什么在这里的单位是不同的呢？这是因为（xc，yc）虽然单位较大但是与（w，h）的值域是相同的，都是属于（0，1），如果将（xc，yc）的单位改成1*1网格的大小则有可能造成梯度消失的现象。下图截取自论文
7*7为图片的分辨率，而30=5+5+20，而 5 = x+y+w+h+c 。先更根据坐标将两个bbox置信c度全部置1，然后将对应类别的概率置1，最后再把两个bbox的四个坐标输入。最终return target。
# 第一个框的置信度
target[int(ij[1]), int(ij[0]), 4] = 1
# 第二个框的置信度
target[int(ij[1]), int(ij[0]), 9] = 1
target[int(ij[1]), int(ij[0]), int(labels[i]) + 9] = 1 # 对应类别的概率yoloLoss.py
文件的任务如文件名所示，即定义损失函数供训练时使用。本文件结构如下
__init__ 部分没什么好说的，定义一些必要的参数。
forward ：在创建对象时会自动调用，这部分计算全部损失函数，此函数会自动进行前向传播因为此类继承了（nn.Module），所以去需要在对象中传入两个参数（对应train.py的第100行），下面解释以下代码的逻辑：传入的两个参数格式为(batch_size*7*7*30)的张量，前者将图片出入神经网络得到的输出值，后者就是上面制作的target也就是ground truth。需要提取ground truth与pred_target的bbox信息，置信度信息以及类别信息，求取损失函数。可以参考此博客(15条消息) 目标检测模型YOLO-V1损失函数详解_智能算法的博客-CSDN博客。
def forward(self, pred_tensor, target_tensor):
'''
pred_tensor: (tensor) size(batchsize,7,7,30)
target_tensor: (tensor) size(batchsize,7,7,30) --- ground truth
'''
N = pred_tensor.size()[0] # batchsize
coo_mask = target_tensor[:, :, :, 4] > 0 # 具有目标标签的索引值 true batchsize*7*7
noo_mask = target_tensor[:, :, :, 4] == 0 # 不具有目标的标签索引值 false batchsize*7*7
coo_mask = coo_mask.unsqueeze(-1).expand_as(target_tensor) # 得到含物体的坐标等信息,复制粘贴 batchsize*7*7*30
noo_mask = noo_mask.unsqueeze(-1).expand_as(target_tensor) # 得到不含物体的坐标等信息 batchsize*7*7*30
coo_pred = pred_tensor[coo_mask].view(-1, 30) # view类似于reshape
box_pred = coo_pred[:, :10].contiguous().view(-1, 5) # 塑造成X行5列（-1表示自动计算），一个box包含5个值
class_pred = coo_pred[:, 10:] # [n_coord, 20]
coo_target = target_tensor[coo_mask].view(-1, 30)
box_target = coo_target[:, :10].contiguous().view(-1, 5)
class_target = coo_target[:, 10:]
# 不包含物体grid ceil的置信度损失
noo_pred = pred_tensor[noo_mask].view(-1, 30)
noo_target = target_tensor[noo_mask].view(-1, 30)
noo_pred_mask = torch.cuda.ByteTensor(noo_pred.size()).bool()
noo_pred_mask.zero_()
noo_pred_mask[:, 4] = 1
noo_pred_mask[:, 9] = 1
noo_pred_c = noo_pred[noo_pred_mask] # noo pred只需要计算 c 的损失 size[-1,2]
noo_target_c = noo_target[noo_pred_mask]
nooobj_loss = F.mse_loss(noo_pred_c, noo_target_c, size_average=False) # 均方误差
# compute contain obj loss
coo_response_mask = torch.cuda.ByteTensor(box_target.size()).bool() # ByteTensor 构建Byte类型的tensor元素全为0
coo_response_mask.zero_() # 全部元素置False bool:将其元素转变为布尔值
no_coo_response_mask = torch.cuda.ByteTensor(box_target.size()).bool() # ByteTensor 构建Byte类型的tensor元素全为0
no_coo_response_mask.zero_() # 全部元素置False bool:将其元素转变为布尔值
box_target_iou = torch.zeros(box_target.size()).cuda()
# box1 = 预测框 box2 = ground truth
for i in range(0, box_target.size()[0], 2): # box_target.size()[0]：有多少bbox，并且一次取两个bbox
box1 = box_pred[i:i + 2] # 第一个grid ceil对应的两个bbox
box1_xyxy = Variable(torch.FloatTensor(box1.size()))
box1_xyxy[:, :2] = box1[:, :2] / float(self.S) - 0.5 * box1[:, 2:4] # 原本(xc,yc)为7*7 所以要除以7
box1_xyxy[:, 2:4] = box1[:, :2] / float(self.S) + 0.5 * box1[:, 2:4]
box2 = box_target[i].view(-1, 5)
box2_xyxy = Variable(torch.FloatTensor(box2.size()))
box2_xyxy[:, :2] = box2[:, :2] / float(self.S) - 0.5 * box2[:, 2:4]
box2_xyxy[:, 2:4] = box2[:, :2] / float(self.S) + 0.5 * box2[:, 2:4]
iou = self.compute_iou(box1_xyxy[:, :4], box2_xyxy[:, :4])
max_iou, max_index = iou.max(0)
max_index = max_index.data.cuda()
coo_response_mask[i + max_index] = 1 # IOU最大的bbox
no_coo_response_mask[i + 1 - max_index] = 1 # 舍去的bbox
# confidence score = predicted box 与 the ground truth 的 IOU
box_target_iou[i + max_index, torch.LongTensor([4]).cuda()] = max_iou.data.cuda()
box_target_iou = Variable(box_target_iou).cuda()
# 置信度误差（含物体的grid ceil的两个bbox与ground truth的IOU较大的一方）
box_pred_response = box_pred[coo_response_mask].view(-1, 5)
box_target_response_iou = box_target_iou[coo_response_mask].view(-1, 5)
# IOU较小的一方
no_box_pred_response = box_pred[no_coo_response_mask].view(-1, 5)
no_box_target_response_iou = box_target_iou[no_coo_response_mask].view(-1, 5)
no_box_target_response_iou[:, 4] = 0 # 保险起见置0（其实原本就是0）
box_target_response = box_target[coo_response_mask].view(-1, 5)
# 含物体grid ceil中IOU较大的bbox置信度损失
contain_loss = F.mse_loss(box_pred_response[:, 4], box_target_response_iou[:, 4], size_average=False)
# 含物体grid ceil中舍去的bbox损失
no_contain_loss = F.mse_loss(no_box_pred_response[:, 4], no_box_target_response_iou[:, 4], size_average=False)
# bbox坐标损失
loc_loss = F.mse_loss(box_pred_response[:, :2], box_target_response[:, :2], size_average=False) + F.mse_loss(
torch.sqrt(box_pred_response[:, 2:4]), torch.sqrt(box_target_response[:, 2:4]), size_average=False)
# 类别损失
class_loss = F.mse_loss(class_pred, class_target, size_average=False)
return (self.l_coord * loc_loss + contain_loss + self.l_noobj * (nooobj_loss + no_contain_loss) + class_loss) / N这里解释一下部分代码的含义及其用法：
unsqueeze()：增加一个维度
expand_as(tensor)：将原本的张量扩充，一般是将通道数扩充，扩充的部分就是将原来的部分复制粘贴。
coo_pred = pred_tensor[coo_mask].view(-1, 30)：其中coo_mask必须为一个布尔值张量，它与pred_tensor的size一致，若coo_mask的某一位置为True则提取pred_tensor对应位置的信息。而view就相当于resize用于改变张量的形状。
noo_pred_mask = torch.cuda.ByteTensor(noo_pred.size()).bool()：创建一个Byte类型的张量，形状与noo_pred一致，然后将所有的参数转化为布尔值。
noo_pred_mask.zero_()：将参数全部置0，也就是False
compute_iou ：输入两个box，输出IOU。IOU指的时交并比，即（交集/并集）。计算公式： iou = inter / (area1 + area2 - inter)
上文提到过，bbox的四个信息为（xc,yc,w,h），（xc，yc）是基于7*7网格的，而（w，h）是基于1*1的网格。而我们计算IOU需要知道bbox的四个顶点的坐标（相对于grid ceil的坐标而不是坐标轴），通过以下公式进行换算。对应代码中的88，89，92，93行。
计算玩IOU之后需要在一个grid ceil中挑选出IOU最大的bbox用于计算置信度损失，另外一个就舍去。然后再去计算坐标损失，置信度损失（含物体grid ceil和不含物体grid ceil）还有类别损失，最后相加即可。
def compute_iou(self, box1, box2): # box1(2,4) box2(1,4)
N = box1.size(0) # 2
M = box2.size(0) # 1
lt = torch.max( # 返回张量所有元素的最大值
# [N,2] -> [N,1,2] -> [N,M,2]
box1[:, :2].unsqueeze(1).expand(N, M, 2),
# [M,2] -> [1,M,2] -> [N,M,2]
box2[:, :2].unsqueeze(0).expand(N, M, 2),
)
rb = torch.min(
# [N,2] -> [N,1,2] -> [N,M,2]
box1[:, 2:].unsqueeze(1).expand(N, M, 2),
# [M,2] -> [1,M,2] -> [N,M,2]
box2[:, 2:].unsqueeze(0).expand(N, M, 2),
)
wh = rb - lt # [N,M,2]
wh[wh < 0] = 0 # clip at 0
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] 重复面积
area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1]) # [N,]
area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1]) # [M,]
area1 = area1.unsqueeze(1).expand_as(inter) # [N,] -> [N,1] -> [N,M]
area2 = area2.unsqueeze(0).expand_as(inter) # [M,] -> [1,M] -> [N,M]
iou = inter / (area1 + area2 - inter)
return iouRnsNet50.py
此文件就是定义resnet50网络，按照YOLOV1的要求输入图像的尺寸为448*448*3，要求输出7*7*30的张量，而ResNet50可以输入224任意倍数的彩色三通道图片（当然也包括448），所以与YOLOV1算法契合度较高。若输入448*448*3的图片经过ResNet50会得到2048*14*14的张量，所以还需要进行后续处理，当然处理的办法有很多这里分享一种处理办法（详情请见代码的160~167行）。
先说明以下代码的结构，一共分为三个类和一个函数。
首先是 Bottleneck 类 ：定义一个基本快，会根据参数downsample的取值不同可以变为Conv Block和Identity Block。
class Bottleneck(nn.Module): # 定义基本块
def __init__(self, in_channel, out_channel, stride, downsample):
super(Bottleneck, self).__init__()
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.in_channel = in_channel
self.out_channel = out_channel
self.bottleneck = Sequential(
Conv2d(in_channel, out_channel, kernel_size=1, stride=stride[0], padding=0, bias=False),
BatchNorm2d(out_channel),
ReLU(inplace=True),
Conv2d(out_channel, out_channel, kernel_size=3, stride=stride[1], padding=1, bias=False),
BatchNorm2d(out_channel),
ReLU(inplace=True),
Conv2d(out_channel, out_channel * 4, kernel_size=1, stride=stride[2], padding=0, bias=False),
BatchNorm2d(out_channel * 4),
)
if self.downsample is False: # 如果 downsample = True则为Conv_Block 为False为Identity_Block
self.shortcut = Sequential()
else:
self.shortcut = Sequential(
Conv2d(self.in_channel, self.out_channel * 4, kernel_size=1, stride=stride[0], bias=False),
BatchNorm2d(self.out_channel * 4)
)
def forward(self, x):
out = self.bottleneck(x)
out += self.shortcut(x)
out = self.relu(out)
return out其次是 output_net 类 ：配合后续的代码将网络的输出调整为7*7*30的格式
class output_net(nn.Module):
# no expansion
# dilation = 2
# type B use 1x1 conv
expansion = 1
def __init__(self, in_planes, planes, stride=1, block_type='A'):
super(output_net, 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()
self.relu = nn.ReLU(inplace=True)
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 = self.relu(self.bn1(self.conv1(x)))
out = self.relu(self.bn2(self.conv2(out)))
out = self.bn3(self.conv3(out))
out += self.downsample(x)
out = self.relu(out)
return out最后是 ResNet 类 ：这个类内置了两个主要的方法分别是_make_layer以及_make_output_layer，前者定义的ResNet50的主要部分，后者将前者输出的张量转化为YOLOV1损失函数需要的张量格式（7*7*30）
class ResNet50(nn.Module):
def __init__(self, block, num_class = 20):
super(ResNet50, self).__init__()
self.block = block
self.layer0 = Sequential(
Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
BatchNorm2d(64),
ReLU(inplace=True),
MaxPool2d(kernel_size=3, stride=2, padding=1)
)
self.layer1 = self.make_layer(self.block, channel=[64, 64], stride1=[1, 1, 1], stride2=[1, 1, 1], n_re=3)
self.layer2 = self.make_layer(self.block, channel=[256, 128], stride1=[2, 1, 1], stride2=[1, 1, 1], n_re=4)
self.layer3 = self.make_layer(self.block, channel=[512, 256], stride1=[2, 1, 1], stride2=[1, 1, 1], n_re=6)
self.layer4 = self.make_layer(self.block, channel=[1024, 512], stride1=[2, 1, 1], stride2=[1, 1, 1], n_re=3)
self.layer5 = self._make_output_layer(in_channels=2048)
self.avgpool = nn.AvgPool2d(2) # kernel_size = 2 , stride = 2
self.conv_end = nn.Conv2d(256, 30, kernel_size=3, stride=1, padding=1, bias=False)
self.bn_end = nn.BatchNorm2d(30)
def make_layer(self, block, channel, stride1, stride2, n_re):
layers = []
for num_layer in range(0, n_re):
if num_layer == 0:
layers.append(block(channel[0], channel[1], stride1, downsample=True))
else:
layers.append(block(channel[1]*4, channel[1], stride2, downsample=False))
return Sequential(*layers)
def _make_output_layer(self, in_channels):
layers = []
layers.append(
output_net(
in_planes=in_channels,
planes=256,
block_type='B'))
layers.append(
output_net(
in_planes=256,
planes=256,
block_type='A'))
layers.append(
output_net(
in_planes=256,
planes=256,
block_type='A'))
return nn.Sequential(*layers)
def forward(self, x):
# print(x.shape) # 3*448*448
out = self.layer0(x)
# print(out.shape) # 64*112*112
out = self.layer1(out)
# print(out.shape) # 256*112*112
out = self.layer2(out)
# print(out.shape) # 512*56*56
out = self.layer3(out)
# print(out.shape) # 1024*28*28
out = self.layer4(out) # 2048*14*14
out = self.layer5(out) # batch_size*256*14*14
out = self.avgpool(out) # batch_size*256*7*7
out = self.conv_end(out) # batch_size*30*7*7
out = self.bn_end(out)
out = torch.sigmoid(out)
out = out.permute(0, 2, 3, 1) # bitch_size*7*7*30
return out最后通过 resnet50 函数 return网络结构。在后续训练的时候只需要调用这个函数即可。
train.py
1、要将数据集用于训练，首先得将数据集打包为训练集以及测试集，此操作之前已经定义过，为、yoloData.py内的yoloDataset类以及DaraLoader类实现。
train_dataset = yoloDataset(img_root=file_root, list_file='voctrain.txt', train=True, transform=[transforms.ToTensor()])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_dataset = yoloDataset(img_root=file_root, list_file='voctest.txt', train=False, transform=[transforms.ToTensor()])
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
2、根据定义的网络导入权重参数，因之前定义的网络后面一部分与ResNet50的结构略有差异，所以并不能完全使用使用torchvision的models库中的resnet50导入权重参数。需要对其权重参数进行一定的筛选。
权重参数导入方法：自己定义的网络以及models库内的网络各自创建一个对象。接着使用state_dict()导入各自的权重参数。网络结构相同的部分将new_state_dict的值赋给op。但是如果自己定义的网络结构的键值与torch自带的库不一直的话，导入权重参数会稍微麻烦一点。这里给出了一种解决办法，参考代码的40至44行。
注意：
state_dict():返回一个类（若直接输出其格式类似与字典）
enumerate: for循环中经常用到，既可以遍历元素又可以遍历索引。
3、定义损失函数优化器并开始训练，网络结构，损失函数以及train_loader返回的图片以及target都要设置为“CUDA”形式最后通过torch.save保存模型参数。
criterion = yoloLoss(7, 2, 5, 0.5)
criterion = criterion.to(device)
net.train() # 训练前需要加入的语句
optimizer = torch.optim.SGD( # 定义优化器 “随机梯度下降”
params, # net.parameters() 为什么不用这个???
lr=learning_rate,
momentum=0.9, # 即更新的时候在一定程度上保留之前更新的方向 可以在一定程度上增加稳定性，从而学习地更快
weight_decay=5e-4) # L2正则化理论中出现的概念
# torch.multiprocessing.freeze_support() # 多进程相关 猜测是使用多显卡训练需要
for epoch in range(num_epochs):
net.train()
if epoch == 30:
learning_rate = 0.0001
if epoch == 40:
learning_rate = 0.00001
for param_group in optimizer.param_groups: # 其中的元素是2个字典；optimizer.param_groups[0]： 长度为6的字典，包括[‘amsgrad’, ‘params’, ‘lr’, ‘betas’, ‘weight_decay’, ‘eps’]这6个参数；
# optimizer.param_groups[1]： 好像是表示优化器的状态的一个字典；
param_group['lr'] = learning_rate # 更改全部的学习率
print('\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, target = images.cuda(), target.cuda()
pred = net(images)
loss = criterion(pred, target)
total_loss += loss.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.item(), total_loss / (i + 1)))
validation_loss = 0.0
net.eval()
for i, (images, target) in enumerate(test_loader): # 导入dataloader 说明开始训练了 enumerate 建立一个迭代序列
images, target = images.cuda(), target.cuda()
pred = net(images) # 将图片输入
loss = criterion(pred, target)
validation_loss += loss.item() # 累加loss值 （固定搭配）
validation_loss /= len(test_loader) # 计算平均loss
best_test_loss = validation_loss
print('get best test loss %.5f' % best_test_loss)
torch.save(net.state_dict(), 'yolo.pth')
predict.py
此为预测代码，执行此代码可以通过输入图片得到识别的结果。从上文可知，将图片输入神经网络会输出一个7*7*30的张量，我们需要对其进行分析得到很多个bbox包括坐标信息以及对应的物体类别信息，最终进行NMS非极大值抑制对bbox进行筛选得到最终结果。
以下为NMS的步骤：
(1)对于类别1， 从概率最大的bbox F开始，分别判断A、B、C、D、E与F的IOU是否大于设定的阈值。
(2) 假设B、D与F的重叠度超过阈值，那么就扔掉B、D（将其置信度置0），然后保留F。
(3) 从剩下的矩形框A、C、E中，选择概率最大的E，然后判断A、C与E的重叠度，重叠度大于一定的阈值，那么就扔掉；并标记E是我们保留下来的第二个矩形框。
(4) 重复这个过程，找到此类别所有被保留下来的矩形框。
(5) 对于类别2，类别3等等…都要重复以上4个步骤。
首先需要导入模型以及参数，并且设置好有关NMS的两个参数：置信度以及IOU最大值。然后就可以开始预测了。首先需要通过opencv读取图片并且将其resize为448*448的RGB图像，将其进行均值处理后输入神经网络得到7*7*30的张量。
然后运行 decode 方法：因为一个grid ceil只预测一个物体，而一个grid ceil生成两个bbox。这里对grid ceil进行以下操作。
1、选择置信度较高的bbox。
2、选择20种类别概率中的最大者作为这个grid ceil预测的类别。
3、置信度乘以物体类别概率作为物体最终的概率。
最终输入一个7*7*6的张量，7*7代表grid ceil 。6=bbox的4个坐标信息+类别概率+类别代号
最后运行 NMS 方法对bbox进行筛选：因为bbox的4个坐标信息为（xc,yc,w,h）需要将其转化为（x,y,w,h）后才能进行非极大值抑制处理。
class Pred():
def __init__(self, model, img_root):
self.model = model
self.img_root = img_root
def result(self):
img = cv2.imread(self.img_root)
h, w, _ = img.shape
image = cv2.resize(img, (448, 448))
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mean = (123, 117, 104) # RGB
img = img - np.array(mean, dtype=np.float32)
transform = ToTensor()
img = transform(img)
img = img.unsqueeze(0) # 输入要求是4维的
Result = self.model(img) # 1*7*7*30
bbox = self.Decode(Result)
bboxes = self.NMS(bbox) # n*6 bbox坐标是基于7*7网格需要将其转换成448
if len(bboxes) == 0:
print("未识别到任何物体")
print("可调小 confident 以及 iou_con")
print("也可能是由于训练不充分，可在训练时将epoch增大")
for i in range(0, len(bboxes)): # bbox坐标将其转换为原图像的分辨率
bboxes[i][0] = bboxes[i][0] * 64
bboxes[i][1] = bboxes[i][1] * 64
bboxes[i][2] = bboxes[i][2] * 64
bboxes[i][3] = bboxes[i][3] * 64
x1 = bboxes[i][0].item() # 后面加item()是因为画框时输入的数据不可一味tensor类型
x2 = bboxes[i][1].item()
y1 = bboxes[i][2].item()
y2 = bboxes[i][3].item()
class_name = bboxes[i][5].item()
print(x1, x2, y1, y2, VOC_CLASSES[int(class_name)])
cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (144, 144, 255)) # 画框
cv2.imshow('img', image)
cv2.waitKey(0)
def Decode(self, result): # x -> 1**7*30
result = result.squeeze() # 7*7*30
grid_ceil1 = result[:, :, 4].unsqueeze(2) # 7*7*1
grid_ceil2 = result[:, :, 9].unsqueeze(2)
grid_ceil_con = torch.cat((grid_ceil1, grid_ceil2), 2) # 7*7*2
grid_ceil_con, grid_ceil_index = grid_ceil_con.max(2) # 按照第二个维度求最大值 7*7 一个grid ceil两个bbox，两个confidence
class_p, class_index = result[:, :, 10:].max(2) # size -> 7*7 找出单个grid ceil预测的物体类别最大者
class_confidence = class_p * grid_ceil_con # 7*7 真实的类别概率
bbox_info = torch.zeros(7, 7, 6)
for i in range(0, 7):
for j in range(0, 7):
bbox_index = grid_ceil_index[i, j]
bbox_info[i, j, :5] = result[i, j, (bbox_index * 5):(bbox_index+1) * 5] # 删选bbox 0-5 或者5-10
bbox_info[:, :, 4] = class_confidence
bbox_info[:, :, 5] = class_index
return bbox_info # 7*7*6 6 = bbox4个信息+类别概率+类别代号
def NMS(self, bbox, iou_con=iou_con):
for i in range(0, 7):
for j in range(0, 7):
# xc = bbox[i, j, 0] # 目前bbox的四个坐标是以grid ceil的左上角为坐标原点 而且单位不一致
# yc = bbox[i, j, 1] # (xc,yc) 单位= 7*7 (w,h) 单位= 1*1
# w = bbox[i, j, 2] * 7
# h = bbox[i, j, 3] * 7
# Xc = i + xc
# Yc = j + yc
# xmin = Xc - w/2 # 计算bbox四个顶点的坐标（以整张图片的左上角为坐标原点）单位7*7
# xmax = Xc + w/2
# ymin = Yc - h/2
# ymax = Yc + h/2 # 更新bbox参数 xmin and ymin的值有可能小于0
xmin = j + bbox[i, j, 0] - bbox[i, j, 2] * 7 / 2 # xmin
xmax = j + bbox[i, j, 0] + bbox[i, j, 2] * 7 / 2 # xmax
ymin = i + bbox[i, j, 1] - bbox[i, j, 3] * 7 / 2 # ymin
ymax = i + bbox[i, j, 1] + bbox[i, j, 3] * 7 / 2 # ymax
bbox[i, j, 0] = xmin
bbox[i, j, 1] = xmax
bbox[i, j, 2] = ymin
bbox[i, j, 3] = ymax
bbox = bbox.view(-1, 6) # 49*6
bboxes = []
ori_class_index = bbox[:, 5]
class_index, class_order = ori_class_index.sort(dim=0, descending=False)
class_index = class_index.tolist() # 从0开始排序到7
bbox = bbox[class_order, :] # 更改bbox排列顺序
a = 0
for i in range(0, CLASS_NUM):
num = class_index.count(i)
if num == 0:
continue
x = bbox[a:a+num, :] # 提取同一类别的所有信息
score = x[:, 4]
score_index, score_order = score.sort(dim=0, descending=True)
y = x[score_order, :] # 同一种类别按照置信度排序
if y[0, 4] >= confident: # 物体类别的最大置信度大于给定值才能继续删选bbox，否则丢弃全部bbox
for k in range(0, num):
y_score = y[:, 4] # 每一次将置信度置零后都重新进行排序，保证排列顺序依照置信度递减
_, y_score_order = y_score.sort(dim=0, descending=True)
y = y[y_score_order, :]
if y[k, 4] > 0:
area0 = (y[k, 1] - y[k, 0]) * (y[k, 3] - y[k, 2])
for j in range(k+1, num):
area1 = (y[j, 1] - y[j, 0]) * (y[j, 3] - y[j, 2])
x1 = max(y[k, 0], y[j, 0])
x2 = min(y[k, 1], y[j, 1])
y1 = max(y[k, 2], y[j, 2])
y2 = min(y[k, 3], y[j, 3])
w = x2 - x1
h = y2 - y1
if w < 0 or h < 0:
w = 0
h = 0
inter = w * h
iou = inter / (area0 + area1 - inter)
# iou大于一定值则认为两个bbox识别了同一物体删除置信度较小的bbox
# 同时物体类别概率小于一定值则认为不包含物体
if iou >= iou_con or y[j, 4] < confident:
y[j, 4] = 0
for mask in range(0, num):
if y[mask, 4] > 0:
bboxes.append(y[mask])
a = num + a
return bboxes运行结果如下（这里使用了自己的数据集）
总结
本文默认使用VOC数据集，但是也可以使用自己的数据集进行训练，值得注意的时若使用自己的训练集训练由于类别总数可能不是20个所以神经网络的输出不再是7*7*30的张量，而是7*7*（10+类别数）的张量。使用自己的数据集训练需要作的修改方法在github中。
