当前位置:   article > 正文

经典目标检测YOLOV1理论基础的编码实现(通过编码理解原理)

经典目标检测YOLOV1理论基础的编码实现(通过编码理解原理)

目录

正文:

1、前期准备

1.1  创建目录结构

1.2  数据集介绍与下载

1.2.1 数据集介绍

1.2.1 数据集下载

1.3. 数据集处理

1.4. 构建数据加载器 

1.4.1定义初始化方法

1.4.2 定义增强图片方法

1.4.3 定义编码器

1.4.4 完整代码

1.5、定义主函数

2构建网络

2.1 ResNet架构介绍

2.2 构建ResNet

2.2.1 构建Base_Block

__init__ 函数

forward 函数

2.2.2 构建Senior_Block

__init__ 初始化函数

forward 前向传播函数

2.2.3 构建Output_Block

__init__ 初始化函数

forward 前向传播函数

2.2.4 构建ResNet

__init__ 初始化函数

_make_layer 方法

_make_out_layer 方法

forward 方法

初始化模型参数部分

2.3 定义模型

2.4 完整代码

 3、定义损失

4、定义训练过程

5、定义预测方法

6、验证图片

7、完整代码和权重下载
​​​​​​​


       为了更详细的了解Yolo的实现原理,通过YOLO的代码形式进行一步一步说明,在阅读代码的过程中,按照下面的目录的步骤去看代码。但在看代码之前建议先去简单的了解YOLO的原理,在了解之后通过下面的内容进一步验证yolo的实现过程,同时也了解了ResNet的实现过程。

       YOLO的原理参考通道1

       YOLO的原理参考通道2

正文:

1、前期准备

准备好目录结构、数据集和关于YOLOv1的基础认知

1.1  创建目录结构

        自己创建项目目录结构,结构目录如下:

network                    CNN Backbone 存放位置
weights                    权重存放的位置
test_images             测试用的图片
utils                          辅助功能的代码存放位置 

models                    保存模型位置

data                         训练的数据集

1.2  数据集介绍与下载

1.2.1 数据集介绍

       首先了解数据集,对数据集了解后方便对数据进行相应处理。数据集详细介绍直通车:https://blog.csdn.net/qq_41946216/article/details/137683750?spm=1001.2014.3001.5501

1.2.1 数据集下载

       本次采用数据集: VOC2012数据集。

       数据集下载方式一http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

       数据集下载方式二:

     下载并构建VOC2012数据集,从:https://gitee.com/ppov-nuc/pascal-vocdataset_-for_-yolo.git, 下载get_data文件generate_csv.py文件到本地,放到创建的目录结构中,修改get_data中下载的内容和相应路径,然后运行批处理文件get_data,在get_dat中会自动执行generate_csv.py,如下图所示。


1.3. 数据集处理

       在utils目录下创建工具类 generate_txt_file.py,主要用于数据集的划分和解析 Annotations/xxxxx.xml 文件中的类别bbox信息,并将信息存入voctrain.txt和voctest.txt文件,如下图所示:

具体代码:

  1. # author: baiCai
  2. # 1. 导包
  3. from xml.etree import ElementTree as ET
  4. import os
  5. import random
  6. # 2. 定义一些基本的参数
  7. # 定义所有的类名
  8. VOC_CLASSES = (
  9. 'aeroplane', 'bicycle', 'bird', 'boat',
  10. 'bottle', 'bus', 'car', 'cat', 'chair',
  11. 'cow', 'diningtable', 'dog', 'horse',
  12. 'motorbike', 'person', 'pottedplant',
  13. 'sheep', 'sofa', 'train', 'tvmonitor')
  14. '''
  15. 读取所有 xml 文件,存入列表
  16. '''
  17. # 要读取的xml文件路径,记得自己修改路径
  18. Annotations = '../data/VOC2012/Annotations/'
  19. # 列出所有的xml文件
  20. xml_files = os.listdir(Annotations)
  21. # 打乱数据集
  22. random.shuffle(xml_files)
  23. '''
  24. 定义训练集和测试比例
  25. 划分Annotations中的训练集和测试集文件列表
  26. '''
  27. # 训练集数量
  28. train_num = int(len(xml_files) * 0.7)
  29. # 训练列表
  30. train_file_list = xml_files[:train_num]
  31. # 测测试列表
  32. test_file_list = xml_files[train_num:]
  33. '''
  34. 定义 xml 解析后的信息存储路径和写对象
  35. '''
  36. # 训练集和测试集文件名字
  37. train_set_path = './voctrain.txt'
  38. test_set_path = './voctest.txt'
  39. # 3. 定义解析xml文件的函数
  40. '''
  41. 主要解析 xml 获取 类别名字和bbox,如
  42. {'name': 'person','bbox': [174, 101, 349, 351]}
  43. '''
  44. def parse_rec(filename):
  45. # 参数:输入xml文件名
  46. # 创建xml对象
  47. tree = ET.parse(filename)
  48. objects = []
  49. # 迭代读取xml文件中的object节点,即物体信息
  50. for obj in tree.findall('object'):
  51. obj_struct = {}
  52. # difficult属性,即这里不需要那些难判断的对象
  53. difficult = int(obj.find('difficult').text)
  54. if difficult == 1: # 若为1则跳过本次循环
  55. continue
  56. # 开始收集信息
  57. obj_struct['name'] = obj.find('name').text
  58. bbox = obj.find('bndbox')
  59. obj_struct['bbox'] =\
  60. [int(float(bbox.find('xmin').text)),
  61. int(float(bbox.find('ymin').text)),
  62. int(float(bbox.find('xmax').text)),
  63. int(float(bbox.find('ymax').text))]
  64. objects.append(obj_struct)
  65. return objects
  66. # 4. 把信息保存入文件中
  67. def write_txt(file_list,set_path):
  68. # # 生成训练集txt
  69. count = 0
  70. with open(set_path, 'w') as wt:
  71. for xml_file in file_list:
  72. count += 1
  73. # 获取图片名字
  74. image_name = xml_file.split('.')[0] + '.jpg' # 图片文件名
  75. # 对xml_file进行解析
  76. results = parse_rec(Annotations + xml_file)
  77. # 如果返回的对象为空,表示张图片难以检测,因此直接跳过
  78. if len(results) == 0:
  79. print(xml_file)
  80. continue
  81. # 否则,则写入文件中
  82. # 先写入图片名字
  83. wt.write(image_name)
  84. # 接着指定下面写入的格式
  85. for result in results:
  86. class_name = result['name']
  87. bbox = result['bbox']
  88. class_name = VOC_CLASSES.index(class_name) # 名字在类别中是下标位置
  89. wt.write(' ' + str(bbox[0]) +
  90. ' ' + str(bbox[1]) +
  91. ' ' + str(bbox[2]) +
  92. ' ' + str(bbox[3]) +
  93. ' ' + str(class_name))
  94. wt.write('\n')
  95. wt.close()
  96. # 5. 运行
  97. if __name__ == '__main__':
  98. write_txt(train_file_list,train_set_path)
  99. write_txt(test_file_list,test_set_path)

1.4. 构建数据加载器 

      在utils目录下创建工具类 yolo_dataset.py,主要用于数据集的构建,包含初始化、图片增强及归一化等功能。具体代码在下面完整代码中。

1.4.1定义初始化方法

       读取xxxx.xml解析后的文件
       对每行数据(每个图片信息)的所有中心点信息以【x,y,w,h】和标签分别存入box列表和label列表。
       当前图片的边界框和标签信息即box列表和label列表,转换为LongTensor格式添加到对应的boxex列表和labels列表。

1.4.2 定义增强图片方法

       增加方法如下列表:

增加方法名称定义的函数
随机翻转图片和边界框random_flip(img, boxes)
随机缩放图片和边界框randomScale(img, boxes)
随机模糊图片randomBlur(img)
随机调整图片亮度RandomBrightness(img)
随机调整图片色调RandomHue(img)
随机调整图片饱和度RandomSaturation(img)
随机移动图片和边界框randomShift(img, boxes, labels)        
随机裁剪图片和边界框randomCrop(img, boxes, labels)
用于从图像中减去均值subMean(self, bgr, mean)
将BGR图像转换为RGB图像BGR2RGB(self, img)
将BGR图像转换为HSV图像BGR2HSV(self, img)
将HSV图像转换为BGR图像HSV2BGR(self, img)

1.4.3 定义编码器

        定义编码器主要目的是用于将边界框(归一化后的边界框信息)和标签编码为目标张量。

1.4.4 完整代码

  1. # author: baiCai
  2. # 导入所需的库
  3. import os
  4. import random
  5. import numpy as np
  6. import torch
  7. import torchvision.transforms as T
  8. from torch.utils.data import Dataset
  9. import cv2
  10. # 定义一个Yolo数据加载器类
  11. class Yolo_Dataset(Dataset):
  12. # 初始化默认的图片大小
  13. image_size = 448
  14. def __init__(self, root, list_file, train=True, transforms=None):
  15. ''"""
  16. 初始化函数
  17. :param root: 图片的根目录
  18. :param list_file: 包含图片路径和标注信息的txt文件
  19. :param train: 是否为训练集,默认为True
  20. :param transforms: 预处理方法,默认为None
  21. """
  22. # 保存传入的参数
  23. self.root = root
  24. self.train = train
  25. self.transform = transforms
  26. self.fnames = [] # 存储图片文件名
  27. self.boxes = [] # 存储边界框信息
  28. self.labels = [] # 存储标签信息
  29. self.mean = (123, 117, 104) # RGB通道的均值,用于归一化
  30. # 打开并读取txt文件
  31. with open(list_file) as f:
  32. lines = f.readlines() # 读取所有行
  33. # 遍历文件中的每一行
  34. for line in lines:
  35. splited = line.strip().split() # 去除行尾的空格并按空格分割
  36. self.fnames.append(splited[0]) # 添加图片文件名到列表中
  37. # 计算当前行包含的对象数量,每五个数据表示一个对象
  38. num_boxes = (len(splited) - 1) // 5
  39. # 初始化当前图片的边界框和标签列表
  40. box = []
  41. label = []
  42. # 遍历当前行的所有对象
  43. for i in range(num_boxes):
  44. # 读取边界框的坐标信息
  45. x = float(splited[1 + 5 * i])
  46. y = float(splited[2 + 5 * i])
  47. x2 = float(splited[3 + 5 * i])
  48. y2 = float(splited[4 + 5 * i])
  49. # 读取对象的类别标签
  50. c = splited[5 + 5 * i]
  51. # 将边界框坐标添加到列表中,并转换为Tensor格式
  52. box.append([x, y, x2, y2])
  53. # 将标签转换为整数并加1(因为标签通常从1开始计数),然后添加到列表中
  54. label.append(int(c) + 1)
  55. # 将当前图片的边界框和标签信息转换为LongTensor格式添加到对应的列表中
  56. self.boxes.append(torch.Tensor(box))
  57. self.labels.append(torch.LongTensor(label))
  58. # 记录数据集中的样本数量
  59. self.num_samples = len(self.boxes)
  60. def __len__(self):
  61. # 返回数据集中的样本数量
  62. return self.num_samples
  63. # 如果是训练模式,需要进行图像的增强,idx表示每个图像在fnames列表中索引
  64. def __getitem__(self, idx):
  65. # 根据索引获取图片的文件名
  66. fname = self.fnames[idx]
  67. # 拼接完整的图片路径并读取图片
  68. img = cv2.imread(os.path.join(self.root, fname))
  69. # 获取当前图片的边界框和标签信息
  70. boxes = self.boxes[idx].clone()
  71. labels = self.labels[idx].clone()
  72. # 如果是训练模式,需要进行图像的增强
  73. # 需要注意的是,同时处理图像和box
  74. if self.train:
  75. # 随机翻转图片和边界框
  76. img, boxes = self.random_flip(img, boxes)
  77. # 随机缩放图片和边界框
  78. img, boxes = self.randomScale(img, boxes)
  79. # 随机模糊图片
  80. img = self.randomBlur(img)
  81. # 随机调整图片亮度
  82. img = self.RandomBrightness(img)
  83. # 随机调整图片色调
  84. img = self.RandomHue(img)
  85. # 随机调整图片饱和度
  86. img = self.RandomSaturation(img)
  87. # 随机移动图片和边界框
  88. img, boxes, labels = self.randomShift(img, boxes, labels)
  89. # 随机裁剪图片和边界框
  90. img, boxes, labels = self.randomCrop(img, boxes, labels)
  91. # 获取图片的高度、宽度和通道数
  92. h, w, _ = img.shape
  93. # 将边界框坐标除以图片宽高,进行归一化
  94. boxes /= torch.Tensor([w, h, w, h]).expand_as(boxes)
  95. # 由于cv2读取图片为BGR格式,需要转换为RGB格式
  96. img = self.BGR2RGB(img)
  97. # 从图片中减去RGB通道的均值
  98. img = self.subMean(img, self.mean)
  99. # 将图片缩放到指定的大小,这里是448x448
  100. img = cv2.resize(img, (self.image_size, self.image_size))
  101. # 将边界框和标签信息编码为YOLOv1需要的格式,即7x7x30的张量
  102. target = self.encoder(boxes, labels)
  103. # 应用预定义的数据变换
  104. for t in self.transform:
  105. img = t(img)
  106. # 返回处理后的图片和编码后的目标信息
  107. return img, target
  108. # 定义一个encoder方法,用于将边界框和标签编码为目标张量
  109. def encoder(self, boxes, labels):
  110. '''
  111. 将边界框和标签编码为7x7x30的目标张量
  112. 参数:
  113. boxes (tensor): 形状为[[x1,y1,x2,y2],[]]的边界框张量
  114. labels (tensor): 标签张量
  115. 返回:
  116. target (tensor): 形状为7x7x30的目标张量
  117. '''
  118. grid_num = 7 # 定义网格数量
  119. # 创建一个形状为(grid_num, grid_num, 30)的全零张量作为目标张量
  120. target = torch.zeros((grid_num, grid_num, 30))
  121. # 计算每个网格的缩放因子
  122. cell_size = 1. / grid_num
  123. print(cell_size)
  124. print(boxes)
  125. # 计算边界框的宽度、高度和中心点坐标,boxes 是归一化后数据
  126. wh = boxes[:, 2:] - boxes[:, :2]
  127. # 即 xmin ymin ,xmax,ymax 中心点坐标 ((xmax-xmin)/2,(ymax-ymin)/2))
  128. cxcy = (boxes[:, 2:] + boxes[:, :2]) / 2
  129. # 遍历每个边界框的中心点
  130. for i in range(cxcy.size()[0]):
  131. cxcy_sample = cxcy[i] # 获取当前边界框的中心点坐标
  132. # 计算中心点所在的网格索引,并进行向上取整和减一操作
  133. ij = (cxcy_sample / cell_size).ceil() - 1
  134. # 标记目标张量中对应网格的objectness为1
  135. target[int(ij[1]), int(ij[0]), 4] = 1
  136. # 标记目标张量中对应网格的类别为1
  137. target[int(ij[1]), int(ij[0]), 9] = 1
  138. # 标记目标张量中对应网格的标签位置为1
  139. target[int(ij[1]), int(ij[0]), int(labels[i]) + 9] = 1
  140. # 计算中心点在当前网格内的相对坐标
  141. xy = ij * cell_size
  142. # 计算中心点相对于网格左上角的偏移量
  143. delta_xy = (cxcy_sample - xy) / cell_size
  144. # 填充目标张量中对应网格的宽度、高度信息
  145. target[int(ij[1]), int(ij[0]), 2:4] = wh[i]
  146. # 填充目标张量中对应网格的中心点偏移量信息
  147. target[int(ij[1]), int(ij[0]), :2] = delta_xy
  148. # 重复填充宽度、高度信息到目标张量的其他位置
  149. target[int(ij[1]), int(ij[0]), 7:9] = wh[i]
  150. # 重复填充中心点偏移量信息到目标张量的其他位置
  151. target[int(ij[1]), int(ij[0]), 5:7] = delta_xy
  152. # 返回填充好的目标张量
  153. return target
  154. # 以下是各种图像预处理方法的定义
  155. # 将BGR图像转换为RGB图像
  156. def BGR2RGB(self, img):
  157. return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  158. # 将BGR图像转换为HSV图像
  159. def BGR2HSV(self, img):
  160. return cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
  161. # 将HSV图像转换为BGR图像
  162. def HSV2BGR(self, img):
  163. return cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
  164. # 随机调整图像的亮度
  165. def RandomBrightness(self, bgr):
  166. # 以0.5的概率调整图像的亮度
  167. if random.random() < 0.5:
  168. hsv = self.BGR2HSV(bgr) # 将BGR图像转换为HSV图像
  169. h, s, v = cv2.split(hsv) # 分割HSV图像的通道
  170. # 随机选择亮度调整系数
  171. adjust = random.choice([0.5, 1.5])
  172. v = v * adjust # 调整亮度通道的值
  173. # 确保亮度值在0到255之间
  174. v = np.clip(v, 0, 255).astype(hsv.dtype)
  175. # 合并通道得到调整后的HSV图像
  176. hsv = cv2.merge((h, s, v))
  177. # 将调整后的HSV图像转换回BGR图像
  178. bgr = self.HSV2BGR(hsv)
  179. return bgr
  180. # 定义一个随机调整图像饱和度的方法
  181. def RandomSaturation(self, bgr):
  182. # 以0.5的概率决定是否调整图像的饱和度
  183. if random.random() < 0.5:
  184. # 将BGR图像转换为HSV图像
  185. hsv = self.BGR2HSV(bgr)
  186. # 分割HSV图像的通道
  187. h, s, v = cv2.split(hsv)
  188. # 随机选择饱和度调整系数
  189. adjust = random.choice([0.5, 1.5])
  190. # 调整饱和度通道的值
  191. s = s * adjust
  192. # 确保饱和度值在0到255之间
  193. s = np.clip(s, 0, 255).astype(hsv.dtype)
  194. # 合并通道得到调整后的HSV图像
  195. hsv = cv2.merge((h, s, v))
  196. # 将调整后的HSV图像转换回BGR图像
  197. bgr = self.HSV2BGR(hsv)
  198. # 返回调整后的BGR图像
  199. return bgr
  200. # 定义一个随机调整图像色调的方法
  201. # 随机调整图片色调
  202. def RandomHue(self, bgr):
  203. # 以0.5的概率决定是否调整图像的色调
  204. if random.random() < 0.5:
  205. # 将BGR图像转换为HSV图像
  206. hsv = self.BGR2HSV(bgr)
  207. # 分割HSV图像的通道
  208. h, s, v = cv2.split(hsv)
  209. # 随机选择色调调整系数
  210. adjust = random.choice([0.5, 1.5])
  211. # 调整色调通道的值
  212. h = h * adjust
  213. # 确保色调值在0到179之间(因为色调是循环的)
  214. h = h % 180
  215. # 确保色调值的数据类型与hsv图像一致
  216. h = h.astype(hsv.dtype)
  217. # 合并通道得到调整后的HSV图像
  218. hsv = cv2.merge((h, s, v))
  219. # 将调整后的HSV图像转换回BGR图像
  220. bgr = self.HSV2BGR(hsv)
  221. # 返回调整后的BGR图像
  222. return bgr
  223. # 定义一个随机模糊图像的方法
  224. def randomBlur(self, bgr):
  225. # 以0.5的概率决定是否对图像进行模糊处理
  226. if random.random() < 0.5:
  227. # 使用cv2.blur函数对图像进行模糊处理,核大小为(5, 5)
  228. bgr = cv2.blur(bgr, (5, 5))
  229. # 返回模糊处理后的图像或原图像
  230. return bgr
  231. # 定义一个随机平移图像的方法,同时更新边界框和标签的位置
  232. def randomShift(self, bgr, boxes, labels):
  233. # 计算边界框的中心点坐标
  234. center = (boxes[:, 2:] + boxes[:, :2]) / 2
  235. # 以0.5的概率决定是否对图像进行平移
  236. if random.random() < 0.5:
  237. # 获取图像的高度、宽度和通道数
  238. height, width, c = bgr.shape
  239. # 创建一个与原图同样大小的零矩阵,并填充为特定的BGR值(这里可能用于填充平移后的空白区域)
  240. after_shfit_image = np.zeros((height, width, c), dtype=bgr.dtype)
  241. after_shfit_image[:, :, :] = (104, 117, 123) # 填充BGR颜色值
  242. # 生成随机的x和y方向的平移量,限制在图像尺寸的20%以内
  243. shift_x = random.uniform(-width * 0.2, width * 0.2)
  244. shift_y = random.uniform(-height * 0.2, height * 0.2)
  245. # 根据平移量的正负值,对图像进行不同方向的平移
  246. if shift_x >= 0 and shift_y >= 0:
  247. # 当x和y都大于等于0时,从原图的右下角开始平移
  248. after_shfit_image[int(shift_y):, int(shift_x):, :] = bgr[:height - int(shift_y), :width - int(shift_x),
  249. :]
  250. elif shift_x >= 0 and shift_y < 0:
  251. # 当x大于等于0且y小于0时,从原图的右上方开始平移
  252. after_shfit_image[:height + int(shift_y), int(shift_x):, :] = bgr[-int(shift_y):, :width - int(shift_x),
  253. :]
  254. elif shift_x < 0 and shift_y >= 0:
  255. # 当x小于0且y大于等于0时,从原图的左下方开始平移
  256. after_shfit_image[int(shift_y):, :width + int(shift_x), :] = bgr[:height - int(shift_y), -int(shift_x):,
  257. :]
  258. elif shift_x < 0 and shift_y < 0:
  259. # 当x和y都小于0时,从原图的左上方开始平移
  260. after_shfit_image[:height + int(shift_y), :width + int(shift_x), :] = bgr[-int(shift_y):,
  261. -int(shift_x):, :]
  262. # 将平移量转换为Tensor格式,并扩展到与中心点坐标相同的形状
  263. shift_xy = torch.FloatTensor([[int(shift_x), int(shift_y)]]).expand_as(center)
  264. # 更新中心点坐标,加上平移量
  265. center = center + shift_xy
  266. # 创建掩码,确保更新后的中心点坐标在图像尺寸范围内
  267. mask1 = (center[:, 0] > 0) & (center[:, 0] < width)
  268. mask2 = (center[:, 1] > 0) & (center[:, 1] < height)
  269. # 合并两个掩码,得到同时满足两个条件的掩码
  270. mask = (mask1 & mask2).view(-1, 1)
  271. # 根据掩码筛选出更新后仍在图像内的边界框
  272. boxes_in = boxes[mask.expand_as(boxes)].view(-1, 4)
  273. # 如果经过平移后,没有边界框留在图像内,则直接返回原图像、原边界框和原标签
  274. if len(boxes_in) == 0:
  275. return bgr, boxes, labels
  276. # 创建一个Tensor,用于存储平移量,其中x和y的平移量分别对应边界框的左上角和右下角
  277. box_shift = torch.FloatTensor([[int(shift_x), int(shift_y), int(shift_x), int(shift_y)]]).expand_as(
  278. boxes_in)
  279. # 将平移量加到更新后的边界框上,注意这里是对边界框的四个坐标都进行了平移
  280. boxes_in = boxes_in + box_shift
  281. # 根据掩码筛选出更新后仍在图像内的标签
  282. labels_in = labels[mask.view(-1)]
  283. # 返回处理后的图像、更新后的边界框和标签
  284. return after_shfit_image, boxes_in, labels_in
  285. # 如果不需要进行平移,或者平移后逻辑没有进入if条件块,则直接返回原图像、原边界框和原标签
  286. return bgr, boxes, labels
  287. # 定义randomScale方法,用于对图像进行随机尺度变换
  288. def randomScale(self, bgr, boxes):
  289. # 如果随机生成的小于0.5的数,则执行尺度变换
  290. if random.random() < 0.5:
  291. # 生成0.8到1.2之间的随机尺度因子
  292. scale = random.uniform(0.8, 1.2)
  293. # 获取图像的高度、宽度和通道数
  294. height, width, c = bgr.shape
  295. # 使用cv2的resize函数,固定高度不变,按照生成的尺度因子调整宽度
  296. bgr = cv2.resize(bgr, (int(width * scale), height))
  297. # 创建一个Tensor,用于存储尺度因子,扩展为与boxes相同的形状
  298. scale_tensor = torch.FloatTensor([[scale, 1, scale, 1]]).expand_as(boxes)
  299. # 将尺度因子应用到中心点信息boxes上,进行尺度变换
  300. boxes = boxes * scale_tensor
  301. # 返回尺度变换后的图像和边界框
  302. return bgr, boxes
  303. # 如果不进行尺度变换,则返回原图像和边界框
  304. return bgr, boxes
  305. # 定义randomCrop方法,用于对图像进行随机裁剪
  306. # 随机裁剪图片和边界框
  307. def randomCrop(self, bgr, boxes, labels):
  308. # 如果随机生成的小于0.5的数,则执行随机裁剪
  309. if random.random() < 0.5:
  310. # 计算边界框的中心点
  311. center = (boxes[:, 2:] + boxes[:, :2]) / 2
  312. # 获取图像的高度、宽度和通道数
  313. height, width, c = bgr.shape
  314. # 生成随机裁剪的高度和宽度,至少为原图像的60%
  315. h = random.uniform(0.6 * height, height)
  316. w = random.uniform(0.6 * width, width)
  317. # 生成随机裁剪的起始坐标
  318. x = random.uniform(0, width - w)
  319. y = random.uniform(0, height - h)
  320. # 将坐标和裁剪的尺寸转换为整数
  321. x, y, h, w = int(x), int(y), int(h), int(w)
  322. # 将裁剪起始坐标转换为Tensor,并从中心点中减去
  323. center = center - torch.FloatTensor([[x, y]]).expand_as(center)
  324. # 创建掩码,筛选出中心点落在裁剪区域内的边界框
  325. mask1 = (center[:, 0] > 0) & (center[:, 0] < w)
  326. mask2 = (center[:, 1] > 0) & (center[:, 1] < h)
  327. # 将两个掩码合并,并展平为一维Tensor
  328. mask = (mask1 & mask2).view(-1, 1)
  329. # 根据掩码筛选出裁剪区域内的边界框
  330. boxes_in = boxes[mask.expand_as(boxes)].view(-1, 4)
  331. # 如果没有边界框落在裁剪区域内,则返回原图像、原边界框和原标签
  332. if (len(boxes_in) == 0):
  333. return bgr, boxes, labels
  334. # 创建一个Tensor,用于存储裁剪起始坐标,扩展为与boxes_in相同的形状
  335. box_shift = torch.FloatTensor([[x, y, x, y]]).expand_as(boxes_in)
  336. # 将裁剪起始坐标从边界框中减去
  337. boxes_in = boxes_in - box_shift
  338. # 确保裁剪后的边界框坐标在裁剪区域内
  339. boxes_in[:, 0] = boxes_in[:, 0].clamp_(min=0, max=w)
  340. boxes_in[:, 2] = boxes_in[:, 2].clamp_(min=0, max=w)
  341. boxes_in[:, 1] = boxes_in[:, 1].clamp_(min=0, max=h)
  342. boxes_in[:, 3] = boxes_in[:, 3].clamp_(min=0, max=h)
  343. # 根据掩码筛选出裁剪区域内的标签
  344. labels_in = labels[mask.view(-1)]
  345. # 裁剪图像
  346. img_croped = bgr[y:y + h, x:x + w, :]
  347. # 返回裁剪后的图像、裁剪区域内的边界框和标签
  348. return img_croped, boxes_in, labels_in
  349. # 如果不进行裁剪,则返回原图像、原边界框和原标签
  350. return bgr, boxes, labels
  351. # 定义subMean方法,用于从图像中减去均值
  352. def subMean(self, bgr, mean):
  353. # 将均值转换为float32类型的numpy数组
  354. mean = np.array(mean, dtype=np.float32)
  355. # 从图像bgr中减去均值
  356. bgr = bgr - mean
  357. # 返回减去均值后的图像
  358. return bgr
  359. # 定义random_flip方法,用于随机对图像进行水平翻转,并相应地更新边界框
  360. # 随机翻转图片和边界框
  361. def random_flip(self, im, boxes):
  362. # 以0.5的概率决定是否进行翻转
  363. if random.random() < 0.5:
  364. # 对图像进行水平翻转
  365. im_lr = np.fliplr(im).copy()
  366. # 获取图像的高度、宽度
  367. h, w, _ = im.shape
  368. # 更新边界框的左边界为原右边界与图像宽度的差值
  369. xmin = w - boxes[:, 2]
  370. # 更新边界框的右边界为原左边界与图像宽度的差值
  371. xmax = w - boxes[:, 0]
  372. # 将更新后的左右边界赋值回边界框
  373. boxes[:, 0] = xmin
  374. boxes[:, 2] = xmax
  375. # 返回翻转后的图像和更新后的边界框
  376. return im_lr, boxes
  377. # 如果不进行翻转,则返回原图像和原边界框
  378. return im, boxes
  379. # 定义random_bright方法,用于随机调整图像的亮度
  380. def random_bright(self, im, delta=16):
  381. # 生成一个0到1之间的随机数
  382. alpha = random.random()
  383. # 如果该随机数大于0.3
  384. if alpha > 0.3:
  385. # 调整图像的亮度,乘以随机数alpha并加上一个-delta到delta之间的随机整数
  386. im = im * alpha + random.randrange(-delta, delta)
  387. # 将调整亮度后的图像像素值限制在0到255之间,并转换为uint8类型
  388. im = im.clip(min=0, max=255).astype(np.uint8)
  389. # 返回调整亮度后的图像
  390. return im
  391. # 3. 调试代码
  392. def main():
  393. from torch.utils.data import DataLoader
  394. import torchvision.transforms as transforms
  395. file_root = '../data/VOC2012/JPEGImages/' # 记得改为自己的路径
  396. train_dataset = Yolo_Dataset(root=file_root,list_file='voctrain.txt',train=True,transforms = [T.ToTensor()] )
  397. train_loader = DataLoader(train_dataset,batch_size=1,shuffle=False,num_workers=0)
  398. train_iter = iter(train_loader)
  399. for i in range(100):
  400. img,target = next(train_iter)
  401. print(img.shape)
  402. print(target.shape)
  403. break
  404. if __name__ == '__main__':
  405. main()

1.5、定义主函数

       在utils目录下创建工具类 yolo_dataset.py,中定义主函数进行测试,包含设置和加载一个自定义的Yolo_Dataset数据集,该数据集来源于VOC2012的JPEGImages文件夹。代码中首先导入了必要的模块和类,然后定义了数据集的根目录和预处理操作。接着,创建了一个Yolo_Dataset对象,并通过DataLoader类将其加载为可迭代的数据集。DataLoader的参数包括数据集、批次大小、是否打乱数据以及使用的子进程数量。最后,通过迭代器遍历数据集,获取图像和对应的目标数据,并打印它们的维度信息。

  1. # 3. 调试代码
  2. # 定义一个名为main的函数
  3. def main():
  4. # 从torch.utils.data模块中导入DataLoader类,用于加载数据集
  5. from torch.utils.data import DataLoader
  6. # 导入torchvision.transforms模块,用于图像预处理
  7. import torchvision.transforms as transforms
  8. # 设置文件根目录路径,这里假设是VOC2012数据集的JPEGImages文件夹
  9. file_root = '../data/VOC2012/JPEGImages/' # 记得改为自己的路径
  10. # 创建一个Yolo_Dataset对象,该对象可能是自定义的,用于加载VOC2012数据集
  11. # 参数包括:
  12. # root: 数据集的根目录
  13. # list_file: 包含训练图像文件名的列表文件
  14. # train: 是否为训练模式
  15. # transforms: 图像预处理操作列表,这里只有一个操作,即将图像转换为Tensor
  16. train_dataset = Yolo_Dataset(root=file_root, list_file='voctrain.txt', train=True,
  17. transforms=[transforms.ToTensor()])
  18. # 使用DataLoader类加载训练数据集
  19. # 参数包括:
  20. # dataset: 要加载的数据集
  21. # batch_size: 每个批次的大小
  22. # shuffle: 是否在每个epoch开始时打乱数据
  23. # num_workers: 加载数据使用的子进程数量
  24. train_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=0)
  25. # 创建一个迭代器,用于遍历train_loader中的数据
  26. train_iter = iter(train_loader)
  27. # 遍历前100个批次的数据(但由于下面的break语句,实际上只会处理一个批次)
  28. for i in range(100):
  29. # 使用迭代器获取下一个批次的数据,包括图像和目标(可能是标签或边界框)
  30. img, target = next(train_iter)
  31. # 打印图像的shape(维度)
  32. print(img.shape)
  33. # 打印目标的shape(维度)
  34. print(target.shape)
  35. # 跳出循环,因此只处理一个批次的数据
  36. break
  37. if __name__ == '__main__':
  38. main()

2构建网络

       在YOLO(You Only Look Once)系列的目标检测算法中,Backbone是一个核心组成部分,主要负责提取输入图像的特征。具体来说,Backbone是一种卷积神经网络,它通过多层次的特征提取,从低层次到高层次逐渐提取出图像的不同特征,如边缘、纹理、形状等,最终生成一个高维度的特征向量。这些特征向量在后续的网络层中被用于目标检测和分类等任务。

       在YOLOv1中,大部分人都选用的ResNet作为自己的backbone,因为可以方便调用官方的预训练权重。

2.1 ResNet架构介绍

      ResNet架构使用残差的模式,即输出值为F(x)+x模式,简单理解就是至少不比上次结果差。不了解可以参考其他人博客.

2.2 构建ResNet

       在创建的目录结构network中创建Yolo_ResNet.py文件,用于创建 Base_Block、Senior_Block、Output_Block 和 ResNet。

      ResNet是一个基于残差块(Block)的卷积神经网络结构。它使用多个层级,每个层级包含多个残差块(Block)。构建残差块时候需要注意的是,ResNet18—ResNet152共有两种类型的Block块,一种不涉及1*1卷积网络,另外一种涉及1*1卷积网络,如下图所示:

       自己实现Block时候也要有涉及1*1卷积网络 和 不涉及1*1卷积网络 两种方式。

2.2.1 构建Base_Block

       Base_Block类是一个神经网络的基础块,它定义了一个具有残差连接的结构,这在深度神经网络中(如ResNet)是非常常见的。这个类中不包含 1*1 的卷积,以下是该类中几个函数的主要作用:

__init__ 函数
  • 作用:初始化Base_Block对象。
  • 功能
    • 接收输入通道数(in_planes)、输出通道数(out_planes)、步长(stride)以及下采样函数(downsample)作为参数。
    • 定义并初始化两个卷积层(conv1 和 conv2)、两个批量归一化层(bn1 和 bn2)以及ReLU激活函数(relu)。其中,conv1 和 conv2  使用 3x3 的卷积核。
    • 保存传入的下采样函数(downsample)和步长(stride)以便后续使用。
forward 函数
  • 作用:定义网络块的前向传播过程。
  • 功能
    • 接收输入数据x
    • 保留输入x作为残差(res),用于后续与卷积处理后的输出相加。
    • 将输入x通过两个卷积层、两个批量归一化层以及ReLU激活函数。
    • 如果提供了下采样函数,则对残差res进行下采样,使其维度与卷积后的输出一致。
    • 将卷积处理后的输出与残差res相加,实现残差连接。
    • 对相加后的结果再次应用ReLU激活函数。
    • 返回前向传播的结果。
  1. '''
  2. # 2. 构建block: 不含有1*1
  3. # 定义一个基础块类,继承自nn.Module,用于构建神经网络的基础结构
  4. '''
  5. class Base_Block(nn.Module):
  6. # 用于扩充的变量,表示扩大几倍
  7. # 在某些网络结构中,如ResNet的Bottleneck块,可能需要通过该变量来扩充通道数
  8. expansion = 1
  9. def __init__(self, in_planes, out_planes, stride=1, downsample=None):
  10. # 初始化函数,传入输入通道数、输出通道数、步长以及下采样函数
  11. '''
  12. :param in_planes: 输入的通道数
  13. :param planes: 输出的通道数
  14. :param stride: 默认步长
  15. :param downsample: 是否进行下采样
  16. '''
  17. # 调用父类nn.Module的初始化函数
  18. super(Base_Block, self).__init__()
  19. # 定义网络结构 + 初始化参数
  20. # 第一个卷积层
  21. self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
  22. # 第一个批量归一化层
  23. self.bn1 = nn.BatchNorm2d(out_planes)
  24. # ReLU激活函数,inplace=True表示直接改变原始参数值,节省内存
  25. self.relu = nn.ReLU(inplace=True)
  26. # 第二个卷积层
  27. self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
  28. # 第二个批量归一化层
  29. self.bn2 = nn.BatchNorm2d(out_planes)
  30. # 下采样函数,用于调整输入x的维度,以便与卷积后的输出维度一致
  31. self.downsample = downsample
  32. # 存储步长,以便后续可能用到
  33. self.stride = stride
  34. # 定义前向传播函数
  35. def forward(self, x):
  36. # 残差,即直接输入x,后续会与卷积处理后的输出相加
  37. res = x
  38. # 正常传播
  39. # 通过第一个卷积层
  40. out = self.conv1(x)
  41. # 通过第一个批量归一化层
  42. out = self.bn1(out)
  43. # 通过ReLU激活函数
  44. out = self.relu(out)
  45. # 通过第二个卷积层
  46. out = self.conv2(out)
  47. # 通过第二个批量归一化层
  48. out = self.bn2(out)
  49. # 判断是否下采样
  50. # 如果需要下采样,则调整残差res的维度
  51. if self.downsample is not None:
  52. res = self.downsample(res)
  53. # 残差相加
  54. # 将残差res与卷积处理后的输出out相加,实现残差连接
  55. out += res
  56. # 再次通过ReLU激活函数
  57. out = self.relu(out)
  58. # 返回结果
  59. return out

2.2.2 构建Senior_Block

        Senior_Block 类是用于构建更复杂的网络结构。类中包含1*1卷积核3*3卷积,下面是该类中几个函数的主要作用:

__init__ 初始化函数
  • 作用:初始化 Senior_Block 对象,设置网络层的结构和参数。
  • 功能
    • 接收输入通道数(in_planes)、中间通道数(planes)、步长(stride)以及下采样方法(downsample)作为参数。
    • 定义三个卷积层(conv1conv2conv3)和对应的批量归一化层(bn1bn2bn3)。其中,conv1 和 conv3 使用 1x1 的卷积核,conv2 使用 3x3 的卷积核。
    • 定义 ReLU 激活函数(relu)。
    • 保存传入的下采样方法(downsample)和步长(stride)以便后续使用。
forward 前向传播函数
  • 作用:定义网络块的前向传播过程,即数据从输入到输出的计算流程。
  • 功能
    • 接收输入数据 x
    • 保留输入 x 作为残差(res),用于后续与卷积处理后的输出相加。
    • 将输入 x 通过三个卷积层(conv1conv2conv3)和对应的批量归一化层(bn1bn2bn3),以及 ReLU 激活函数。
    • 如果提供了下采样方法,则对残差 res 进行下采样,使其维度与卷积后的输出一致。
    • 将卷积处理后的输出与残差 res 相加,实现残差连接。
    • 对相加后的结果再次应用 ReLU 激活函数。
    • 返回前向传播的结果。
  1. '''
  2. # 3. 构建Block:含有1*1
  3. # 定义高级块类,继承自nn.Module,用于构建更复杂的网络结构
  4. '''
  5. class Senior_Block(nn.Module):
  6. expansion = 4
  7. def __init__(self, in_planes, planes, stride=1, downsample=None):
  8. '''
  9. :param in_planes: 输入通道数
  10. :param planes: 中间通道数,最终的输出通道数还需要乘以扩大系数,即expansion
  11. :param stride: 步长
  12. :param downsample: 下采样方法
  13. '''
  14. super(Senior_Block, self).__init__()
  15. self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
  16. self.bn1 = nn.BatchNorm2d(planes)
  17. self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
  18. self.bn2 = nn.BatchNorm2d(planes)
  19. self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
  20. self.bn3 = nn.BatchNorm2d(planes * 4)
  21. self.relu = nn.ReLU(inplace=True)
  22. self.downsample = downsample
  23. self.stride = stride
  24. def forward(self, x):
  25. # 残差
  26. res = x
  27. # 前向传播
  28. out = self.conv1(x)
  29. out = self.bn1(out)
  30. out = self.relu(out)
  31. out = self.conv2(out)
  32. out = self.bn2(out)
  33. out = self.relu(out)
  34. out = self.conv3(out)
  35. out = self.bn3(out)
  36. # 是否下采样
  37. if self.downsample is not None:
  38. res = self.downsample(x)
  39. # 相加
  40. out += res
  41. out = self.relu(out)
  42. return out

2.2.3 构建Output_Block

       Output_Block 类是用于构建神经网络的一部分。下面是该类中几个函数的主要作用:

__init__ 初始化函数
  • 作用:初始化 Output_Block 对象,设置网络层的结构和参数。
  • 功能
    • 接收输入通道数(in_planes)、中间通道数(planes)、步长(stride)以及块类型(block_type)作为参数。
    • 定义三个卷积层(conv1conv2conv3)和对应的批量归一化层(bn1bn2bn3)。其中,conv1 和 conv3 使用 1x1 的卷积核,conv2 使用 3x3 的卷积核。。
    • 根据步长、输入通道数和块类型,判断是否需要进行下采样,并初始化下采样层(downsample)。
forward 前向传播函数
  • 作用:定义网络块的前向传播过程,即数据从输入到输出的计算流程。
  • 功能
    • 接收输入数据 x
    • 将输入 x 通过三个卷积层(conv1conv2conv3)和对应的批量归一化层(bn1bn2bn3),并在每个卷积层之后应用 ReLU 激活函数。
    • 将输入 x 通过下采样层(downsample),以便与卷积处理后的输出 out 的维度相匹配。
    • 将卷积处理后的输出 out 与下采样后的输入相加,实现残差连接。
    • 返回前向传播的结果。

  1. '''
  2. # 4. 构建输出层
  3. # 定义输出块类,继承自nn.Module,用于构建更复杂的网络结构
  4. '''
  5. class Output_Block(nn.Module):
  6. # 定义扩大系数,该块最终输出的通道数将与中间通道数相同
  7. expansion = 1
  8. def __init__(self, in_planes, planes, stride=1, block_type='A'):
  9. # 初始化函数,传入输入通道数、中间通道数、步长以及块类型
  10. '''
  11. :param in_planes: 输入通道数
  12. :param planes: 中间通道数
  13. :param stride: 步长
  14. :param block_type: 块类型,为'A'表示不需要下采样,为'B'则需要
  15. '''
  16. # 调用父类nn.Module的初始化函数
  17. super(Output_Block, self).__init__()
  18. # 定义第一个卷积层,用于改变通道数
  19. self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
  20. # 定义第一个批量归一化层
  21. self.bn1 = nn.BatchNorm2d(planes)
  22. # 定义第二个卷积层,使用3x3的卷积核,带有2的空洞率(dilation),可以增大感受野
  23. self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=2, bias=False, dilation=2)
  24. # 定义第二个批量归一化层
  25. self.bn2 = nn.BatchNorm2d(planes)
  26. # 定义第三个卷积层,用于调整通道数至最终的输出通道数
  27. self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
  28. # 定义第三个批量归一化层
  29. self.bn3 = nn.BatchNorm2d(self.expansion * planes)
  30. # 初始化下采样层,默认是一个空的序列
  31. self.downsample = nn.Sequential()
  32. # 判断是否需要下采样,这里综合了步长、通道数以及块类型
  33. if stride != 1 or in_planes != self.expansion * planes or block_type == 'B':
  34. # 如果需要下采样,则设置下采样层
  35. self.downsample = nn.Sequential(
  36. # 定义一个卷积层,用于调整输入的通道数和维度
  37. nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
  38. # 定义一个批量归一化层
  39. nn.BatchNorm2d(self.expansion * planes)
  40. )
  41. # 定义前向传播函数
  42. def forward(self, x):
  43. # 通过第一个卷积层和批量归一化层,并应用ReLU激活函数
  44. out = F.relu(self.bn1(self.conv1(x)))
  45. # 通过第二个卷积层和批量归一化层,并再次应用ReLU激活函数
  46. out = F.relu(self.bn2(self.conv2(out)))
  47. # 通过第三个卷积层和批量归一化层
  48. out = self.bn3(self.conv3(out))
  49. # 将输入x通过下采样层,以便与out的维度相匹配
  50. # 如果不需要下采样,则下采样层是一个空操作
  51. out += self.downsample(x)
  52. # 应用ReLU激活函数
  53. out = F.relu(out)
  54. # 返回输出结果
  55. return out

2.2.4 构建ResNet

   ResNet类定义了一个基于残差块(block)的卷积神经网络结构。它使用了多个层级(layer1layer4)和一个输出层(layer5),以及一个平均池化层和一个最终的卷积层。下面是该类中几个函数的主要作用:

__init__ 初始化函数
  • 作用:初始化ResNet对象,设置网络的结构和参数。
  • 功能
    • 设置初始通道数self.inplanes为64。
    • 定义初始卷积层self.conv1,用于从输入图像中提取特征。
    • 定义批标准化层self.bn1和ReLU激活函数self.relu,用于处理conv1的输出。
    • 定义最大池化层self.maxpool,用于进一步减小特征图的尺寸。
    • 使用_make_layer方法创建四个层级(layer1layer4),每个层级包含多个残差块。
    • 使用_make_out_layer方法创建输出层self.layer5
    • 定义平均池化层self.avgpool,用于减小特征图的尺寸。
    • 定义最终的卷积层self.conv_end和批标准化层self.bn_end,用于输出层的处理。
    • 初始化模型参数,使用正态分布初始化卷积层的权重,并将批标准化层的权重初始化为1,偏置初始化为0。
_make_layer 方法
  • 作用:根据输入的block类型、输出通道数、block数量和步长,创建并返回一个层级(layer)。
  • 功能
    • 循环创建指定数量的block对象,并将它们添加到一个nn.Sequential容器中。
    • 在第一个block中,根据步长设置是否进行下采样。
    • 返回包含多个block的层级。
_make_out_layer 方法
  • 作用:创建输出层。
  • 功能
    • 根据输入的通道数,定义输出层的结构。
    • 输出层的具体结构在代码中没有给出,但通常会包含卷积层、批标准化层等。
forward 方法
  • 作用:定义网络的前向传播过程,即从输入到输出的计算流程。
  • 功能
    • 接收输入数据,并通过初始卷积层、批标准化层、ReLU激活函数和最大池化层进行处理。
    • 将处理后的数据依次通过四个层级(layer1layer4)和输出层(layer5)。
    • 通过平均池化层减小特征图的尺寸。
    • 通过最终的卷积层和批标准化层进行处理,得到网络的输出。
初始化模型参数部分
  • 作用:对模型中的参数进行初始化。
  • 功能
    • 遍历模型中的所有模块。
    • 对于卷积层,使用正态分布初始化权重,标准差根据卷积核的大小和输出通道数计算得出。
    • 对于批标准化层,将权重初始化为1,偏置初始化为0。
  1. '''
  2. # 5. 构建ResNet
  3. # 定义ResNet类,继承自nn.Module,用于构建更复杂的网络结构
  4. '''
  5. class ResNet(nn.Module):
  6. def __init__(self, block, layers):
  7. '''
  8. 初始化ResNet网络。
  9. :param block: 基本的Block块对象,如BasicBlock或Bottleneck。
  10. :param layers: 不同层级的block数量,例如ResNet50为[3, 4, 6, 3]。
  11. '''
  12. super(ResNet, self).__init__() # 调用父类nn.Module的初始化方法。
  13. # 初始通道数设置为64
  14. self.inplanes = 64
  15. # 初始卷积层,输入通道为3(RGB图像),输出通道为64,卷积核大小为7x7,步长为2,填充为3,无偏置。
  16. self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
  17. # 批标准化层,用于conv1的输出
  18. self.bn1 = nn.BatchNorm2d(64)
  19. # ReLU激活函数,inplace=True表示直接在原变量上进行操作,节省内存。
  20. self.relu = nn.ReLU(inplace=True)
  21. # 最大池化层,池化核大小为3x3,步长为2,填充为1。
  22. self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
  23. # 根据输入的block和layers创建不同的层级
  24. # 创建四个层级(layer1至layer4),每个层级包含多个残差块
  25. self.layer1 = self._make_layer(block, 64, layers[0]) # 第一层,输出通道为64,block数量由layers[0]决定。
  26. self.layer2 = self._make_layer(block, 128, layers[1], stride=2) # 第二层,输出通道为128,block数量由layers[1]决定,步长为2。
  27. self.layer3 = self._make_layer(block, 256, layers[2], stride=2) # 第三层,输出通道为256,block数量由layers[2]决定,步长为2。
  28. self.layer4 = self._make_layer(block, 512, layers[3], stride=2) # 第四层,输出通道为512,block数量由layers[3]决定,步长为2。
  29. # 创建输出层
  30. self.layer5 = self._make_out_layer(in_channels=2048) # 输出层的输入通道为2048。
  31. # 平均池化层,池化核大小和步长均为2。
  32. self.avgpool = nn.AvgPool2d(2)
  33. # 最后的卷积层,将输出通道变为30,卷积核大小为3x3,步长为1,填充为1,无偏置。
  34. self.conv_end = nn.Conv2d(256, 30, kernel_size=3, stride=1, padding=1, bias=False)
  35. # 批标准化层,用于conv_end的输出
  36. self.bn_end = nn.BatchNorm2d(30)
  37. # 初始化模型参数
  38. for m in self.modules(): # 遍历模型中的所有模块
  39. if isinstance(m, nn.Conv2d): # 如果是卷积层
  40. n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels # 计算权重初始化的标准差
  41. m.weight.data.normal_(0, math.sqrt(2. / n)) # 使用正态分布初始化权重
  42. elif isinstance(m, nn.BatchNorm2d): # 如果是批标准化层
  43. m.weight.data.fill_(1) # 权重初始化为1
  44. m.bias.data.zero_() # 偏置初始化为0
  45. # 根据传入的layer个数和block创建
  46. # 输入层的构建
  47. def _make_layer(self, block, planes, blocks, stride=1):
  48. '''
  49. 根据给定的block类型、输出通道数、block数量和步长来创建一系列block组成的层。
  50. :param block: 要使用的block类型,是一个类,如BasicBlock或Bottleneck
  51. :param planes: 每个block的输出通道数
  52. :param blocks: 需要搭建的block数量
  53. :param stride: 第一个block的卷积层步长,用于控制是否进行下采样
  54. '''
  55. # 初始化下采样变量,用于第一个block可能需要的下采样操作
  56. downsample = None
  57. # 判断是否需要进行下采样
  58. # 如果步长不为1,或者当前层的输入通道数(self.inplanes)与block输出通道数(planes * block.expansion)不匹配,则进行下采样
  59. if stride != 1 or self.inplanes != planes * block.expansion:
  60. downsample = nn.Sequential(
  61. # 使用1x1的卷积层进行通道数转换和步长调整,使得残差连接可以相加
  62. nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
  63. # 批量归一化,稳定训练过程
  64. nn.BatchNorm2d(planes * block.expansion),
  65. )
  66. # 初始化层的列表
  67. layers = []
  68. # 添加第一个block,并可能进行下采样
  69. # 第一个block的输入通道数是self.inplanes,输出通道数是planes,步长是stride,可能包含下采样操作
  70. layers.append(block(self.inplanes, planes, stride, downsample))
  71. # 更新当前层的输入通道数为block的输出通道数
  72. self.inplanes = planes * block.expansion
  73. # 添加剩余的blocks,它们的输入通道数已经是更新后的self.inplanes,输出通道数保持不变
  74. for i in range(1, blocks):
  75. layers.append(block(self.inplanes, planes))
  76. # 将层的列表转换为Sequential模型,便于作为神经网络的一部分进行前向传播
  77. return nn.Sequential(*layers) # *操作符用于将列表中的元素解包为函数参数
  78. # 输出层的构建
  79. def _make_out_layer(self, in_channels):
  80. '''
  81. 根据输入通道数构建输出层,包含多个Output_Block。
  82. :param in_channels: 输入通道数
  83. '''
  84. # 初始化层的列表
  85. layers = []
  86. # 添加第一个Output_Block,并指定其输入通道数和输出通道数
  87. # 注意:这里使用了block_type参数来区分不同的Output_Block类型
  88. layers.append(Output_Block(in_planes=in_channels, planes=256, block_type='B'))
  89. # 添加两个相同类型的Output_Block,输入通道数和输出通道数都保持不变
  90. layers.append(Output_Block(in_planes=256, planes=256, block_type='A'))
  91. layers.append(Output_Block(in_planes=256, planes=256, block_type='A'))
  92. # 将层的列表转换为Sequential模型
  93. return nn.Sequential(*layers) # *操作符用于将列表中的元素解包为函数参数
  94. def forward(self, x):
  95. # x 为网络的输入,这里的前向传播方法描述了网络如何从输入x获得最终的输出
  96. # 经历共有的卷积和池化层
  97. # 第一个卷积层,对输入x进行卷积操作
  98. x = self.conv1(x)
  99. # 批量归一化层,对卷积后的结果进行归一化处理
  100. x = self.bn1(x)
  101. # ReLU激活函数,对归一化后的结果进行非线性变换
  102. x = self.relu(x)
  103. # 最大池化层,对激活后的结果进行下采样
  104. x = self.maxpool(x)
  105. # 经历各个block块
  106. # 第一个block层,对池化后的结果进行进一步的特征提取
  107. x = self.layer1(x)
  108. # 第二个block层,继续提取特征
  109. x = self.layer2(x)
  110. # 第三个block层
  111. x = self.layer3(x)
  112. # 第四个block层
  113. x = self.layer4(x)
  114. # 第五个block层,这里需要注意,通常标准的ResNet结构中只有4个layer,这里可能是自定义的扩展
  115. x = self.layer5(x)
  116. # 经历最终的输出
  117. # 平均池化层,对最后一个block的输出进行全局平均池化
  118. x = self.avgpool(x)
  119. # 卷积层,对池化后的结果进行进一步的卷积变换
  120. x = self.conv_end(x)
  121. # 批量归一化层,对卷积后的结果进行归一化处理
  122. x = self.bn_end(x)
  123. # Sigmoid激活函数,将特征映射到0-1之间
  124. x = F.sigmoid(x) # 归一化到0-1
  125. # 将输出构建为正确的shape
  126. # permute函数用于改变张量的维度顺序,这里将通道维度移动到最后一个位置
  127. # 假设输出shape原本为(-1, C, H, W),经过permute后变为(-1, H, W, C)
  128. x = x.permute(0, 2, 3, 1) # (-1,7,7,30)
  129. # 返回最终的输出x
  130. return x

2.3 定义模型

      构建不同的ResNet函数。即不同的ResNet模型。如下:

  1. # 6. 构建不同的ResNet函数
  2. # 预训练下载链接
  3. model_urls = {
  4. 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
  5. 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
  6. 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
  7. 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
  8. 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
  9. }
  10. # 构建ResNet18
  11. def resnet18(pretrained=False, **kwargs):
  12. model = ResNet(Base_Block, [2, 2, 2, 2], **kwargs)
  13. # 是否预训练
  14. if pretrained:
  15. model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
  16. return model
  17. # 构建ResNet34
  18. def resnet34(pretrained=False, **kwargs):
  19. model = ResNet(Base_Block, [3, 4, 6, 3], **kwargs)
  20. if pretrained:
  21. model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
  22. return model
  23. # 构建ResNet50
  24. def resnet50(pretrained=False, **kwargs):
  25. model = ResNet(Senior_Block, [3, 4, 6, 3], **kwargs)
  26. if pretrained:
  27. model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
  28. return model
  29. # 构建ResNet101
  30. def resnet101(pretrained=False, **kwargs):
  31. model = ResNet(Senior_Block, [3, 4, 23, 3], **kwargs)
  32. if pretrained:
  33. model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
  34. return model
  35. # 构建ResNet152
  36. def resnet152(pretrained=False, **kwargs):
  37. model = ResNet(Senior_Block, [3, 8, 36, 3], **kwargs)
  38. if pretrained:
  39. model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
  40. return model

2.4 完整代码

Yolo_ResNet卷积神经网络

  1. ''
  2. '''
  3. # 1. 导入所需要的包
  4. '''
  5. import torch
  6. import math
  7. from torch import nn
  8. import torch.utils.model_zoo as model_zoo
  9. import torch.nn.functional as F
  10. '''
  11. # 2. 构建block: 不含有1*1
  12. # 定义一个基础块类,继承自nn.Module,用于构建神经网络的基础结构
  13. '''
  14. class Base_Block(nn.Module):
  15. # 用于扩充的变量,表示扩大几倍
  16. # 在某些网络结构中,如ResNet的Bottleneck块,可能需要通过该变量来扩充通道数
  17. expansion = 1
  18. def __init__(self, in_planes, out_planes, stride=1, downsample=None):
  19. # 初始化函数,传入输入通道数、输出通道数、步长以及下采样函数
  20. '''
  21. :param in_planes: 输入的通道数
  22. :param planes: 输出的通道数
  23. :param stride: 默认步长
  24. :param downsample: 是否进行下采样
  25. '''
  26. # 调用父类nn.Module的初始化函数
  27. super(Base_Block, self).__init__()
  28. # 定义网络结构 + 初始化参数
  29. # 第一个卷积层
  30. self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
  31. # 第一个批量归一化层
  32. self.bn1 = nn.BatchNorm2d(out_planes)
  33. # ReLU激活函数,inplace=True表示直接改变原始参数值,节省内存
  34. self.relu = nn.ReLU(inplace=True)
  35. # 第二个卷积层
  36. self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
  37. # 第二个批量归一化层
  38. self.bn2 = nn.BatchNorm2d(out_planes)
  39. # 下采样函数,用于调整输入x的维度,以便与卷积后的输出维度一致
  40. self.downsample = downsample
  41. # 存储步长,以便后续可能用到
  42. self.stride = stride
  43. # 定义前向传播函数
  44. def forward(self, x):
  45. # 残差,即直接输入x,后续会与卷积处理后的输出相加
  46. res = x
  47. # 正常传播
  48. # 通过第一个卷积层
  49. out = self.conv1(x)
  50. # 通过第一个批量归一化层
  51. out = self.bn1(out)
  52. # 通过ReLU激活函数
  53. out = self.relu(out)
  54. # 通过第二个卷积层
  55. out = self.conv2(out)
  56. # 通过第二个批量归一化层
  57. out = self.bn2(out)
  58. # 判断是否下采样
  59. # 如果需要下采样,则调整残差res的维度
  60. if self.downsample is not None:
  61. res = self.downsample(res)
  62. # 残差相加
  63. # 将残差res与卷积处理后的输出out相加,实现残差连接
  64. out += res
  65. # 再次通过ReLU激活函数
  66. out = self.relu(out)
  67. # 返回结果
  68. return out
  69. '''
  70. # 3. 构建Block:含有1*1
  71. # 定义高级块类,继承自nn.Module,用于构建更复杂的网络结构
  72. '''
  73. class Senior_Block(nn.Module):
  74. expansion = 4
  75. def __init__(self, in_planes, planes, stride=1, downsample=None):
  76. '''
  77. :param in_planes: 输入通道数
  78. :param planes: 中间通道数,最终的输出通道数还需要乘以扩大系数,即expansion
  79. :param stride: 步长
  80. :param downsample: 下采样方法
  81. '''
  82. super(Senior_Block, self).__init__()
  83. self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
  84. self.bn1 = nn.BatchNorm2d(planes)
  85. self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
  86. self.bn2 = nn.BatchNorm2d(planes)
  87. self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
  88. self.bn3 = nn.BatchNorm2d(planes * 4)
  89. self.relu = nn.ReLU(inplace=True)
  90. self.downsample = downsample
  91. self.stride = stride
  92. def forward(self, x):
  93. # 残差
  94. res = x
  95. # 前向传播
  96. out = self.conv1(x)
  97. out = self.bn1(out)
  98. out = self.relu(out)
  99. out = self.conv2(out)
  100. out = self.bn2(out)
  101. out = self.relu(out)
  102. out = self.conv3(out)
  103. out = self.bn3(out)
  104. # 是否下采样
  105. if self.downsample is not None:
  106. res = self.downsample(x)
  107. # 相加
  108. out += res
  109. out = self.relu(out)
  110. return out
  111. '''
  112. # 4. 构建输出层
  113. # 定义输出块类,继承自nn.Module,用于构建更复杂的网络结构
  114. '''
  115. class Output_Block(nn.Module):
  116. # 定义扩大系数,该块最终输出的通道数将与中间通道数相同
  117. expansion = 1
  118. def __init__(self, in_planes, planes, stride=1, block_type='A'):
  119. # 初始化函数,传入输入通道数、中间通道数、步长以及块类型
  120. '''
  121. :param in_planes: 输入通道数
  122. :param planes: 中间通道数
  123. :param stride: 步长
  124. :param block_type: 块类型,为'A'表示不需要下采样,为'B'则需要
  125. '''
  126. # 调用父类nn.Module的初始化函数
  127. super(Output_Block, self).__init__()
  128. # 定义第一个卷积层,用于改变通道数
  129. self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
  130. # 定义第一个批量归一化层
  131. self.bn1 = nn.BatchNorm2d(planes)
  132. # 定义第二个卷积层,使用3x3的卷积核,带有2的空洞率(dilation),可以增大感受野
  133. self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=2, bias=False, dilation=2)
  134. # 定义第二个批量归一化层
  135. self.bn2 = nn.BatchNorm2d(planes)
  136. # 定义第三个卷积层,用于调整通道数至最终的输出通道数
  137. self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
  138. # 定义第三个批量归一化层
  139. self.bn3 = nn.BatchNorm2d(self.expansion * planes)
  140. # 初始化下采样层,默认是一个空的序列
  141. self.downsample = nn.Sequential()
  142. # 判断是否需要下采样,这里综合了步长、通道数以及块类型
  143. if stride != 1 or in_planes != self.expansion * planes or block_type == 'B':
  144. # 如果需要下采样,则设置下采样层
  145. self.downsample = nn.Sequential(
  146. # 定义一个卷积层,用于调整输入的通道数和维度
  147. nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
  148. # 定义一个批量归一化层
  149. nn.BatchNorm2d(self.expansion * planes)
  150. )
  151. # 定义前向传播函数
  152. def forward(self, x):
  153. # 通过第一个卷积层和批量归一化层,并应用ReLU激活函数
  154. out = F.relu(self.bn1(self.conv1(x)))
  155. # 通过第二个卷积层和批量归一化层,并再次应用ReLU激活函数
  156. out = F.relu(self.bn2(self.conv2(out)))
  157. # 通过第三个卷积层和批量归一化层
  158. out = self.bn3(self.conv3(out))
  159. # 将输入x通过下采样层,以便与out的维度相匹配
  160. # 如果不需要下采样,则下采样层是一个空操作
  161. out += self.downsample(x)
  162. # 应用ReLU激活函数
  163. out = F.relu(out)
  164. # 返回输出结果
  165. return out
  166. '''
  167. # 5. 构建ResNet
  168. # 定义ResNet类,继承自nn.Module,用于构建更复杂的网络结构
  169. '''
  170. class ResNet(nn.Module):
  171. def __init__(self, block, layers):
  172. '''
  173. 初始化ResNet网络。
  174. :param block: 基本的Block块对象,如BasicBlock或Bottleneck。
  175. :param layers: 不同层级的block数量,例如ResNet50为[3, 4, 6, 3]。
  176. '''
  177. super(ResNet, self).__init__() # 调用父类nn.Module的初始化方法。
  178. # 初始通道数设置为64
  179. self.inplanes = 64
  180. # 初始卷积层,输入通道为3(RGB图像),输出通道为64,卷积核大小为7x7,步长为2,填充为3,无偏置。
  181. self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
  182. # 批标准化层,用于conv1的输出
  183. self.bn1 = nn.BatchNorm2d(64)
  184. # ReLU激活函数,inplace=True表示直接在原变量上进行操作,节省内存。
  185. self.relu = nn.ReLU(inplace=True)
  186. # 最大池化层,池化核大小为3x3,步长为2,填充为1。
  187. self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
  188. # 根据输入的block和layers创建不同的层级
  189. # 创建四个层级(layer1至layer4),每个层级包含多个残差块
  190. self.layer1 = self._make_layer(block, 64, layers[0]) # 第一层,输出通道为64,block数量由layers[0]决定。
  191. self.layer2 = self._make_layer(block, 128, layers[1], stride=2) # 第二层,输出通道为128,block数量由layers[1]决定,步长为2。
  192. self.layer3 = self._make_layer(block, 256, layers[2], stride=2) # 第三层,输出通道为256,block数量由layers[2]决定,步长为2。
  193. self.layer4 = self._make_layer(block, 512, layers[3], stride=2) # 第四层,输出通道为512,block数量由layers[3]决定,步长为2。
  194. # 创建输出层
  195. self.layer5 = self._make_out_layer(in_channels=2048) # 输出层的输入通道为2048。
  196. # 平均池化层,池化核大小和步长均为2。
  197. self.avgpool = nn.AvgPool2d(2)
  198. # 最后的卷积层,将输出通道变为30,卷积核大小为3x3,步长为1,填充为1,无偏置。
  199. self.conv_end = nn.Conv2d(256, 30, kernel_size=3, stride=1, padding=1, bias=False)
  200. # 批标准化层,用于conv_end的输出
  201. self.bn_end = nn.BatchNorm2d(30)
  202. # 初始化模型参数
  203. for m in self.modules(): # 遍历模型中的所有模块
  204. if isinstance(m, nn.Conv2d): # 如果是卷积层
  205. n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels # 计算权重初始化的标准差
  206. m.weight.data.normal_(0, math.sqrt(2. / n)) # 使用正态分布初始化权重
  207. elif isinstance(m, nn.BatchNorm2d): # 如果是批标准化层
  208. m.weight.data.fill_(1) # 权重初始化为1
  209. m.bias.data.zero_() # 偏置初始化为0
  210. # 根据传入的layer个数和block创建
  211. # 输入层的构建
  212. def _make_layer(self, block, planes, blocks, stride=1):
  213. '''
  214. 根据给定的block类型、输出通道数、block数量和步长来创建一系列block组成的层。
  215. :param block: 要使用的block类型,是一个类,如BasicBlock或Bottleneck
  216. :param planes: 每个block的输出通道数
  217. :param blocks: 需要搭建的block数量
  218. :param stride: 第一个block的卷积层步长,用于控制是否进行下采样
  219. '''
  220. # 初始化下采样变量,用于第一个block可能需要的下采样操作
  221. downsample = None
  222. # 判断是否需要进行下采样
  223. # 如果步长不为1,或者当前层的输入通道数(self.inplanes)与block输出通道数(planes * block.expansion)不匹配,则进行下采样
  224. if stride != 1 or self.inplanes != planes * block.expansion:
  225. downsample = nn.Sequential(
  226. # 使用1x1的卷积层进行通道数转换和步长调整,使得残差连接可以相加
  227. nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
  228. # 批量归一化,稳定训练过程
  229. nn.BatchNorm2d(planes * block.expansion),
  230. )
  231. # 初始化层的列表
  232. layers = []
  233. # 添加第一个block,并可能进行下采样
  234. # 第一个block的输入通道数是self.inplanes,输出通道数是planes,步长是stride,可能包含下采样操作
  235. layers.append(block(self.inplanes, planes, stride, downsample))
  236. # 更新当前层的输入通道数为block的输出通道数
  237. self.inplanes = planes * block.expansion
  238. # 添加剩余的blocks,它们的输入通道数已经是更新后的self.inplanes,输出通道数保持不变
  239. for i in range(1, blocks):
  240. layers.append(block(self.inplanes, planes))
  241. # 将层的列表转换为Sequential模型,便于作为神经网络的一部分进行前向传播
  242. return nn.Sequential(*layers) # *操作符用于将列表中的元素解包为函数参数
  243. # 输出层的构建
  244. def _make_out_layer(self, in_channels):
  245. '''
  246. 根据输入通道数构建输出层,包含多个Output_Block。
  247. :param in_channels: 输入通道数
  248. '''
  249. # 初始化层的列表
  250. layers = []
  251. # 添加第一个Output_Block,并指定其输入通道数和输出通道数
  252. # 注意:这里使用了block_type参数来区分不同的Output_Block类型
  253. layers.append(Output_Block(in_planes=in_channels, planes=256, block_type='B'))
  254. # 添加两个相同类型的Output_Block,输入通道数和输出通道数都保持不变
  255. layers.append(Output_Block(in_planes=256, planes=256, block_type='A'))
  256. layers.append(Output_Block(in_planes=256, planes=256, block_type='A'))
  257. # 将层的列表转换为Sequential模型
  258. return nn.Sequential(*layers) # *操作符用于将列表中的元素解包为函数参数
  259. def forward(self, x):
  260. # x 为网络的输入,这里的前向传播方法描述了网络如何从输入x获得最终的输出
  261. # 经历共有的卷积和池化层
  262. # 第一个卷积层,对输入x进行卷积操作
  263. x = self.conv1(x)
  264. # 批量归一化层,对卷积后的结果进行归一化处理
  265. x = self.bn1(x)
  266. # ReLU激活函数,对归一化后的结果进行非线性变换
  267. x = self.relu(x)
  268. # 最大池化层,对激活后的结果进行下采样
  269. x = self.maxpool(x)
  270. # 经历各个block块
  271. # 第一个block层,对池化后的结果进行进一步的特征提取
  272. x = self.layer1(x)
  273. # 第二个block层,继续提取特征
  274. x = self.layer2(x)
  275. # 第三个block层
  276. x = self.layer3(x)
  277. # 第四个block层
  278. x = self.layer4(x)
  279. # 第五个block层,这里需要注意,通常标准的ResNet结构中只有4个layer,这里可能是自定义的扩展
  280. x = self.layer5(x)
  281. # 经历最终的输出
  282. # 平均池化层,对最后一个block的输出进行全局平均池化
  283. x = self.avgpool(x)
  284. # 卷积层,对池化后的结果进行进一步的卷积变换
  285. x = self.conv_end(x)
  286. # 批量归一化层,对卷积后的结果进行归一化处理
  287. x = self.bn_end(x)
  288. # Sigmoid激活函数,将特征映射到0-1之间
  289. x = F.sigmoid(x) # 归一化到0-1
  290. # 将输出构建为正确的shape
  291. # permute函数用于改变张量的维度顺序,这里将通道维度移动到最后一个位置
  292. # 假设输出shape原本为(-1, C, H, W),经过permute后变为(-1, H, W, C)
  293. x = x.permute(0, 2, 3, 1) # (-1,7,7,30)
  294. # 返回最终的输出x
  295. return x
  296. # 6. 构建不同的ResNet函数
  297. # 预训练下载链接
  298. model_urls = {
  299. 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
  300. 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
  301. 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
  302. 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
  303. 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
  304. }
  305. # 构建ResNet18
  306. def resnet18(pretrained=False, **kwargs):
  307. model = ResNet(Base_Block, [2, 2, 2, 2], **kwargs)
  308. # 是否预训练
  309. if pretrained:
  310. model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
  311. return model
  312. # 构建ResNet34
  313. def resnet34(pretrained=False, **kwargs):
  314. model = ResNet(Base_Block, [3, 4, 6, 3], **kwargs)
  315. if pretrained:
  316. model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
  317. return model
  318. # 构建ResNet50
  319. def resnet50(pretrained=False, **kwargs):
  320. model = ResNet(Senior_Block, [3, 4, 6, 3], **kwargs)
  321. if pretrained:
  322. model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
  323. return model
  324. # 构建ResNet101
  325. def resnet101(pretrained=False, **kwargs):
  326. model = ResNet(Senior_Block, [3, 4, 23, 3], **kwargs)
  327. if pretrained:
  328. model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
  329. return model
  330. # 构建ResNet152
  331. def resnet152(pretrained=False, **kwargs):
  332. model = ResNet(Senior_Block, [3, 8, 36, 3], **kwargs)
  333. if pretrained:
  334. model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
  335. return model

 3、定义损失

       在创建的目录结构 network 中创建Yolo_Loss.py文件,用于创建损失。

       YOLO的损失函数通常包括以下几个部分:

  1. 定位损失(Localization Loss):即长宽损失和坐标损失,用于衡量模型预测的目标边界框与真实边界框之间的差异。这通常使用均方误差(MSE)或其他类似的回归损失函数来计算。定位损失关注于边界框的中心点坐标(x, y)、宽度(w)和高度(h)的预测精度。
  2. 置信度损失(Confidence Loss):用于衡量模型预测的目标置信度与真实置信度之间的差异。置信度表示模型对预测目标存在的信心程度。对于每个边界框,模型会预测一个置信度分数,这个分数与真实标签(目标存在与否)进行比较,计算损失。
  3. 分类损失(Classification Loss):用于衡量模型预测的目标类别与真实类别之间的差异。这通常使用交叉熵损失函数来计算。对于每个目标,模型会预测一个类别概率分布,这个分布与真实标签进行比较,计算损失。

      在YOLO的损失函数中,这些部分通常会根据它们的重要性进行加权。例如,由于定位的准确性对于目标检测至关重要,因此定位损失通常会赋予较高的权重。而置信度损失和分类损失的权重则可以根据具体任务进行调整。

       此外,YOLO的损失函数还需要考虑如何处理没有目标的背景区域。这通常通过在损失函数中加入一个额外的项来实现,该项用于惩罚模型在背景区域产生高置信度的预测。

总的来说,YOLO的损失函数是一个综合了定位、置信度和分类损失的复合函数,旨在全面地衡量模型在目标检测任务上的性能。通过优化这个损失函数,我们可以训练出更加准确和鲁棒的目标检测模型。

  1. # 1. 导包
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. from torch.autograd import Variable
  6. # 2. 损失函数类
  7. class Yolo_Loss(nn.Module):
  8. def __init__(self, S=7, B=2, l_coord=5, l_noobj=0.5):
  9. '''
  10. :param S: Yolov1论文中的S,即划分的网格,默认为7
  11. :param B: Yolov1论文中的B,即多少个框预测,默认为2
  12. :param l_coord: 损失函数中的超参数,默认为5
  13. :param l_noobj: 同上,默认为0.5
  14. '''
  15. super(Yolo_Loss, self).__init__()
  16. # 初始化各个参数
  17. self.S = S
  18. self.B = B
  19. self.l_coord = l_coord
  20. self.l_noobj = l_noobj
  21. # 前向传播
  22. def forward(self, pred_tensor, target_tensor):
  23. # 获取batchsize大小 ,即每批次数据的个数
  24. N = pred_tensor.size()[0]
  25. # target_tensor[:, :, :, 4]是目标张量中用于标记 目标 是否存在的部分(通常第五列是置信度或目标存在标志)。
  26. # 这行代码生成一个布尔掩码,其中值为True的位置表示对应网格中存在目标,值为False则表示不存在。
  27. # 具有目标标签的索引值,此时shape为[batch,7,7]
  28. coo_mask = target_tensor[:, :, :, 4] > 0
  29. # 与上一行相反,这行代码生成一个布尔掩码,表示哪些网格中没有目标。
  30. # 不具有目标的标签索引值,此时shape为[batch,7,7]
  31. noo_mask = target_tensor[:, :, :, 4] == 0
  32. # unsqueeze(-1)在最后一个维度上增加一个维度,使得coo_mask的形状从[batch, 7, 7]变为[batch, 7, 7, 1]。
  33. # expand_as(target_tensor)将这个新维度扩展至与target_tensor相同的形状,这样便于后续操作。
  34. # 将shape变为[batch,7,7,30]
  35. coo_mask = coo_mask.unsqueeze(-1).expand_as(target_tensor)
  36. # 同上,对noo_mask也进行相同的操作。
  37. noo_mask = noo_mask.unsqueeze(-1).expand_as(target_tensor)
  38. # 使用coo_mask从pred_tensor中索引出存在目标的预测值,然后通过view操作改变形状,使得每个存在目标的预测框成为一个长度为30的向量。
  39. # 获取预测值中包含对象的所有点(共7*7个点),并转为[x,30]的形式,其中x表示有多少点中的框包含有对象
  40. coo_pred = pred_tensor[coo_mask].view(-1, 30)
  41. # 对上面获取的值进行处理
  42. # 从coo_pred中取前10个元素(对应预测框的坐标和宽度高度信息),然后改变形状为[2x, 5],其中x表示存在目标的网格数量。
  43. # 1. 转为box形式:box[x1,y1,w1,h1,c1],shape为[2x,5],因为每个单元格/点有两个预测框
  44. box_pred = coo_pred[:, :10].contiguous().view(-1, 5)
  45. # 从coo_pred中取后20个元素(对应类别预测信息)。
  46. # 2. 转为class信息,即30中后面的20个值
  47. class_pred = coo_pred[:, 10:]
  48. # 对目标张量进行与预测张量相同的操作,获取存在目标的真实值。
  49. # 同理,对真实值进行操作,方便对比计算损失值
  50. coo_target = target_tensor[coo_mask].view(-1, 30)
  51. # 从coo_target中取出坐标和宽度高度信息。
  52. box_target = coo_target[:, :10].contiguous().view(-1, 5)
  53. # 从coo_target中取出类别信息。
  54. class_target = coo_target[:, 10:]
  55. # 同上的操作,获取不包含对象的预测值、真实值
  56. # 使用noo_mask从pred_tensor中索引出不存在目标的预测值,并改变形状。
  57. noo_pred = pred_tensor[noo_mask].view(-1, 30)
  58. # 对目标张量执行相同的操作,获取不存在目标的真实值。
  59. noo_target = target_tensor[noo_mask].view(-1, 30)
  60. # 注意:上述代码只是将预测值和真实值进行了预处理,以便后续计算损失。实际损失的计算(比如坐标损失、置信度损失和类别损失)在这段代码之后,通常还需要根据YOLO的具体定义来编写。
  61. # 不包含物体grid ceil的置信度损失:即图中的D部分
  62. # 这部分代码计算的是那些不包含目标物体的网格中预测框的置信度损失。
  63. # 1. 自己创建一个索引
  64. # 创建一个与noo_pred形状相同的全零的ByteTensor(布尔掩码),用于标识需要计算的预测框置信度。
  65. noo_pred_mask = torch.cuda.ByteTensor(noo_pred.size())
  66. # 将这个掩码的所有元素初始化为0(即False),意味着初始时所有预测框的置信度都不被考虑。
  67. noo_pred_mask.zero_() # 将全部元素变为Flase的意思
  68. # 2. 将其它位置的索引置为0,唯独两个框的置信度位置变为1
  69. # 将所有网格中第一个预测框的置信度位置(索引为4,因为通常置信度是预测框向量的最后一个元素,且前面有4个坐标值)设置为1。
  70. noo_pred_mask[:, 4] = 1
  71. # 将所有网格中第二个预测框的置信度位置(索引为9,因为每个网格有两个预测框,每个预测框有5个值:4个坐标和1个置信度)设置为1。
  72. noo_pred_mask[:, 9] = 1
  73. # 3. 获取对应的值
  74. # 使用上面创建的掩码从noo_pred中提取出所有预测框的置信度值。
  75. noo_pred_c = noo_pred[noo_pred_mask] # noo pred只需要计算 c 的损失 size[-1,2]
  76. # 同样地,使用相同的掩码从noo_target中提取出目标置信度值,用于与预测值比较。
  77. noo_target_c = noo_target[noo_pred_mask]
  78. # 4. 计算损失值:均方误差
  79. # 使用均方误差(Mean Squared Error, MSE)计算预测置信度与目标置信度之间的差异,作为不包含物体网格的置信度损失。
  80. nooobj_loss = F.mse_loss(noo_pred_c, noo_target_c, size_average=False)
  81. # 计算包含物体的损失值
  82. # 创建几个全为False/0的变量,用于后期存储值
  83. # 创建一个与box_target形状相同的全零的ByteTensor,用于标识负责预测物体的网格中的预测框。
  84. coo_response_mask = torch.cuda.ByteTensor(box_target.size()) # 负责预测框
  85. # 将这个掩码的所有元素初始化为0。
  86. coo_response_mask.zero_()
  87. # 创建一个与box_target形状相同的全零的ByteTensor,用于标识不负责预测物体的网格中的预测框。
  88. coo_not_response_mask = torch.cuda.ByteTensor(box_target.size()) # 不负责预测的框的索引(因为一个cell两个预测框,而只有IOU最大的负责索引)
  89. # 将这个掩码的所有元素初始化为0。
  90. coo_not_response_mask.zero_()
  91. # 创建一个与box_target形状相同的全零的Tensor,用于存放每个预测框与目标框之间的IOU(交并比)值。这个值在后续计算损失时会用到。
  92. box_target_iou = torch.zeros(box_target.size()).cuda() # 具体的IOU值存放处
  93. # 由于一个单元格两个预测框,因此step=2
  94. # 在YOLO系列算法中,每个网格单元通常预测固定数量的预测框。这里假设每个单元有两个预测框,因此迭代步长为2。
  95. # 遍历所有目标框,因为每个目标框与一个单元格关联,且每个单元格有两个预测框。
  96. # 选择具有最佳IOU的预测框。
  97. for i in range(0, box_target.size()[0], 2): # choose the best iou box
  98. # 获取预测值中的两个box
  99. # 从预测框集合中取出当前单元格的两个预测框。这里box1是一个包含两个预测框的tensor,每个预测框有五个值:[x, y, w, h, c],分别代表中心坐标、宽度、高度和置信度。
  100. box1 = box_pred[i:i + 2] # [x,y,w,h,c]
  101. # 创建一个临时变量box1形状相同的tensor,用于存放左上角和右下角角坐标值,因为计算IOU需要
  102. box1_xyxy = Variable(torch.FloatTensor(box1.size()))
  103. # 将预测框的中心坐标和宽高转换为左上角和右下角的坐标,并进行归一化处理。self.S通常表示网格的大小,比如7x7的网格。
  104. box1_xyxy[:, :2] = box1[:, :2] / float(self.S) - 0.5 * box1[:, 2:4] # 原本(xc,yc)为7*7 所以要除以7
  105. # 同样地,计算预测框的右下角坐标。
  106. box1_xyxy[:, 2:4] = box1[:, :2] / float(self.S) + 0.5 * box1[:, 2:4]
  107. # 用同样的思路对真实值进行处理,不过不同的是真实值一个对象只有一个框
  108. # 从目标框集合中取出与当前单元格关联的目标框,并将其reshape为[-1, 5]的形式。
  109. box2 = box_target[i].view(-1, 5)
  110. # 创建一个与目标框形状相同的tensor,用于存放左上角和右下角坐标形式的目标框。
  111. box2_xyxy = Variable(torch.FloatTensor(box2.size()))
  112. # 将目标框的中心坐标和宽高转换为左上角和右下角的坐标,并进行归一化处理。
  113. box2_xyxy[:, :2] = box2[:, :2] / float(self.S) - 0.5 * box2[:, 2:4]
  114. # 计算目标框的右下角坐标。
  115. box2_xyxy[:, 2:4] = box2[:, :2] / float(self.S) + 0.5 * box2[:, 2:4]
  116. # 计算两者的IOU
  117. # 调用compute_iou函数计算两个预测框与目标框之间的IOU值。这里box1_xyxy[:, :4]包含两个预测框的坐标,box2_xyxy[:, :4]包含一个目标框的坐标。
  118. iou = self.compute_iou(box1_xyxy[:, :4], box2_xyxy[:, :4]) # 前者shape为[2,4],后者为[1,4]
  119. # 获取两者IOU最大的值(max_iou)和对应的索引(max_index),因为一个cell有两个预测框,一般而言取IOU最大的作为预测框
  120. max_iou, max_index = iou.max(0)
  121. # 将索引移至GPU上(如果之前不在GPU上的话)。
  122. max_index = max_index.data.cuda()
  123. # 将IOU最大的索引设置为1,即表示这个框负责预测
  124. coo_response_mask[i + max_index] = 1
  125. # 将box_target_iou转为PyTorch的Variable,并移到GPU上(如果有的话),以支持自动微分。
  126. box_target_iou = Variable(box_target_iou).cuda()
  127. # 获取负责预测框的值、IOU值和真实框的值
  128. # 从box_pred中筛选出负责预测的框(即IOU最大的框),并改变其形状为[-1, 5]。
  129. box_pred_response = box_pred[coo_response_mask].view(-1, 5)
  130. # 从box_target_iou中筛选出负责预测的框的IOU值,并改变其形状。
  131. box_target_response_iou = box_target_iou[coo_response_mask].view(-1, 5)
  132. # 从box_target中筛选出负责预测的框的真实值,并改变其形状。
  133. box_target_response = box_target[coo_response_mask].view(-1, 5)
  134. # 这个对应的是图中的部分C,负责预测框的损失
  135. # 计算负责预测框的置信度损失,使用均方误差损失函数(MSE Loss)。
  136. contain_loss = F.mse_loss(box_pred_response[:, 4], box_target_response_iou[:, 4], size_average=False)
  137. # 1. 计算坐标损失,即图中的A和B部分
  138. # 计算负责预测框的坐标损失,包括中心坐标和宽高的损失。注意,宽高损失计算时使用了开方,这通常是为了改善损失函数的性能。
  139. loc_loss = F.mse_loss(box_pred_response[:, :2], box_target_response[:, :2], size_average=False) + F.mse_loss(
  140. torch.sqrt(box_pred_response[:, 2:4]), torch.sqrt(box_target_response[:, 2:4]), size_average=False)
  141. # 获取不负责预测框的值、真实值
  142. # 从box_pred中筛选出不负责预测的框,并改变其形状。
  143. box_pred_not_response = box_pred[coo_not_response_mask].view(-1, 5)
  144. # 从box_target中筛选出不负责预测的框的真实值,并改变其形状。
  145. box_target_not_response = box_target[coo_not_response_mask].view(-1, 5)
  146. # 将不负责预测框的置信度真实值设为0,因为这些框不应包含任何对象。
  147. box_target_not_response[:, 4] = 0 # 将真实值置为0
  148. # 2. 计算不负责预测框的损失值,即图中的部分C
  149. # 计算不负责预测框的置信度损失。
  150. not_contain_loss = F.mse_loss(box_pred_not_response[:, 4], box_target_not_response[:, 4], size_average=False)
  151. # 3. 类别损失,即图中的E部分
  152. # 计算类别损失,使用均方误差损失函数。
  153. class_loss = F.mse_loss(class_pred, class_target, size_average=False)
  154. # 返回总损失,它是坐标损失、负责预测框的置信度损失、不负责预测框的置信度损失、类别损失以及一个未显示的`nooobj_loss`的加权和,最后除以N进行归一化。
  155. return (self.l_coord * loc_loss + contain_loss + not_contain_loss + self.l_noobj * nooobj_loss + class_loss) / N
  156. # 计算IOU的函数
  157. def compute_iou(self, box1, box2):
  158. '''
  159. :param box1: 预测的box,一般为[2,4]
  160. :param box2: 真实的box,一般为[1,4]
  161. :return:
  162. '''
  163. # 获取各box个数
  164. N = box1.size(0)
  165. M = box2.size(0)
  166. # 计算两者中左上角左边较大的
  167. lt = torch.max(
  168. box1[:, :2].unsqueeze(1).expand(N, M, 2), # [N,2] -> [N,1,2] -> [N,M,2]
  169. box2[:, :2].unsqueeze(0).expand(N, M, 2), # [M,2] -> [1,M,2] -> [N,M,2]
  170. )
  171. # 计算两者右下角左边较小的
  172. rb = torch.min(
  173. box1[:, 2:].unsqueeze(1).expand(N, M, 2), # [N,2] -> [N,1,2] -> [N,M,2]
  174. box2[:, 2:].unsqueeze(0).expand(N, M, 2), # [M,2] -> [1,M,2] -> [N,M,2]
  175. )
  176. # 计算两者相交部分的长、宽
  177. wh = rb - lt # [N,M,2]
  178. # 如果长、宽中有小于0的,表示可能没有相交趋于,置为0即可
  179. wh[wh < 0] = 0 # clip at 0
  180. inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
  181. # 计算各个的面积
  182. # box1的面积
  183. area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1]) # [N,]
  184. # box2的面积
  185. area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1]) # [M,]
  186. area1 = area1.unsqueeze(1).expand_as(inter) # [N,] -> [N,1] -> [N,M]
  187. area2 = area2.unsqueeze(0).expand_as(inter) # [M,] -> [1,M] -> [N,M]
  188. # IOu值,交集除以并集,其中并集为两者的面积和减去交集部分
  189. iou = inter / (area1 + area2 - inter)
  190. return iou

4、定义训练过程

      在创建的目录结构 network 中创建train.py文件,用于进行训练。

  1. 导入所需的包
  2. 定义基本参数
  3. 创建模型并使用预训练参数
  4. 损失函数,优化器,并将模型、损失函数放入GPU中
  5. 加载数据
  6. 训练:循环50次,第20和第35次调整学习率,每次对数据集中按批次执行训练和验证
  1. # 1. 导入所需的包
  2. import warnings
  3. from YOLOv1.network.Yolo_ResNet import resnet50
  4. from YOLOv1.utils.yolo_dataset import Yolo_Dataset
  5. from YOLOv1_pytorch.utils.Yolo_Loss import Yolo_Loss
  6. warnings.filterwarnings("ignore")
  7. from tqdm import tqdm
  8. import torch
  9. from torch.utils.data import DataLoader
  10. import torchvision.transforms as T
  11. from torchvision import models
  12. # 2. 定义基本参数
  13. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  14. batch_size = 2 # 根据自己的电脑设定
  15. epochs = 50
  16. lr = 0.01
  17. file_root = './data/VOC2012/JPEGImages/' # 需要根据的实际路径修改
  18. # 3. 创建模型并使用预训练参数
  19. pytorch_resnet = models.resnet50(pretrained=True) # 官方的resnet50预训练模型
  20. model = resnet50() # 创建自己的resnet50,
  21. # 接下来就是让自己的模型去继承官方的权重参数
  22. pytorch_state_dict = pytorch_resnet.state_dict()
  23. model_state_dict = model.state_dict()
  24. for k in pytorch_state_dict.keys():
  25. # 调试: 看看模型哪些有没有问题
  26. # print(k)
  27. # 如果自己的模型和官方的模型key相同,并且不是fc层,则继承过来
  28. if k in model_state_dict.keys() and not k.startswith('fc'):
  29. model_state_dict[k] = pytorch_state_dict[k]
  30. # 4. 损失函数,优化器,并将模型、损失函数放入GPU中
  31. loss = Yolo_Loss()
  32. optimizer = torch.optim.SGD(model.parameters(),lr=lr,momentum=0.9,weight_decay=5e-4)
  33. model.to(device)
  34. loss.to(device)
  35. # 5. 加载数据
  36. train_dataset = Yolo_Dataset(root=file_root, list_file='./utils/voctrain.txt', train=True, transforms = [T.ToTensor()])
  37. train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,drop_last=True)
  38. test_dataset = Yolo_Dataset(root=file_root, list_file='./utils/voctest.txt', train=False, transforms = [T.ToTensor()])
  39. test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=True,drop_last=True)
  40. # 6. 训练
  41. # 打印一些基本的信息
  42. print('starting train the model')
  43. print('the train_dataset has %d images' % len(train_dataset))
  44. print('the batch_size is ',batch_size)
  45. # 定义一个最佳损失值
  46. best_test_loss = 0
  47. # 开始训练
  48. for e in range(epochs):
  49. model.train()
  50. # 调整学习率
  51. if e == 20:
  52. print('change the lr')
  53. optimizer.param_groups[0]['lr'] /= 10
  54. if e == 35:
  55. print('change the lr')
  56. optimizer.param_groups[0]['lr'] /= 10
  57. # 进度条显示
  58. tqdm_tarin = tqdm(train_loader)
  59. # 定义损失变量
  60. total_loss = 0.
  61. for i,(images,target) in enumerate(tqdm_tarin):
  62. # 将变量放入设备中
  63. images,target = images.to(device),target.to(device)
  64. # 训练--损失等
  65. pred = model(images)
  66. loss_value = loss(pred,target)
  67. total_loss += loss_value.item()
  68. #在开始新的梯度计算之前,调用优化器的zero_grad()方法将模型参数的梯度清零。这是因为在PyTorch中,梯度是累积的,如果不清零,新的梯度会累加到旧的梯度上。
  69. optimizer.zero_grad()
  70. #对损失值loss_value调用backward()方法,计算损失相对于模型参数的梯度。这些梯度将被用于接下来的参数更新步骤。
  71. loss_value.backward()
  72. #调用优化器的step()方法,根据计算出的梯度更新模型的参数。这一步是模型训练的关键,它使得模型在训练过程中逐渐优化。
  73. optimizer.step()
  74. # 打印一下损失值
  75. if (i+1) % 5 == 0:
  76. tqdm_tarin.desc = 'train epoch[{}/{}] loss:{:.6f}'.format(e+1,epochs,total_loss/(i+1))
  77. # 启用验证模式
  78. model.eval()
  79. validation_loss = 0.0
  80. tqdm_test = tqdm(test_loader)
  81. for i, (images, target) in enumerate(tqdm_test):
  82. images, target = images.to(device),target.to(device)
  83. pred = model(images)
  84. loss_value = loss(pred, target)
  85. validation_loss += loss_value.item()
  86. validation_loss /= len(test_loader)
  87. # 显示验证集的损失值
  88. print('In the test step,the average loss is %.6f' % validation_loss)
  89. # 如果最佳损失值大于验证集的损失,意味着当前训练很好
  90. # 这一点需要设置好最佳的损失值,不容易设置
  91. # 是否启用看大家心情
  92. # if best_test_loss > validation_loss:
  93. # best_test_loss = validation_loss
  94. # print('get best test loss %.5f' % best_test_loss)
  95. # torch.save(model.state_dict(), './weights/best.pth')
  96. # 记得最后保存参数
  97. torch.save(model.state_dict(), './weights/yolo.pth')

5、定义预测方法

       在创建的目录结构 network 中创建predict.py文件,用于进行预测。同时在目录test_images中存入预测的图片(在最下面,自己进行保存,并进行命名为1、2、3、4、5、6)

  1. # author: baiCai
  2. # 1. 导包
  3. import os
  4. import random
  5. import torch
  6. from torch.autograd import Variable
  7. import torchvision.transforms as transforms
  8. import cv2
  9. from matplotlib import pyplot as plt
  10. import numpy as np
  11. import warnings
  12. from YOLOv1_pytorch.network.My_ResNet import resnet50
  13. warnings.filterwarnings('ignore')
  14. # 2. 定义一些基本的参数
  15. # 类别索引
  16. VOC_CLASSES = (
  17. 'aeroplane', 'bicycle', 'bird', 'boat',
  18. 'bottle', 'bus', 'car', 'cat', 'chair',
  19. 'cow', 'diningtable', 'dog', 'horse',
  20. 'motorbike', 'person', 'pottedplant',
  21. 'sheep', 'sofa', 'train', 'tvmonitor')
  22. # 画矩形框的时候用到的颜色变量
  23. Color = [[0, 0, 0],
  24. [128, 0, 0],
  25. [0, 128, 0],
  26. [128, 128, 0],
  27. [0, 0, 128],
  28. [128, 0, 128],
  29. [0, 128, 128],
  30. [128, 128, 128],
  31. [64, 0, 0],
  32. [192, 0, 0],
  33. [64, 128, 0],
  34. [192, 128, 0],
  35. [64, 0, 128],
  36. [192, 0, 128],
  37. [64, 128, 128],
  38. [192, 128, 128],
  39. [0, 64, 0],
  40. [128, 64, 0],
  41. [0, 192, 0],
  42. [128, 192, 0],
  43. [0, 64, 128]]
  44. # 3. 解码函数
  45. def decoder(pred):
  46. '''
  47. :param pred: batchx7x7x30,但是预测的时候一般一张图片一张的放,因此batch=1
  48. :return: box[[x1,y1,x2,y2]] label[...]
  49. '''
  50. # 定义一些基本的参数
  51. grid_num = 7 # 网格划分标准大小
  52. #初始化三个空列表,用于存储解码后的边界框坐标、类别索引和对应的概率值。
  53. boxes=[]
  54. cls_indexs=[]
  55. probs = []
  56. #计算网格单元的大小(或称为缩放因子),这对于将模型输出的坐标从网格空间转换到图像空间是必要的。
  57. cell_size = 1./grid_num # 缩放因子
  58. # 获取一些值
  59. pred = pred.data # 预测值的数据:1*7*7*30
  60. pred = pred.squeeze(0) # 预测值的数据:7x7x30
  61. contain1 = pred[:,:,4].unsqueeze(2) # 先获取第一个框的置信度,然后升维变为7*7*1
  62. contain2 = pred[:,:,9].unsqueeze(2) # 同上,只是为第二个框
  63. contain = torch.cat((contain1,contain2),2) # 拼接在一起,变为7*7*2
  64. mask1 = contain > 0.1 #大于阈值0.1,设置为True
  65. mask2 = (contain==contain.max()) # 找出置信度最大的,设置为True
  66. mask = (mask1+mask2).gt(0) # 将mask1+mask2,让其中大于0的设置为True
  67. # 开始迭代每个单元格,即7*7个
  68. for i in range(grid_num):
  69. for j in range(grid_num):
  70. # 迭代两个预测框
  71. for b in range(2):
  72. # 如果mask为1,表示这个框是最大的置信度框
  73. if mask[i,j,b] == 1:
  74. # 获取坐标值
  75. box = pred[i,j,b*5:b*5+4]
  76. # 获取置信度值
  77. contain_prob = torch.FloatTensor([pred[i,j,b*5+4]])
  78. # 将7*7的坐标,归一化
  79. xy = torch.FloatTensor([j,i])*cell_size #cell左上角 up left of cell
  80. #
  81. box[:2] = box[:2]*cell_size + xy
  82. # 将[cx,cy,w,h]转为[x1,xy1,x2,y2]
  83. box_xy = torch.FloatTensor(box.size()) # 重新创建一个变量存储值
  84. box_xy[:2] = box[:2] - 0.5*box[2:] # 这个就是中心坐标加减宽度/高度得到左上角/右下角坐标
  85. box_xy[2:] = box[:2] + 0.5*box[2:]
  86. # 获取最大的概率和类别索引值
  87. max_prob,cls_index = torch.max(pred[i,j,10:],0)
  88. # 如果置信度 * 类别概率 > 0.1,即说明有一定的可信度
  89. # 那么把值加入各个变量列表中
  90. if float((contain_prob*max_prob)[0]) > 0.1:
  91. boxes.append(box_xy.view(1,4))
  92. cls_indexs.append(torch.tensor([cls_index.item()]))
  93. probs.append(contain_prob*max_prob)
  94. # 如果boxes为0,表示没有框,返回0
  95. if len(boxes) ==0:
  96. boxes = torch.zeros((1,4))
  97. probs = torch.zeros(1)
  98. cls_indexs = torch.zeros(1)
  99. # 否则,进行处理,就是简单把原来的列表值[tensor,tensor]转为tensor的形式
  100. # 里面的值不变
  101. else:
  102. boxes = torch.cat(boxes,0) #(n,4)
  103. probs = torch.cat(probs,0) #(n,)
  104. cls_indexs = torch.cat(cls_indexs,0) #(n,)
  105. # 后处理——NMS
  106. keep = mns(boxes,probs)
  107. # 返回值
  108. return boxes[keep],cls_indexs[keep],probs[keep]
  109. # 4. NMS处理
  110. def mns(bboxes,scores,threshold=0.5):
  111. '''
  112. :param bboxes: bboxes(tensor) [N,4]
  113. :param scores: scores(tensor) [N,]
  114. :param threshold: 阈值
  115. :return: 返回过滤后的框
  116. '''
  117. # 获取各个框的坐标值
  118. x1 = bboxes[:,0]
  119. y1 = bboxes[:,1]
  120. x2 = bboxes[:,2]
  121. y2 = bboxes[:,3]
  122. # 计算面积
  123. areas = (x2-x1) * (y2-y1)
  124. # 将置信度按照降序排序,并获取排序后的各个置信度在这个顺序中的索引
  125. _,order = scores.sort(0,descending=True)
  126. keep = []
  127. # 判断order中的元素个数是否大于0
  128. while order.numel() > 0:
  129. # 如果元素个数只剩下一个了,结束循环
  130. if order.numel() == 1:
  131. i = order.item()
  132. keep.append(i)
  133. break
  134. # 获取最大置信度的索引
  135. i = order[0]
  136. keep.append(i)
  137. # 对后面的元素坐标进行截断处理
  138. xx1 = x1[order[1:]].clamp(min=x1[i]) # min指的是小于它的设置为它的值,大于它的不管
  139. yy1 = y1[order[1:]].clamp(min=y1[i])
  140. xx2 = x2[order[1:]].clamp(max=x2[i])
  141. yy2 = y2[order[1:]].clamp(max=y2[i])
  142. # 此时的xx1,yy1等是排除了目前选中的框的,即假设x1有三个元素,那么xx1只有2个元素
  143. # 获取排序后的长和宽以及面积,如果小于0则设置为0
  144. w = (xx2-xx1).clamp(min=0)
  145. h = (yy2-yy1).clamp(min=0)
  146. inter = w*h
  147. # 准备更新order、
  148. # 计算选中的框和剩下框的IOU值
  149. ovr = inter / (areas[i] + areas[order[1:]] - inter)
  150. # 如果 IOU小于设定的阈值,说明需要保存下来继续筛选(NMS原理)
  151. ids = (ovr<=threshold).nonzero().squeeze()
  152. if ids.numel() == 0:
  153. break
  154. order = order[ids+1]
  155. return torch.LongTensor(keep)
  156. # 5. 预测函数
  157. def predict_single(model, image_name, root_path=''):
  158. result = [] # 保存结果的变量
  159. # 打开图片
  160. image = cv2.imread(root_path + image_name)
  161. h, w, _ = image.shape
  162. # resize为模型的输入大小,即448*448
  163. img = cv2.resize(image, (448, 448))
  164. # 由于我们模型那里定义的颜色模式为RGB,因此这里需要转换
  165. mean = (123, 117, 104) # RGB均值
  166. img = img - np.array(mean, dtype=np.float32)
  167. # 预处理
  168. transform = transforms.Compose([transforms.ToTensor(), ])
  169. img = transform(img)
  170. img = Variable(img[None, :, :, :], volatile=True)
  171. img = img.cuda()
  172. # 开始预测
  173. pred = model(img) # 1x7x7x30
  174. pred = pred.cpu()
  175. # 解码
  176. boxes, cls_indexs, probs = decoder(pred)
  177. # 开始迭代每个框
  178. for i, box in enumerate(boxes):
  179. # 获取相关坐标,只是需要把原来归一化后的坐标转回去
  180. x1 = int(box[0] * w)
  181. x2 = int(box[2] * w)
  182. y1 = int(box[1] * h)
  183. y2 = int(box[3] * h)
  184. # 获取类别索引、概率等值
  185. cls_index = cls_indexs[i]
  186. cls_index = int(cls_index) # convert LongTensor to int
  187. prob = probs[i]
  188. prob = float(prob)
  189. # 把这些值集中放入一个变量中返回
  190. result.append([(x1, y1), (x2, y2), VOC_CLASSES[cls_index], image_name, prob])
  191. return result
  192. if __name__ == '__main__':
  193. # 慢慢的显示
  194. import time
  195. # 创建模型,加载参数
  196. model = resnet50()
  197. model.load_state_dict(torch.load('./weights/yolo.pth'))
  198. model.eval()
  199. model.cuda()
  200. # 设置图片路径
  201. base_path = './test_images/'
  202. # base_path = '../data/VOC2012/JPEGImages/'
  203. image_name_list = [base_path+i for i in os.listdir(base_path)]
  204. # 打乱顺序
  205. random.shuffle(image_name_list)
  206. print('stating predicting....')
  207. for image_name in image_name_list:
  208. image = cv2.imread(image_name)
  209. image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  210. result = predict_single(model, image_name)
  211. # 画矩形框和对应的类别信息
  212. for left_up, right_bottom, class_name, _, prob in result:
  213. # 获取颜色
  214. color = Color[VOC_CLASSES.index(class_name)]
  215. # 画矩形
  216. cv2.rectangle(image, left_up, right_bottom, color, 2)
  217. # 获取类型信息和对应概率,此时为str类型
  218. label = class_name + str(round(prob, 2))
  219. # 把类别和概率信息写上,还要为这个信息加上一个矩形框
  220. text_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
  221. p1 = (left_up[0], left_up[1] - text_size[1])
  222. cv2.rectangle(image, (p1[0] - 2 // 2, p1[1] - 2 - baseline), (p1[0] + text_size[0], p1[1] + text_size[1]),
  223. color, -1)
  224. cv2.putText(image, label, (p1[0], p1[1] + baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, 8)
  225. # 显示图片
  226. plt.figure()
  227. plt.imshow(image)
  228. plt.show()
  229. time.sleep(2)
  230. # 是否保存结果图片
  231. cv2.imwrite('./test_images/result.jpg', image)

6、验证图片

 

7、完整代码和权重下载

链接:https://pan.baidu.com/s/1nJbfAxTUshCy4giFPoy7sA 
提取码:yov1 

8、总结

需要改进:

  首先,肯定是batch_size的大小,如果有得选,我也想把batch_size改大一点_
  其次,是优化器的选择和其参数的设置,本次选用的SGD,参数都是默认设置的,我在想是否改为Adam会好一点
  另外,本次学习率从0.01开始,分别在20和35epoch时除以10,感觉初始的学习率有点大了,是否可以减小学习率?
  最后,所谓微调,我们继承了ResNet50的权重,但是这个ResNet50是在尺度为224*224下训练的,而我们的目标检测将分辨率改为了448*448,是否需要先进行一定程度的微调才来训练yolov1值得思考。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/从前慢现在也慢/article/detail/480555
推荐阅读
相关标签
  

闽ICP备14008679号