(本文进行训练和预测代码的记录)
直接上代码:
因为我们最后输出的是N 24 H W的格式,其中24 = 3×8
8表示:置信度,回归,分类
置信度采取:二分类交叉熵 回归采取:BCE
分类采取:多分类交叉熵
3表示:3个框。
from torch import nn,optim
import torch
from net import *
from torch.utils.data import DataLoader
from Dataset import *
import os
from torch.utils.tensorboard import SummaryWriter
def loss_fun(output,target,c):
output = output.permute(0,2,3,1)
"""
将原来的N 24 13 13转换成N 13 13 24(24可拆)
转换后可用前面的数据查询后面的数据
"""
output = output.reshape(output.size(0),output.size(1),output.size(2),3,-1)
#改变格式为N H W 3个框 最后一位用“-1”代替(自动计算)
#计算置信度
mask_obj = target[..., 0] > 0
mask_no_obj = target[..., 0] == 0
#target[..., 0]取最后一个维度,判断是否有目标
loss_p_fun = nn.BCELoss()
# 二值交叉熵
loss_p = loss_p_fun(torch.sigmoid(output[...,0]),target[...,0])
#激活函数,归一化到0-1之间,方便计算也防止精度爆炸
loss_box_fun = nn.MSELoss()
# 均方差
loss_box = loss_box_fun(output[mask_obj][...,1:5],target[mask_obj][...,1:5])
# target[mask_obj][...,1:5]取最后的3×15中15的1-4
loss_segment_fun = nn.CrossEntropyLoss()
# 交叉熵
loss_segment = loss_segment_fun(output[mask_obj][...,5:],torch.argmax(target[mask_obj][...,5:],dim = 1,keepdim = True).squeeze(dim = 1))
'''
因为多分类交叉熵自带Onehot编码,所以要把已经处理好的独热编码数据恢复成原状
'''
loss = c*loss_p+(1-c)*0.5*loss_box+(1-c)*0.5*loss_segment
"""
因为正负样本的比例相差较大,所以要给相应的偏重进行调整
所有的损失相加,对全局调整
"""
return loss
if __name__ == '__main__':
summary_write = SummaryWriter('logs')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#判断GPU是否可用,不可用则用CPU
dataset = YoloDataset()
data_loader = DataLoader(dataset,batch_size=2,shuffle=True)
# 调用数据加载器,不调用的话数据无法加载进来
weight_path = 'params/net.pt'
# 保存权重
net = Yolo_v3_Net().to(device)
# 加载网络进GPU
if os.path.exists(weight_path):
net.load_state_dict(torch.load(weight_path))
#如果权重存在,则进行权重的调用
opt = optim.Adam(net.parameters())
#优化器(学习率不需要我们给)
epoch = 0
index = 0
while True:
for target_13,target_26,target_52,img_data in data_loader:
target_13,target_26,target_52,img_data = target_13.to(device),target_26.to(device),target_52.to(device),img_data.to(device)
#将数据存储到GPU
output_13,output_26,output_52 = net(img_data)
loss_13 = loss_fun(output_13.float(), target_13.float(), 0.7)
loss_26 = loss_fun(output_26.float(), target_26.float(), 0.7)
loss_52 = loss_fun(output_52.float(), target_52.float(), 0.7)
loss = loss_13+loss_26+loss_52
opt.zero_grad()
#清空梯度
loss.backward()
#计算梯度
opt.step()
print(f'loss{epoch}=={index}',loss.item())
index+=1
summary_write.add_scalar('train_loss',loss,index)
# 写入到tensorboard
torch.save(net.state_dict(),'params/net.pt')
# 保存模型
print('SAVE SUSSESSFULLY!')
epoch+=1
tensorboard效果:
首先,输入一张图片以后,需要给一个置信度约束那些置信度很低的值,比如>0.5就可舍去<0.5的部分。
这里以13×13举例:
X、Y的偏移量是由:X/32 = R(整数—索引)+S(小数部分—偏移量)
原图真实坐标:(R+S)×32
W、H的偏移量是由:W/W’ = 偏移量log(其中W’为建议框的W)
原图 W:e^(网络输出值)×W’
import torch
from torch import nn
from net import Yolo_v3_Net
from PIL import Image,ImageDraw
from utils import *
from config import *
from Dataset import *
class_num = {
0:'person',
1:'horse',
2:'bicycle'
}
device = torch.device('cpu')
# 调用CPU
class Detect(nn.Module):
def __init__(self):
super(Detect,self).__init__()
self.weight = 'params/net.pt'
# 调用模型
self.net = Yolo_v3_Net().to(device)
# 导入网络到GPU/CPU
if os.path.exists(self.weight):
self.net.load_state_dict(torch.load(self.weight))
print('加载权重成功')
self.net.eval()
# 测试模型前使用net.eval(),加载标准化后的值(一定要加)
# net.load_state_dict()加载权重
def forward(self,input,thresh,anchors,case):
"""
在13×13 26×26 52×52进行相关计算
"""
output_13,output_26,output_52 = self.net(input)
index_13,bias_13 = self.get_index_and_bias(output_13,thresh)
boxes_13 = self.get_true_positon(index_13,bias_13,32,anchors[13],case)
index_26,bias_26 = self.get_index_and_bias(output_26,thresh)
boxes_26 = self.get_true_positon(index_26,bias_26,16,anchors[26],case)
index_52,bias_52 = self.get_index_and_bias(output_52,thresh)
boxes_52 = self.get_true_positon(index_52,bias_52,8,anchors[52],case)
return torch.cat([boxes_13,boxes_26,boxes_52],dim = 0)
# 返回真实的数据值
def get_index_and_bias(self,output,thresh):
# 索取索引和偏移量
output = output.permute(0,2,3,1)
# 转换成N H W 3 8
output = output.reshape(output.size(0),output.size(1),output.size(2),3,-1)
##改变格式为N H W 3个框 最后一位用“-1”代替(自动计算)
mask = output[...,0]>thresh
# 保存大小为thresh以上的置信度
index = mask.nonzero()
# 返回为True的坐标索引
bias = output[mask]
# 返回偏移量
return index,bias
def get_true_positon(self,index,bias,t,anchors):
#得到原图真实的位置,按上述讲解计算
anchors = torch.Tensor(anchors)
# 转换张量
a = index[:,3]
# 得到N H W 3中的3
cy = (index[:,1].float+bias[:,2].float())*t/case
cx = (index[:,2].float+bias[:,1].float())*t/case
# 计算中心点
w = anchors[a,0]*torch.exp(bias[:,3])/case
h = anchors[a, 1] * torch.exp(bias[:, 4])/case
# 计算W H
p = bias[:,0]
# 置信度
cls_p = bias[:,5:]
# 计算分类值
cls_index = torch.argmax(cls_p,dim = 1)
return torch.stack([torch.sigmoid(p),cx,cy,w,h,cls_index],dim = 1)
if __name__ == '__main__':
detector = Detect()
img = Image.open('images/1.jpg')
_img = make_416_Image('images/1.jpg')
#这里为图片的存储路径
temp = max(_img.size)
case = 416/temp
_img = _img.resize((416,416))
_img = torch.unsqueeze(_img,dim = 0)
# 降一个维度,使数据量对应
_img = tf(_img).to(device)
# 导入到GPU/CPU
results = detector(_img,0.3,antors,case)
# 存储结果,置信度在0.3以下的舍弃
draw = ImageDraw.Draw(img)
# 框出目标
for rst in results:
x1,y1,x2,y2 = rst[1]-0.5*rst[3],rst[2]-0.5*rst[4],rst[1]+0.5*rst[3],rst[2]+0.5*rst[4]
# 计算最终真实坐标
print(x1,y1,x2,y2)
print('class',class_num[int(rst[5])])
draw.text((x1,y1),str(class_num[int(rst[5].item())])+str(rst[0].item())[:4])
# 类别转换
draw.rectangle((x1,y1,x2,y2),outline='red',width = 1)
img.show()
这里对mask.nonzero()进行说明,以下用了一个测试代码:
import torch
x = torch.randn(5,13,13,3,8)
mask = x[...,0]>1
index = mask.nonzero()
print(mask)
print(mask.shape)
print(index.shape)
print(x[mask].shape)
代码效果如下:
mask.nonzero()就是把图中为True的索引返回。
注意:这里需要使用1.9版本以上的pytorch,否则argmax会报错张量空白。
11.16学习笔记
前两篇笔记:
YoloV3笔记(二):_风声向寂的博客-CSDN博客
(19条消息) YoloV3学习笔记_风声向寂的博客-CSDN博客