一.代码资源下载：
1.代码下载：
链接
2.打开下载的项目，并新建文件夹data
3.下载pascal VOC 2007 数据集和 small_yolo.ckpt文件，将下载好的文件分别放入data文件夹下的pascal_voc和weight文件夹下。
数据集是训练时需要的，
如果不想训练直接用别人训练好的模型可以只下载权重文件:
权重文件链接：https://pan.baidu.com/s/12uZLQ5HJQB6OAqDEKB_4Ww
提取码：rxo5
数据文件链接：https://pan.baidu.com/s/1HZ4jgpj_DjWARbofqpZKng
提取码：8w3g
二、源码解析
项目全部文件：
2.1 data文件夹
data文件夹下：pascal_voc存放数据、weight文件下存放权重。
2.2 test文件夹
存放测试的图片数据
2.3 utils文件夹
(1) pascal_voc.py
总体预览：本py文件用来获取数据集的图片和标注，共包含
6个函数，本py文件需要config.py配置文件中的一些参数。
import os
import xml.etree.ElementTree as ET
import numpy as np
import cv2
import pickle
import copy
import yolo.config as cfg #导入config.py文件
#本py文件用来获取训练图片和标签，为了便于理解将函数的定义顺序进行了调整
class pascal_voc(object):
'''
准备训练或者测试的数据
args:
phase：传入字符串 ‘train‘：表示训练,‘test‘：测试
rebuild：是否重新创建数据集的标签文件，保存在缓存文件夹下
'''
# 初始化设置
def __init__(self, phase, rebuild=False):
# config.py中：PASCAL_PATH设置----data\pascal_voc
self.devkil_path = os.path.join(cfg.PASCAL_PATH, 'VOCdevkit')#devkil_path-----》data\pascal_voc\VOCdevkit
self.data_path = os.path.join(self.devkil_path, 'VOC2007')# data_path-----》data\pascal_voc\VOCdevkit\VOC2007
# config.py中：CACHE_PATH设置----data\pascal_voc\cache
# BATCH_SIZE设置----BATCH_SIZE = 45
# IMAGE_SIZE设置----IMAGE_SIZE = 448
# CELL_SIZE设置----CELL_SIZE = 7
# CLASSES设置----CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus','car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
# 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
# FLIPPED设置----FLIPPED = True
self.cache_path = cfg.CACHE_PATH##catch文件所在路径
self.batch_size = cfg.BATCH_SIZE
self.image_size = cfg.IMAGE_SIZE
self.cell_size = cfg.CELL_SIZE
self.classes = cfg.CLASSES#数据集中的所有类列表
self.flipped = cfg.FLIPPED# ##图片是否采用水平镜像扩充训练集
#类别名->索引的dict.将类列表，转换为字典，字典的键为类名，字典的值为序号。例{‘猫’：1，‘狗’：2}
self.class_to_ind = dict(zip(self.classes, range(len(self.classes))))
#__init__（）传入的参数：phase，rebuild
self.phase = phase#取train或者test，表示训练或者测试
self.rebuild = rebuild# #是否重新创建数据集标签文件
#从gt_labels加载数据，cursor表明当前读取到第几个
self.cursor = 0
self.epoch = 1
# gt_labels存放数据集标签 是一个list 每一个元素都是一个dict，对应一个图片 及其标注
self.gt_labels = None
# 加载数据集标签 初始化gt_labels
self.prepare()
#获取一张图片的标注信息,和该图片中目标的数量
def load_pascal_annotation(self, index):
'''
:param index:
index：图片文件的index
:return:
label：标签 [7,7,25]
0:1：置信度，表示这个地方是否有目标
1:5：目标边界框 目标中心，宽度和高度(这里是实际值，没有归一化)
5:25：目标的类别
len(objs)：objs对象长度
'''
# 获取图片文件名路径并读取
# data_path-----》data\pascal_voc\VOCdevkit\VOC2007
# imname-----》data\pascal_voc\VOCdevkit\VOC2007\JPEGImages\index.jpg 这里index是函数的参数，传入的是图片的名字
imname = os.path.join(self.data_path, 'JPEGImages', index + '.jpg')#获取图片路径
# 读取数据
im = cv2.imread(imname)#读取图片
# 宽和高缩放比例
h_ratio = 1.0 * self.image_size / im.shape[0]#resize后的图片是原图的几倍，后续边界框要根据这个resize的比例进行缩放
w_ratio = 1.0 * self.image_size / im.shape[1]#resize后的图片是原图的几倍
im = cv2.resize(im, [self.image_size, self.image_size])
# 图片文件的标注xml文件
# data_path-----》data\pascal_voc\VOCdevkit\VOC2007
# filename-----》data\pascal_voc\VOCdevkit\VOC2007\Annotations\index.xml
filename = os.path.join(self.data_path, 'Annotations', index + '.xml')#获取图片标注文件的路径
tree = ET.parse(filename)
objs = tree.findall('object')#获取xml文件里的object
for obj in objs:#遍历所有的object
#获取object标签里面的bndbox标签，取边界框信息
bbox = obj.find('bndbox')
#当图片缩放到image_size时，边界框也进行同比例缩放
#因为原图片进行了resize比率为 w_ratio和 h_ratio，因此标注框也要进行同比率缩放
x1 = max(min((float(bbox.find('xmin').text) - 1) * w_ratio, self.image_size - 1), 0)
y1 = max(min((float(bbox.find('ymin').text) - 1) * h_ratio, self.image_size - 1), 0)
x2 = max(min((float(bbox.find('xmax').text) - 1) * w_ratio, self.image_size - 1), 0)
y2 = max(min((float(bbox.find('ymax').text) - 1) * h_ratio, self.image_size - 1), 0)
# 根据图片的分类名 ->类别index 转换
#获取object标签里面的name标签：获取该object的类别名
cls_ind = self.class_to_ind[obj.find('name').text.lower().strip()]#将该object的类别名通过class_to_ind字典key转变为索引
'''计算边框（x,y,w,h）(没有归一化),xml中的边界框信息是左上和右下坐标'''
boxes = [(x2 + x1) / 2.0, (y2 + y1) / 2.0, x2 - x1, y2 - y1]#边界框信息：中心坐标（x,y）以及边界框宽和高
'''计算中心点位于哪个网格单元'''
x_ind = int(boxes[0] * self.cell_size / self.image_size)
y_ind = int(boxes[1] * self.cell_size / self.image_size)
label = np.zeros((self.cell_size, self.cell_size, 25)) #（7,7,25）的张量
if label[y_ind, x_ind, 0] == 1:#判断格子里是否有目标了如果有目标就跳过这个格子
continue
#''' 边界框信息：网格单元行，网格单元列，[置信度，边界框信息，类别]'''
# 置信度，表示这个地方有物体
label[y_ind, x_ind, 0] = 1
# 物体边界框
label[y_ind, x_ind, 1:5] = boxes
# 物体的类别
label[y_ind, x_ind, 5 + cls_ind] = 1
return label, len(objs)#返回标签(属于哪个网格单元，边界框信息，类别信息)和目标个数
#获取数据集标签：函数返回一个列表，列表的内容为字典格式：{'imname': 图片路径, 'label': 图片标注信息,'flipped': False}
def load_labels(self):
'''
加载数据集标签
:return:
gt_labels：是一个list 每一个元素对应一张图片，是一个dict
imname：图片文件路径
label：图片文件对应的标签 [7,7,25]的矩阵
flipped：是否使用水平镜像？ 设置为False
'''
# 缓冲文件名：即用来保存数据集标签的文件
#cathe.path----data\pascal_voc\cache
#cache_file----data\pascal_voc\cache\pascal_train_gt_labels.pkl 假设phase为train
cache_file = os.path.join(self.cache_path, 'pascal_' + self.phase + '_gt_labels.pkl')
#如果 cache_file文件存在并且不需要rebulid那么直接从cache中读取文件
if os.path.isfile(cache_file) and not self.rebuild:
#print('Loading gt_labels from: ' + cache_file)
with open(cache_file, 'rb') as f:
gt_labels = pickle.load(f)#loads 将pickle数据转换为python的数据结构
return gt_labels
#如果cache_file文件不存在，就创建这个文件
if not os.path.exists(self.cache_path):
os.makedirs(self.cache_path)
#加载训练或者测试集 图片的名字，放入列表image_index中
if self.phase == 'train':
# data_path-----》data\pascal_voc\VOCdevkit\VOC2007
#加载训练集的图片名字集合文件：trainval.txt里面每一行数据是一张训练图片的名字
txtname = os.path.join(self.data_path, 'ImageSets', 'Main', 'trainval.txt')
else:
# data_path-----》data\pascal_voc\VOCdevkit\VOC2007
#加载测试集的图片名字集合文件：test.txt里面每一行数据是一张测试图片的名字
txtname = os.path.join(self.data_path, 'ImageSets', 'Main', 'test.txt')
with open(txtname, 'r') as f:#将训练集或测试集的图片名放入列表image_index
self.image_index = [x.strip() for x in f.readlines()]#strip() 方法用于移除字符串头尾指定的字符(默认为空格或换行符)或字符序列
# 存放图片的标签，图片路径，是否使用水平镜像？
gt_labels = []
for index in self.image_index:#遍历训练集的每个图片
# 读取某图片的标签label [7,7,25]
label, num = self.load_pascal_annotation(index) #调用load_pascal_annotation()函数读取index这一图片的标签和所含目标的数量
if num == 0:
continue
# 某图片路径
imname = os.path.join(self.data_path, 'JPEGImages', index + '.jpg')#图片路径
#保存该图片的信息：将图片路径、图片标签、图片是否翻转这三个信息存入字典，并将这一字典存入列表中
gt_labels.append({'imname': imname, 'label': label,'flipped': False})
# 保存，将数据集标签信息保存到cache_file
with open(cache_file, 'wb') as f:#将 gt_labels写入cache_file文件
pickle.dump(gt_labels, f)
return gt_labels#返回整个数据集的标签
# 获取数据集标签：返回的是经过数据增强后的数据集标签，prepare()函数调用load_labels()函数，
# 加载所有数据集的标签，保存在遍历gt_labels集合中
def prepare(self):
'''
初始化数据集的标签，保存在变量gt_labels中
return:
gt_labels:返回数据集的标签 是一个list 每一个元素对应一张图片，是一个dict
imname：图片文件路径
label：图片文件对应的标签 [7,7,25]的矩阵
flipped：是否使用水平镜像？ 设置为False
本函数实质是对训练集进行图像增强后（翻转），训练集的标签信息加载
'''
#加载数据集标签：调用load_labels()函数
gt_labels = self.load_labels()#调用load_labels()函数，加载训练集图像的信息（图片，图片标注，是否翻转）
#如果水平镜像，则追加一倍的训练数据集
if self.flipped:#如果对图片翻转进行下面操作
print('Appending horizontally-flipped training examples ...')
gt_labels_cp = copy.deepcopy(gt_labels)#深度拷贝
# 遍历每一个图片标签
for idx in range(len(gt_labels_cp)):
gt_labels_cp[idx]['flipped'] = True #设置flipped属性为True
gt_labels_cp[idx]['label'] =gt_labels_cp[idx]['label'][:, ::-1, :] #目标所在格子也进行水平镜像 [7,7,25]
for i in range(self.cell_size):
for j in range(self.cell_size):
if gt_labels_cp[idx]['label'][i, j, 0] == 1: #置信度==1，表示这个格子有目标
gt_labels_cp[idx]['label'][i, j, 1] = self.image_size - 1 -gt_labels_cp[idx]['label'][i, j, 1] #中心的x坐标水平镜像
# 追加数据集的标签 后面的是由原数据集标签扩充的水平镜像数据集标签
gt_labels += gt_labels_cp
np.random.shuffle(gt_labels)#随机打乱数据集的标签
self.gt_labels = gt_labels
'''
gt_labels:数据集的标签 是一个list 每一个元素对应一张图片，是一个dict
imname：图片文件路径
label：图片文件对应的标签 [7,7,25]的矩阵
flipped：是否使用水平镜像？ 设置为False
'''
return gt_labels
#读取批量图片和标签的函数：get()函数用在训练的时候，每次从gt_labels集合随机读取batch大小的图片以及图片对应的标签。
# 这个函数用到了load_labels(self)函数返回的数据集标签信息，通过数据集标签信息获取数据集中的图片和对应的标注
def get(self):
'''
加载数据集 每次读取batch大小的图片以及图片对应的标签
:return:
images:读取到的图片数据 [45,448,448,3]
labels:对应的图片标签 [45,7,7,25]
'''
#images为（45,448,448,3）的张量 labels为（45,7,7,25）的张量
images = np.zeros((self.batch_size, self.image_size, self.image_size, 3))
labels = np.zeros((self.batch_size, self.cell_size, self.cell_size, 25))
count = 0
# 一次加载batch_size个图片数据
while count < self.batch_size:
imname = self.gt_labels[self.cursor]['imname']#获取图片路径
flipped = self.gt_labels[self.cursor]['flipped']#是否使用水平镜像？
images[count, :, :, :] = self.image_read(imname, flipped)#读取图片数据,调用image_read()函数，读取图片并将图片存入images
labels[count, :, :, :] = self.gt_labels[self.cursor]['label']# #读取图片标签,读取图片标签并存入labels
count += 1
self.cursor += 1
# 如果读取完一轮数据，则当前cursor置为0，当前训练轮数+1
if self.cursor >= len(self.gt_labels):
np.random.shuffle(self.gt_labels)
self.cursor = 0
self.epoch += 1
return images, labels
#图片读取函数，先读取图片，然后缩放，转换为RGB格式，再对数据进行归一化处理。
def image_read(self, imname, flipped=False):
image = cv2.imread(imname)#读取图片
image = cv2.resize(image, (self.image_size, self.image_size))#图片resize至（448,448）
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)#opencv读取的图片格式为BGR格式，将图片转换成RGB格式
# 归一化处理 [-1.0,1.0]
image = (image / 255.0) * 2.0 - 1.0
if flipped:#即水平镜像
#image = image[:, ::-1, :] 表示将图像向右翻转180°
#image = image[::-1,: , :]表示将图像向下翻转180°
image = image[:, ::-1, :]#将图像向右翻转180
return image
(2) timer.py
import time
import datetime
#本py文件只是用来定义显示时间的操作，不重要
class Timer(object):
def __init__(self):
self.init_time = time.time()
self.total_time = 0.
self.calls = 0
self.start_time = 0.
self.diff = 0.
self.average_time = 0.
self.remain_time = 0.
def tic(self):
self.start_time = time.time()
def toc(self, average=True):
self.diff = time.time() - self.start_time
self.total_time += self.diff
self.calls += 1
self.average_time = self.total_time / self.calls
if average:
return self.average_time
else:
return self.diff
def remain(self, iters, max_iters):
if iters == 0:
self.remain_time = 0
else:
self.remain_time = (time.time() - self.init_time) * (max_iters - iters) / iters
return str(datetime.timedelta(seconds=int(self.remain_time)))
2.4 yolo文件夹
(1) config.py
一些参数化设置
import os
#本py文件是一些配置信息
#路径
DATA_PATH = 'data'
PASCAL_PATH = os.path.join(DATA_PATH, 'pascal_voc')#data\pascal_voc
CACHE_PATH = os.path.join(PASCAL_PATH, 'cache')#data\pascal_voc\cache
OUTPUT_DIR = os.path.join(PASCAL_PATH, 'output')#data\pascal_voc\output
WEIGHTS_DIR = os.path.join(PASCAL_PATH, 'weights')#data\pascal_voc\weights
WEIGHTS_FILE = None
WEIGHTS_FILE = os.path.join(DATA_PATH, 'weights', 'YOLO_small.ckpt')
#训练集参数
CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
'train', 'tvmonitor']
FLIPPED = True
#模型参数
IMAGE_SIZE = 448
CELL_SIZE = 7
BOXES_PER_CELL = 2
ALPHA = 0.1
DISP_CONSOLE = False
OBJECT_SCALE = 1.0
NOOBJECT_SCALE = 1.0
CLASS_SCALE = 2.0
COORD_SCALE = 5.0
#训练参数
GPU = ''
LEARNING_RATE = 0.0001
DECAY_STEPS = 30000
DECAY_RATE = 0.1
STAIRCASE = True
BATCH_SIZE = 45
MAX_ITER = 15000
SUMMARY_ITER = 10
SAVE_ITER = 1000
#测试参数
THRESHOLD = 0.2
IOU_THRESHOLD = 0.5
(2) yolo_net.py
用来构建网络的文件，只包含5个函数。
import numpy as np
import tensorflow as tf
import yolo.config as cfg
slim = tf.contrib.slim
#本py文件是YOLOv1网络的搭建
class YOLONet(object):
#初始化函数
def __init__(self, is_training=True):
#数据参数、模型参数、训练参数的设置
self.classes = cfg.CLASSES#数据集的类名称集合
self.num_class = len(self.classes)#类的数量
self.image_size = cfg.IMAGE_SIZE#图片大小448
self.cell_size = cfg.CELL_SIZE# 整张输入图片划分为cell_size * cell_size的网格
self.boxes_per_cell = cfg.BOXES_PER_CELL#每个网格单元的边界框数量
'''网络输出的大小 S*S*(B*5 + C) = 1470'''
self.output_size = (self.cell_size * self.cell_size) *(self.num_class + self.boxes_per_cell * 5)#网络输出的张量size
self.scale = 1.0 * self.image_size / self.cell_size#图片缩放比例
'''# 将网络输出分离为类别和置信度以及边界框，输出维度为7*7*20 + 7*7*2 + 7*7*2*4=1470'''
self.boundary1 = self.cell_size * self.cell_size * self.num_class#7X7X20
self.boundary2 = self.boundary1 +self.cell_size * self.cell_size * self.boxes_per_cell#7X7X20+7X7X2
'''代价函数权重'''
self.object_scale = cfg.OBJECT_SCALE #OBJECT_SCALE = 1.0
self.noobject_scale = cfg.NOOBJECT_SCALE#NOOBJECT_SCALE = 1.0
self.class_scale = cfg.CLASS_SCALE#CLASS_SCALE = 2.0
self.coord_scale = cfg.COORD_SCALE#COORD_SCALE = 5.0
#训练参数
self.learning_rate = cfg.LEARNING_RATE#LEARNING_RATE = 0.0001
self.batch_size = cfg.BATCH_SIZE#BATCH_SIZE = 45
self.alpha = cfg.ALPHA#ALPHA = 0.1
#偏置形状[7,7,2]
self.offset = np.transpose(np.reshape(np.array( [np.arange(self.cell_size)] * self.cell_size * self.boxes_per_cell),
(self.boxes_per_cell, self.cell_size, self.cell_size)), (1, 2, 0))
#输入图片占位符 [NONE,image_size,image_size,3]
#为图片设置placeholder：placeholder()函数是在神经网络构建graph的时候在模型中的占位，此时并没有把要输入的数据传入模型，它只会分配必要的内存。等建立session，在会话中，运行模型的时候通过feed_dict()函数向占位符喂入数据。
self.images = tf.placeholder(tf.float32, [None, self.image_size, self.image_size, 3], name='images')
##调用构建网络函数：获取YOLO网络的输出(不经过激活函数的输出) 形状[None,1470】
#logits获取网络输出：调用构建网络的函数build_network（），输出图片经过网络后的输出，即网络输出
self.logits = self.build_network( self.images, num_outputs=self.output_size, alpha=self.alpha,is_training=is_training)
if is_training:
# 如果是在训练，为标注标签设置placeholder，标签占位符
self.labels = tf.placeholder(tf.float32,[None, self.cell_size, self.cell_size, 5 + self.num_class])
self.loss_layer(self.logits, self.labels)#调用loss_layer（）计算损失，参数（网络预测输出，真实标注），用来计算损失
self.total_loss = tf.losses.get_total_loss()##加入权重正则化之后的损失函数，返回张量代表总损失
tf.summary.scalar('total_loss', self.total_loss)#用来显示标量信息， #将损失以标量形式显示，该变量命名为total_loss
#网络构建函数，网络的输入维度是[None,448,448,3]，输出维度为[None,1470]。
def build_network(self,images,num_outputs,alpha,keep_prob=0.5,is_training=True,scope='yolo'):
#定义变量命名空间
with tf.variable_scope(scope):
# 定义共享参数 使用l2正则化
with slim.arg_scope( [slim.conv2d, slim.fully_connected],activation_fn=leaky_relu(alpha),
weights_regularizer=slim.l2_regularizer(0.0005),
weights_initializer=tf.truncated_normal_initializer(0.0, 0.01) ):
# 图片填充变为454x454x3
net = tf.pad(images, np.array([[0, 0], [3, 3], [3, 3], [0, 0]]),name='pad_1')#对images进行填充，tf. pad详解https://blog.csdn.net/weixin_38517705/article/details/84559348?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522158703330619724839209907%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fall.%2522%257D&request_id=158703330619724839209907&biz_id=0&utm_source=distribute.pc_search_result.none-task-blog-2~all~first_rank_v2~rank_v25-1
#slim.conv2d的解析https://blog.csdn.net/c20081052/article/details/80238090?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522158703383919724835831749%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fall.%2522%257D&request_id=158703383919724835831749&biz_id=0&utm_source=distribute.pc_search_result.none-task-blog-2~all~first_rank_v2~rank_v25-1
net = slim.conv2d(net, 64, 7, 2, padding='VALID', scope='conv_2')# 224x224x64；参数(输入，卷积核个数，卷积核维度，步长)
net = slim.max_pool2d(net, 2, padding='SAME', scope='pool_3')#112x112x64；参数（输入，池化窗口维度）
net = slim.conv2d(net, 192, 3, scope='conv_4')#112x112x192
net = slim.max_pool2d(net, 2, padding='SAME', scope='pool_5')# 56x56x192
net = slim.conv2d(net, 128, 1, scope='conv_6')#56x56x128
net = slim.conv2d(net, 256, 3, scope='conv_7')#56x56x256
net = slim.conv2d(net, 256, 1, scope='conv_8')# 56x56x256
net = slim.conv2d(net, 512, 3, scope='conv_9')#56x56x512
net = slim.max_pool2d(net, 2, padding='SAME', scope='pool_10')#28x28x512
net = slim.conv2d(net, 256, 1, scope='conv_11')#28x28x256
net = slim.conv2d(net, 512, 3, scope='conv_12')# 28x28x512
net = slim.conv2d(net, 256, 1, scope='conv_13')#28x28x256
net = slim.conv2d(net, 512, 3, scope='conv_14')#28x28x512
net = slim.conv2d(net, 256, 1, scope='conv_15')#28x28x256
net = slim.conv2d(net, 512, 3, scope='conv_16')#28x28x512
net = slim.conv2d(net, 256, 1, scope='conv_17')#28x28x256
net = slim.conv2d(net, 512, 3, scope='conv_18')#28x28x512
net = slim.conv2d(net, 512, 1, scope='conv_19')#28x28x512
net = slim.conv2d(net, 1024, 3, scope='conv_20')#28x28x1024
net = slim.max_pool2d(net, 2, padding='SAME', scope='pool_21')#14x14x1024
net = slim.conv2d(net, 512, 1, scope='conv_22')#14x14x512
net = slim.conv2d(net, 1024, 3, scope='conv_23')#14x14x1024
net = slim.conv2d(net, 512, 1, scope='conv_24')#14x14x512
net = slim.conv2d(net, 1024, 3, scope='conv_25')#14x14x1024
net = slim.conv2d(net, 1024, 3, scope='conv_26')#14x14x1024
net = tf.pad( net, np.array([[0, 0], [1, 1], [1, 1], [0, 0]]),name='pad_27')#16x16x2014；pad
net = slim.conv2d(net, 1024, 3, 2, padding='VALID', scope='conv_28')#7x7x1024
net = slim.conv2d(net, 1024, 3, scope='conv_29')# 7x7x1024
net = slim.conv2d(net, 1024, 3, scope='conv_30')#7x7x1024
net = tf.transpose(net, [0, 3, 1, 2], name='trans_31')#[None,1024,7,7]，改变不同维度的顺序，从而改变输入张量的shape
net = slim.flatten(net, scope='flat_32')#展开 50176，将输入扁平化但保留batch_size
net = slim.fully_connected(net, 512, scope='fc_33')# 512
net = slim.fully_connected(net, 4096, scope='fc_34')# 4096
net = slim.dropout(net, keep_prob=keep_prob, is_training=is_training,scope='dropout_35')#4096
net = slim.fully_connected(net, num_outputs, activation_fn=None, scope='fc_36')#1470；全连接层输出长向量，长向量的长度(self.cell_size * self.cell_size) *(self.num_class + self.boxes_per_cell * 5)
return net#返回网络最后一层，激活函数处理之前的值 形状[None,1470]
#IOU计算函数：计算两个 bounding box 之间的 IoU。输入是两个 5 维的bounding box,输出的两个 bounding Box 的IoU
def calc_iou(self, boxes1, boxes2, scope='iou'):
"""calculate ious
Args:
boxes1: 5维 tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL, 4] ====> (x_center, y_center, w, h)
boxes2: 5维 tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL, 4] ===> (x_center, y_center, w, h)
参数x_center, y_center, w, h都是归一到[0,1]之间的，分别表示预测边界框的中心相对整张图片的坐标，宽和高
Return:
iou: 4维 tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL]
"""
with tf.variable_scope(scope):
'''参数x_center, y_center, w, h都是归一到[0,1]之间的，分别表示预测边界框的中心相对整张图片的坐标，宽和高'''
#将边界框信息由(x_center, y_center, w, h) 转变为 (x1, y1, x2, y2)
boxes1_t = tf.stack([boxes1[..., 0] - boxes1[..., 2] / 2.0,#左上角x
boxes1[..., 1] - boxes1[..., 3] / 2.0,#左上角y
boxes1[..., 0] + boxes1[..., 2] / 2.0,#右下角x
boxes1[..., 1] + boxes1[..., 3] / 2.0],#右下角y
axis=-1)
boxes2_t = tf.stack([boxes2[..., 0] - boxes2[..., 2] / 2.0,
boxes2[..., 1] - boxes2[..., 3] / 2.0,
boxes2[..., 0] + boxes2[..., 2] / 2.0,
boxes2[..., 1] + boxes2[..., 3] / 2.0],
axis=-1)
# 两个框相交部分矩形的左上角和右下角点
lu = tf.maximum(boxes1_t[..., :2], boxes2_t[..., :2])# #两个框相交的矩形的左上角(x1,y1)
rd = tf.minimum(boxes1_t[..., 2:], boxes2_t[..., 2:])# #两个框相交的矩形的右下角(x2,y2)
# 求相交矩形的面积
intersection = tf.maximum(0.0, rd - lu)#相交矩形的长和宽：加一个tf.maximum是因为删除那些不合理的框，比如两个框没交集，左上角坐标比右下角还大
inter_square = intersection[..., 0] * intersection[..., 1]#求面积了，就是长乘以宽。
# 求两个框的面积
square1 = boxes1[..., 2] * boxes1[..., 3]
square2 = boxes2[..., 2] * boxes2[..., 3]
#两个框的交面积，因为如果两个框的面积相加，那就会重复了相交的部分，
#所以减去相交的部分，外面有个tf.maximum这个就是保证相交面积不为0,因为后面要做分母
union_square = tf.maximum(square1 + square2 - inter_square, 1e-10)
#tf.clip_by_value( t,clip_value_min, clip_value_max,name=None)
#将一个张量的值限制在给定的最小值和最大值之间。对于给定的张量t，返回的张量与之有着相同的类型和相同的大小，
#只是它的值在clip_value_min和clip_value_max之间。任何比clip_value_min小的数设置成clip_value_min，
#任何比clip_value_max大的数被设置成clip_value_max.
return tf.clip_by_value(inter_square / union_square, 0.0, 1.0)#如果你的交并比大于1,那么就让它等于1,如果小于0,那么就
#让他变为0,因为交并比在0-1之间
#损失函数定义函数
def loss_layer(self, predicts, labels, scope='loss_layer'):
# predicts：Yolo网络的输出形状[None, 1470] 1470=7*7*（20+5*2）
with tf.variable_scope(scope):
#预测信息[None,1470]
#获取预测的类别[0:7*7*20]：形状变为[45,7,7,20]
predict_classes = tf.reshape(
predicts[:, :self.boundary1],#self.boundary1 = self.cell_size * self.cell_size * self.num_class#7X7X20
[self.batch_size, self.cell_size, self.cell_size, self.num_class])
#获取预测的置信度[7*7*20 : 7*7*20 + 7*7*2]：形状变为[45,7,7,2]
predict_scales = tf.reshape(
predicts[:, self.boundary1:self.boundary2],# self.boundary2 = self.boundary1 +self.cell_size * self.cell_size * self.boxes_per_cell#7X7X20+7X7
[self.batch_size, self.cell_size, self.cell_size, self.boxes_per_cell])
#预测的边界框[7*7*20 + 7*7*2：1470];形状变为[45,7,7,2,4]目标中心是相对于当前格子的，宽度和高度的开根号是相对当前整张图像的(归一化的)
predict_boxes = tf.reshape(
predicts[:, self.boundary2:],
[self.batch_size, self.cell_size, self.cell_size, self.boxes_per_cell, 4])
# 真值信息\标签信息：[None,7,7,25]
#置信度【0；1】，表示这个格子是否有目标 形状[45,7,7,1]
response = tf.reshape(
labels[..., 0],
[self.batch_size, self.cell_size, self.cell_size, 1])
#目标边界框【1:5】:目标中心，宽度和高度
boxes = tf.reshape(#形状[45,7,7,1，4]
labels[..., 1:5],
[self.batch_size, self.cell_size, self.cell_size, 1, 4])
boxes = tf.tile(boxes, [1, 1, 1, self.boxes_per_cell, 1]) / self.image_size#标签的边界框 归一化后 张量沿着axis=3重复两边，扩充后[45,7,7,2,4]
# 目标类别【5；25】
classes = labels[..., 5:] #目标类别
'''
offset:init函数中初始化为[7,7,2]的矩阵，每一行都是[7,2]的矩阵，值为[[0,0],[1,1],[2,2],[3,3],[4,4],[5,5],[6,6]]
这个变量是为了将每个cell的坐标对齐，后一个框比前一个框要多加1
比如我们预测了cell_size的每个中心点坐标，那么我们这个中心点落在第几个cell_size
就对应坐标要加几，这个用法比较巧妙，构造了这样一个数组，让他们对应位置相加
'''
#offset reshape为[1,7,7,2] 如果忽略axis=0，则每一行都是 [[0,0],[1,1],[2,2],[3,3],[4,4],[5,5],[6,6]]
offset = tf.reshape(
tf.constant(self.offset, dtype=tf.float32),
[1, self.cell_size, self.cell_size, self.boxes_per_cell])
##shape为[45,7,7,2]
offset = tf.tile(offset, [self.batch_size, 1, 1, 1])
#shape为[45,7,7,2] 如果忽略axis=0 第i行为[[i,i],[i,i],[i,i],[i,i],[i,i],[i,i],[i,i]]
offset_tran = tf.transpose(offset, (0, 2, 1, 3))
'''offset变量用于把预测边界框predict_boxes中的坐标中心(x,y)由相对当前格子转换为相对当前整个图片'''
#计算每个格子中的预测边界框坐标(x,y)相对于整个图片的位置 而不是相对当前格子
# 假设当前格子为(3,3)，当前格子的预测边界框为(x0,y0)，则计算坐标(x,y) = ((x0,y0)+(3,3))/7
predict_boxes_tran = tf.stack(
[(predict_boxes[..., 0] + offset) / self.cell_size,#x
(predict_boxes[..., 1] + offset_tran) / self.cell_size,#y
tf.square(predict_boxes[..., 2]),#w
tf.square(predict_boxes[..., 3])], axis=-1)#h
# offset变量用于把真值框坐标相对整个图像->相对当前格子
boxes_tran = tf.stack(
[boxes[..., 0] * self.cell_size - offset,
boxes[..., 1] * self.cell_size - offset_tran,
tf.sqrt(boxes[..., 2]),
tf.sqrt(boxes[..., 3])], axis=-1)
# 计算每个格子预测边界框与真实边界框之间的IOU [45,7,7,2]
iou_predict_truth = self.calc_iou(predict_boxes_tran, boxes)#calc_iou（）调用，计算预测与真值的iou值
# calculate I tensor [BATCH_SIZE, CELL_SIZE, CELL_SIZE, BOXES_PER_CELL]
#YOLO每个网格单元预测两个边界框，一个类别。在训练时，每个目标只需要一个预测器来负责，哪个预测器与真实值之间具有当前最高的IOU
# 哪个预测器来预测目标。
# 所以object_mask就表示每个格子中的哪个边界框负责该格子中目标预测？哪个边界框取值为1，哪个边界框就负责目标预测
# 当格子中的确有目标时，取值为[1,1]，[1,0],[0,1]
# 比如某一个格子的值为[1,0]，表示第一个边界框负责该格子目标的预测 [0,1]：表示第二个边界框负责该格子目标的预测
# 当格子没有目标时，取值为[0,0]
object_mask = tf.reduce_max(iou_predict_truth, 3, keep_dims=True)#reduce_max函数是求按axis方向的最值
object_mask = tf.cast((iou_predict_truth >= object_mask), tf.float32) * response#cast（）将数据格式转化成dtype
# calculate no_I tensor [CELL_SIZE, CELL_SIZE, BOXES_PER_CELL]
# noobject_mask就表示每个边界框不负责该目标的置信度，
# 当格子没有目标时，取值为[1,1]
# 使用tf.onr_like，使得全部为1,再减去有目标的，也就是有目标的对应坐标为1,这样一减，就变为没有的了。[45,7,7,2]
noobject_mask = tf.ones_like(object_mask, dtype=tf.float32) - object_mask#创建一个和object_mask维度一样数值为1的张量
# 分类损失，如果目标出现在网格中 response为1，否则response为0
class_delta = response * (predict_classes - classes)
class_loss = tf.reduce_mean(tf.reduce_sum(tf.square(class_delta), axis=[1, 2, 3]), name='class_loss') * self.class_scale
#置信度损失
# 有目标物体存在的置信度预测损失：当格子中有目标时，负责该目标预测的边界框的置信度越越接近预测的边界框与实际边界框之间的IOU，代价值越小
object_delta = object_mask * (predict_scales - iou_predict_truth)
object_loss = tf.reduce_mean(tf.reduce_sum(tf.square(object_delta), axis=[1, 2, 3]),name='object_loss') * self.object_scale
# 没有目标物体存在的置信度的损失：(此时iou_predict_truth为0)当格子中没有目标时，预测的两个边界框的置信度越接近0，代价值越小
noobject_delta = noobject_mask * predict_scales
noobject_loss = tf.reduce_mean(tf.reduce_sum(tf.square(noobject_delta), axis=[1, 2, 3]), name='noobject_loss') * self.noobject_scale
# 边界框坐标损失 ：shape 为 [batch_size, 7, 7, 2, 1]，当格子中有目标时，预测的边界框越接近实际边界框，代价值越小
coord_mask = tf.expand_dims(object_mask, 4)
boxes_delta = coord_mask * (predict_boxes - boxes_tran)
coord_loss = tf.reduce_mean( tf.reduce_sum(tf.square(boxes_delta), axis=[1, 2, 3, 4]),name='coord_loss') * self.coord_scale
#汇总所有损失
tf.losses.add_loss(class_loss)
tf.losses.add_loss(object_loss)
tf.losses.add_loss(noobject_loss)
tf.losses.add_loss(coord_loss)
#损失添加到日志记录
tf.summary.scalar('class_loss', class_loss)
tf.summary.scalar('object_loss', object_loss)
tf.summary.scalar('noobject_loss', noobject_loss)
tf.summary.scalar('coord_loss', coord_loss)
tf.summary.histogram('boxes_delta_x', boxes_delta[..., 0])
tf.summary.histogram('boxes_delta_y', boxes_delta[..., 1])
tf.summary.histogram('boxes_delta_w', boxes_delta[..., 2])
tf.summary.histogram('boxes_delta_h', boxes_delta[..., 3])
tf.summary.histogram('iou', iou_predict_truth)
#激活函数定义函数
def leaky_relu(alpha):
def op(inputs):
return tf.nn.leaky_relu(inputs, alpha=alpha, name='leaky_relu')
return op
2.5 test.py
import os
import cv2
from utils.timer import Timer
import argparse
import numpy as np
import tensorflow as tf
import yolo.config as cfg
from yolo.yolo_net import YOLONet
class Detector(object):
#初始化函数
def __init__(self, net, weight_file):
self.net = net#网络模型
self.weights_file = weight_file#网络模型参数
self.classes = cfg.CLASSES#数据集类列表
self.num_class = len(self.classes)#数据集类的数量20
self.image_size = cfg.IMAGE_SIZE#图片尺寸
self.cell_size = cfg.CELL_SIZE#图片划分为cell_size*cell_size的网格
self.boxes_per_cell = cfg.BOXES_PER_CELL#每个网格单元的边界框数量
self.threshold = cfg.THRESHOLD#阈值
self.iou_threshold = cfg.IOU_THRESHOLD#iou阈值
# 将网络输出分离为类别和置信度以及边界框的大小，输出维度为7*7*20 + 7*7*2 + 7*7*2*4=1470‘‘‘
self.boundary1 = self.cell_size * self.cell_size * self.num_class#7*7*20
self.boundary2 = self.boundary1 +self.cell_size * self.cell_size * self.boxes_per_cell#7*7*20+7*7*2
# 运行图之前，初始化变量
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())#初始化模型的参数
print('Restoring weights from: ' + self.weights_file)
#https://www.cnblogs.com/bevishe/p/10359993.html tf.train.Saver()的使用
self.saver = tf.train.Saver()#实例化一个Saver对象 saver = tf.train.Saver() 。在训练过程中，定期调用saver.save方法，像文件夹中写入包含当前模型中所有可训练变量的checkpoint文件 saver.save(sess,FLAGG.train_dir,global_step=step)
# 恢复模型
self.saver.restore(self.sess, self.weights_file)#使用saver.restore()方法，重载模型的参数，继续训练或者用于测试数据 saver.restore(sess,FLAGG.train_dir)
#检测结果显示函数:在原图上绘制边界框，以及附加信息
def draw_result(self, img, result):
for i in range(len(result)):#遍历每一个边界框
x = int(result[i][1])#x_center
y = int(result[i][2])#y_center
w = int(result[i][3] / 2) #w/2
h = int(result[i][4] / 2)#h/2
# 绘制矩形框(目标边界框) 矩形左上角，矩形右下角
cv2.rectangle(img, (x - w, y - h), (x + w, y + h), (0, 255, 0), 2)
# 绘制矩形框，用于存放类别名称，使用灰度填充
cv2.rectangle(img, (x - w, y - h - 20),(x + w, y - h), (125, 125, 125), -1)
# 线型
lineType = cv2.LINE_AA if cv2.__version__ > '3' else cv2.CV_AA
#绘制文本信息 写上类别名和置信度
cv2.putText(
img, result[i][0] + ' : %.2f' % result[i][5],
(x - w + 5, y - h - 7), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(0, 0, 0), 1, lineType)
# 图片目标检测:返回result包含类别名,x_center,y_center,w,h,置信度
def detect(self, img):#图片张量
'''
:param img:原始图片数据
:return:返回检测到的边界框，list类型 每一个元素对应一个目标框
包含{类别名,x_center,y_center,w,h,置信度}
'''
img_h, img_w, _ = img.shape#图片高和宽
inputs = cv2.resize(img, (self.image_size, self.image_size))#图片resize至规定的大小
inputs = cv2.cvtColor(inputs, cv2.COLOR_BGR2RGB).astype(np.float32)#opencv的是BGR，将图片由BGR变为RGB形式
inputs = (inputs / 255.0) * 2.0 - 1.0#图片规范化
inputs = np.reshape(inputs, (1, self.image_size, self.image_size, 3))#reshape操作
# 获取网络输出第一项(即第一张图片) [1,1470]
result = self.detect_from_cvmat(inputs)[0]
# 对检测的图片的边界框进行缩放处理，一张图片可以有多个边界框
for i in range(len(result)):
# x_center, y_center, w, h都是真实值，分别表示预测边界框的中心坐标，宽和高，都是浮点型
result[i][1] *= (1.0 * img_w / self.image_size)
result[i][2] *= (1.0 * img_h / self.image_size)
result[i][3] *= (1.0 * img_w / self.image_size)
result[i][4] *= (1.0 * img_h / self.image_size)
return result
#运行yolo网络，开始检测
def detect_from_cvmat(self, inputs):
'''
:param inputs:输入数据 [None,448,448,3]
:return:返回目标检测的结果，每一个元素对应一个测试图片，每个元素包含着若干个边界框
'''
''' 返回网络最后一层，激活函数处理之前的值 形状[None,1470]'''
net_output = self.sess.run(self.net.logits,feed_dict={self.net.images: inputs})
results = []
# 对网络输出每一行数据进行处理
for i in range(net_output.shape[0]):
results.append(self.interpret_output(net_output[i]))
return results
#yolo网络输出的结果预处理函数：提取出有目标的边界框，方便后续的处理。
def interpret_output(self, output):
'''
:param output:yolo网络输出的每一行数据 大小为[1470,]
:return:result:yolo网络目标检测到的边界框，list类型 每一个元素对应一个目标框
包含{类别名,x_center,y_center,w,h,置信度}
'''
# [7,7,2,20]
probs = np.zeros((self.cell_size, self.cell_size, self.boxes_per_cell, self.num_class))
# 类别概率 [7,7,20]
class_probs = np.reshape(output[0:self.boundary1],(self.cell_size, self.cell_size, self.num_class))
# 置信度 [7,7,2]
scales = np.reshape(
output[self.boundary1:self.boundary2],
(self.cell_size, self.cell_size, self.boxes_per_cell))
# 边界框 [7,7,2,4]
boxes = np.reshape(
output[self.boundary2:],
(self.cell_size, self.cell_size, self.boxes_per_cell, 4))
offset = np.array(
[np.arange(self.cell_size)] * self.cell_size * self.boxes_per_cell)
# [7,7,2] 每一行都是 [[0,0],[1,1],[2,2],[3,3],[4,4],[5,5],[6,6]]
offset = np.transpose(np.reshape( offset, [self.boxes_per_cell, self.cell_size, self.cell_size]),(1, 2, 0))
# 目标中心相对于整个图片
boxes[:, :, :, 0] += offset
boxes[:, :, :, 1] += np.transpose(offset, (1, 0, 2))
boxes[:, :, :, :2] = 1.0 * boxes[:, :, :, 0:2] / self.cell_size
# 宽度、高度相对整个图片的
boxes[:, :, :, 2:] = np.square(boxes[:, :, :, 2:])
# 转换成实际的编辑框(没有归一化的)
boxes *= self.image_size
# 遍历每一个边界框的置信度
for i in range(self.boxes_per_cell):
for j in range(self.num_class):#遍历每一个类别
'''在测试时，条件类概率和单个盒子的置信度预测相乘，
这些分数编码了j类出现在框i中的概率以及预测框拟合目标的程度。'''
probs[:, :, i, j] = np.multiply(class_probs[:, :, j], scales[:, :, i])
# [7,7,2,20] 如果第i个边界框检测到类别j 则[;,;,i,j]=1
filter_mat_probs = np.array(probs >= self.threshold, dtype='bool')
# 返回filter_mat_probs非0值的索引 返回4个List，每个list长度为n 即检测到的边界框的个数
filter_mat_boxes = np.nonzero(filter_mat_probs)
# 获取检测到目标的边界框 [n,4] n表示边界框的个数
boxes_filtered = boxes[filter_mat_boxes[0],
filter_mat_boxes[1], filter_mat_boxes[2]]
# 获取检测到目标的边界框的置信度 (n,)
probs_filtered = probs[filter_mat_probs]
# 获取检测到目标的边界框对应的目标类别 (n,)
classes_num_filtered = np.argmax(
filter_mat_probs, axis=3)[
filter_mat_boxes[0], filter_mat_boxes[1], filter_mat_boxes[2]]
# 按置信度倒序排序，返回对应的索引
argsort = np.array(np.argsort(probs_filtered))[::-1]
boxes_filtered = boxes_filtered[argsort]
probs_filtered = probs_filtered[argsort]
classes_num_filtered = classes_num_filtered[argsort]
for i in range(len(boxes_filtered)):
if probs_filtered[i] == 0:
continue
for j in range(i + 1, len(boxes_filtered)):
'''非极大值抑制'''
# 计算n各边界框，两两之间的IoU是否大于阈值，非极大值抑制
if self.iou(boxes_filtered[i], boxes_filtered[j]) > self.iou_threshold:
probs_filtered[j] = 0.0
# 非极大值抑制后的输出
filter_iou = np.array(probs_filtered > 0.0, dtype='bool')
boxes_filtered = boxes_filtered[filter_iou]
probs_filtered = probs_filtered[filter_iou]
classes_num_filtered = classes_num_filtered[filter_iou]
# 遍历每一个边界框，存入结果列表
result = []
for i in range(len(boxes_filtered)):
result.append(
[self.classes[classes_num_filtered[i]],
boxes_filtered[i][0],
boxes_filtered[i][1],
boxes_filtered[i][2],
boxes_filtered[i][3],
probs_filtered[i]])
return result
#计算两个边界框的IoU函数
def iou(self, box1, box2):
'''
计算两个边界框的IoU
args:
box1：边界框1 [4,] 真实值
box2：边界框2 [4,] 真实值
'''
tb = min(box1[0] + 0.5 * box1[2], box2[0] + 0.5 * box2[2]) - max(box1[0] - 0.5 * box1[2], box2[0] - 0.5 * box2[2])
lr = min(box1[1] + 0.5 * box1[3], box2[1] + 0.5 * box2[3]) - max(box1[1] - 0.5 * box1[3], box2[1] - 0.5 * box2[3])
inter = 0 if tb < 0 or lr < 0 else tb * lr
return inter / (box1[2] * box1[3] + box2[2] * box2[3] - inter)
#摄像头实现实时目标检测，并展示结果的函数：
def camera_detector(self, cap, wait=10):
detect_timer = Timer()#测试时间
ret, _ = cap.read()#读取一帧
while ret:
# 读取一帧
ret, frame = cap.read()
# 测试时间
detect_timer.tic()
# 执行测试
result = self.detect(frame)
# 测试结束时间
detect_timer.toc()
print('Average detecting time: {:.3f}s'.format( detect_timer.average_time))
# 绘制边界框，以及添加附加信息
self.draw_result(frame, result)
cv2.imshow('Camera', frame)
# 显示
cv2.waitKey(wait)
ret, frame = cap.read()
#图片进行目标检测函数，并显示检测结果的函数：调用detect(self, img)
def image_detector(self, imname, wait=0):#传入图片名
detect_timer = Timer()#实例化时间类
image = cv2.imread(imname)#根据路径读取图片
detect_timer.tic()
result = self.detect(image)#调用detect函数
detect_timer.toc()
print('Average detecting time: {:.3f}s'.format(detect_timer.average_time))
self.draw_result(image, result)
cv2.imshow('Image', image)
cv2.waitKey(wait)
def main():
parser = argparse.ArgumentParser()##创建一个解析器对象，并告诉它将会有些什么参数。当程序运行时，该解析器就可以用于处理命令行参数。
parser.add_argument('--weights', default="YOLO_small.ckpt", type=str)
parser.add_argument('--weight_dir', default='weights', type=str)
parser.add_argument('--data_dir', default="data", type=str)
parser.add_argument('--gpu', default='', type=str)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
yolo = YOLONet(False) #创建YOLO网络对象
# 加载检查点文件
weight_file = os.path.join(args.data_dir, args.weight_dir, args.weights)
# 创建测试对象
detector = Detector(yolo, weight_file)
# detect from camera
# cap = cv2.VideoCapture(-1)
# detector.camera_detector(cap)
# detect from image file
imname = 'test/person.jpg'
detector.image_detector(imname)#调用图片检测函数
if __name__ == '__main__':
main()
2.6 train.py
import os
import argparse
import datetime
import tensorflow as tf
import yolo.config as cfg
from yolo.yolo_net import YOLONet
from utils.timer import Timer
from utils.pascal_voc import pascal_voc
# from utils import timer
# from timer import Timer
# from utils import pascal_voc
# from pascal_voc import pascal_voc
slim = tf.contrib.slim
#本py文件用于训练YOLO网络模型
class Solver(object):
#初始化函数
def __init__(self, net, data):
'''
:param net: YOLONet对象
:param data: pascal_voc对象
'''
self.net = net #yolo网络模型
self.data = data#voc数据集
self.weights_file = cfg.WEIGHTS_FILE#检查点文件路径，权重文件夹
self.max_iter = cfg.MAX_ITER#训练最大迭代次数
self.initial_learning_rate = cfg.LEARNING_RATE#初始学习率
self.decay_steps = cfg.DECAY_STEPS#退化学习率衰减步数
self.decay_rate = cfg.DECAY_RATE #衰减率
self.staircase = cfg.STAIRCASE
self.summary_iter = cfg.SUMMARY_ITER#日志文件保存间隔步
self.save_iter = cfg.SAVE_ITER#模型保存间隔步
# 输出文件夹路径
self.output_dir = os.path.join(cfg.OUTPUT_DIR, datetime.datetime.now().strftime('%Y_%m_%d_%H_%M'))
if not os.path.exists(self.output_dir):
os.makedirs(self.output_dir)
self.save_cfg() #保存配置信息
''' 指定保存的张量 这里指定所有变量'''
self.variable_to_restore = tf.global_variables()
self.saver = tf.train.Saver(self.variable_to_restore, max_to_keep=None)
# 指定保存的模型名称
self.ckpt_file = os.path.join(self.output_dir, 'yolo')
# 合并所有的summary
self.summary_op = tf.summary.merge_all()
# 创建writer，指定日志文件路径，用于写日志文件
self.writer = tf.summary.FileWriter(self.output_dir, flush_secs=60)
# 创建变量，保存当前迭代次数
self.global_step = tf.train.create_global_step()
# 退化学习率
self.learning_rate = tf.train.exponential_decay(
self.initial_learning_rate, self.global_step, self.decay_steps,
self.decay_rate, self.staircase, name='learning_rate')
#创建优化器
self.optimizer = tf.train.GradientDescentOptimizer(learning_rate=self.learning_rate)
# 创建train_op
'''train_op 的作用：[1] - 计算 loss.[2] - 应用梯度，更新权重.[3] - 返回 loss 值.'''
self.train_op = slim.learning.create_train_op(self.net.total_loss, self.optimizer, global_step=self.global_step)
# 设置GPU使用资源
gpu_options = tf.GPUOptions()
# 按需分配GPU使用的资源
config = tf.ConfigProto(gpu_options=gpu_options)
self.sess = tf.Session(config=config)
# 运行图之前，初始化变量
self.sess.run(tf.global_variables_initializer())
# 恢复模型
if self.weights_file is not None:
print('Restoring weights from: ' + self.weights_file)
self.saver.restore(self.sess, self.weights_file)
# 将图写入日志文件
self.writer.add_graph(self.sess.graph)
#训练函数
def train(self):
train_timer = Timer() #训练时间
load_timer = Timer()#数据集加载时间
# 开始迭代
for step in range(1, self.max_iter + 1):
load_timer.tic() #计算每次迭代加载数据的起始时间
images, labels = self.data.get()# #加载数据集 每次读取batch大小的图片以及图片对应的标签
load_timer.toc() #计算这次迭代加载数据集所使用的时间
feed_dict = {self.net.images: images,
self.net.labels: labels}
# 迭代summary_iter次，保存一次日志文件，迭代summary_iter*10次，输出一次的迭代信息
if step % self.summary_iter == 0:
if step % (self.summary_iter * 10) == 0:
train_timer.tic()#计算每次迭代训练的起始时间
# 开始迭代训练，每一次迭代后global_step自加1
summary_str, loss, _ = self.sess.run(
[self.summary_op, self.net.total_loss, self.train_op],
feed_dict=feed_dict)
train_timer.toc()
# 输出信息
log_str = '''{} Epoch: {}, Step: {}, Learning rate: {},'''
''' Loss: {:5.3f}\nSpeed: {:.3f}s/iter,'''
'''' Load: {:.3f}s/iter, Remain: {}'''.format(
datetime.datetime.now().strftime('%m-%d %H:%M:%S'),
self.data.epoch,
int(step),
round(self.learning_rate.eval(session=self.sess), 6),
loss,
train_timer.average_time,
load_timer.average_time,
train_timer.remain(step, self.max_iter))
print(log_str)
else:
train_timer.tic()#计算每次迭代训练的起始时间
# 开始迭代训练，每一次迭代后global_step自加1,train_op 的作用：[1] - 计算 loss.[2] - 应用梯度，更新权重.[3] - 返回 loss 值.
summary_str, _ = self.sess.run([self.summary_op, self.train_op],feed_dict=feed_dict)
train_timer.toc()#计算这次迭代训练所使用的时间
self.writer.add_summary(summary_str, step) #将summary写入文件
else:
train_timer.tic() #计算每次迭代训练的起始时间
# 开始迭代训练，每一次迭代后global_step自加1
self.sess.run(self.train_op, feed_dict=feed_dict)
train_timer.toc()
# 每迭代save_iter次，保存一次模型
if step % self.save_iter == 0:
print('{} Saving checkpoint file to: {}'.format(
datetime.datetime.now().strftime('%m-%d %H:%M:%S'),
self.output_dir))
self.saver.save(
self.sess, self.ckpt_file, global_step=self.global_step)
#保存配置参数函数
def save_cfg(self):
with open(os.path.join(self.output_dir, 'config.txt'), 'w') as f:
cfg_dict = cfg.__dict__
for key in sorted(cfg_dict.keys()):
if key[0].isupper():
cfg_str = '{}: {}\n'.format(key, cfg_dict[key])
f.write(cfg_str)
#设定数据集路径，以及检查点文件路径
def update_config_paths(data_dir, weights_file):
'''
:param data_dir:数据文件夹 数据集放在pascal_voc目录下
:param weights_file:检查点文件名 该文件放在数据集目录下的weights文件夹下
:return:
'''
cfg.DATA_PATH = data_dir#数据所在文件夹
cfg.PASCAL_PATH = os.path.join(data_dir, 'pascal_voc')#VOC数据所在文件夹
cfg.CACHE_PATH = os.path.join(cfg.PASCAL_PATH, 'cache') #保存生成的数据集标签缓冲文件所在的文件夹
cfg.OUTPUT_DIR = os.path.join(cfg.PASCAL_PATH, 'output') #保存生成的网络模型和日志文件所在的文件夹
cfg.WEIGHTS_DIR = os.path.join(cfg.PASCAL_PATH, 'weights') #检查点文件所在的目录
cfg.WEIGHTS_FILE = os.path.join(cfg.WEIGHTS_DIR, weights_file)
def main():
#创建一个解析器对象，并告诉它将会有些什么参数。那么当你的程序运行时，
# 该解析器就可以用于处理命令行参数。
parser = argparse.ArgumentParser()
parser.add_argument('--weights', default="YOLO_small.ckpt", type=str)
parser.add_argument('--data_dir', default="data", type=str)
parser.add_argument('--threshold', default=0.2, type=float)
parser.add_argument('--iou_threshold', default=0.5, type=float)
parser.add_argument('--gpu', default='', type=str)
# 定义了所有参数之后，你就可以给 parse_args() 传递一组参数字符串来解析命令行。
# 默认情况下，参数是从 sys.argv[1:] 中获取
# parse_args() 的返回值是一个命名空间，包含传递给命令的参数。该对象将参数保存其属性
args = parser.parse_args()
if args.gpu is not None:
cfg.GPU = args.gpu
#设定数据集路径，以及检查点文件路径
if args.data_dir != cfg.DATA_PATH:
update_config_paths(args.data_dir, args.weights)
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU
yolo = YOLONet() #创建YOLO网络对象
pascal = pascal_voc('train')#数据集对象
solver = Solver(yolo, pascal) #求解器对象
print('Start training ...')
solver.train()#开始训练
print('Done training.')
if __name__ == '__main__':
# python train.py --weights YOLO_small.ckpt --gpu 0
main()
版权声明：本文为qq_36926037原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。