YoloV3学习笔记(三):


文章目录

  • YoloV3学习笔记(三):
  • 1.训练部分
  • 2.侦测部分代码



  直接上代码:

1.训练部分

   因为我们最后输出的是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

效果如下:

csdn yolo模型权重文件_算法

tensorboard效果:

csdn yolo模型权重文件_csdn yolo模型权重文件_02

2.侦测部分代码

    首先,输入一张图片以后,需要给一个置信度约束那些置信度很低的值,比如>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)

代码效果如下:

csdn yolo模型权重文件_机器学习_03

mask.nonzero()就是把图中为True的索引返回。
注意:这里需要使用1.9版本以上的pytorch,否则argmax会报错张量空白。

11.16学习笔记