钱包网站开发产品线上营销方案
pytorch目标检测通用教程(包含目标检测基础知识汇总以及SSD的介绍)
之前写了很多分类网络,一直没时间写个目标检测的教程。(因为懒惰)
如果你也正在研究目标检测,可以直接套用这套代码,直接使用或者说是换成自己需要的网络。
最近正好复习一下之前写过的代码,就写一个通用的目标检测教程
之后如果需要更换训练的模型只需要替换其中的部分模块就可以了
PS:复习真的很重要,我最近常常复习之前写过的代码,收获颇丰
通过这个教程我们可以学会目标检测的代码的编写与SSD的基本知识,开箱即用,也可以根据自身的需求更换其中的模型文件。
废话不多说了,让我们一起来学习把。fighting!!!!
学习之前先介绍目标检测的基础知识
1.Single-Shot Detection(单阶段检测)
Single-Shot Detection将定位和检测任务封装在网络的单一前向扫描中,从而大大加快了检测速度,同时可以部署在更轻的硬件上。
PS:这里科普一下单阶段检测,就算是定位物体与分类在一个步骤中完成,2阶检测例如Faster-RCNN就是有RPN区域特征网络进行定位在特征图上投射出目标框,并在后续的网络中进行分类。
2.Multiscale Feature Maps(多尺度特征图)
在目标检测中,由于多次卷积小的目标在后面的特征图中已经丢失,最后导致不能检测细小的目标。
在分类网络中,我们可能只需要最后的最细节的特征来对图像进行分类,但是在目标检查中,早期的特征图也能够帮助我们进行预测。
3.Priors box(先验框)
根据特征的长宽比与比例在特征图上生成预先假设的框,以此来对真实框进行回归预测,这样子比从无到有生成一个框更便由于目标的定位。
3.Multiboxbox(多个预测框)
有多个特征图进行预测
4.Non-Maximum Suppression.(非极大抑制)
针对一个目标我的不止得到一个预测框,于是我们根据同一个物体上所有的框进行置信度排序,只保留最高置信度的框,同时也保障每个物体都能有一个框。
SSD网络介绍
SSD网络分为两个版本,他们是SSD300 and the SSD512. 后面的数字对应输入特征图的大小。
SSD根据特征提取网络的多个特征层进行预测,有效地解决了多尺度的问题。
由于篇幅问题这里不详细介绍SSD,只是简单介绍基本信息。
SSD的特征提取网络由VGG-16 构成
- 输入图像的大小将是300,300
- 将尺寸减半的第三个池化层在确定输出大小时将使用数学上的
ceiling
function,而不是默认的池化层。这只有在前一个特征图的尺寸是奇数而不是偶数时才有意义。通过看上面的图像,你可以计算出,对于我们输入的图像大小为300,300,conv3_3特征图的横截面为75,75,这将减半为38,38,而不是不方便的37,37 - 我们将第5池化层从2,2卷积核和2步幅修改为3,3卷积核和1步幅。这样做的效果是,feature map(特征图)的尺寸不再从之前的卷积层减半。
- 完全删除了原来特征提取网络里的全连接层,因为它们在这里没有任何作用。我们将完全抛弃fc8,但选择将fc6和fc7重新制作为卷积层conv6和conv7。(fc8,6,7)都是原来的vgg的全连接层。
代码编写部分
先构建ssd模型
from torch import nn
from utils import *
import torch.nn.functional as F
from math import sqrt
from itertools import product as product
import torchvision
import torchdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")#GPU的选用#构建vgg特征提取网络
class VGGBase(nn.Module):def __init__(self):super(VGGBase, self).__init__()self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1)self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True) # ceiling 函数self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1)self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1)self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)self.pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1) # 特殊池化层方便规定特征层的尺寸self.conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6) # atrous convolutionself.conv7 = nn.Conv2d(1024, 1024, kernel_size=1)self.load_pretrained_layers()def forward(self, image):out = F.relu(self.conv1_1(image)) # (N, 64, 300, 300)out = F.relu(self.conv1_2(out)) # (N, 64, 300, 300)out = self.pool1(out) # (N, 64, 150, 150)out = F.relu(self.conv2_1(out)) # (N, 128, 150, 150)out = F.relu(self.conv2_2(out)) # (N, 128, 150, 150)out = self.pool2(out) # (N, 128, 75, 75)out = F.relu(self.conv3_1(out)) # (N, 256, 75, 75)out = F.relu(self.conv3_2(out)) # (N, 256, 75, 75)out = F.relu(self.conv3_3(out)) # (N, 256, 75, 75)out = self.pool3(out) # (N, 256, 38, 38) 采用了ceiling 函数out = F.relu(self.conv4_1(out)) # (N, 512, 38, 38)out = F.relu(self.conv4_2(out)) # (N, 512, 38, 38)out = F.relu(self.conv4_3(out)) # (N, 512, 38, 38)conv4_3_feats = out # (N, 512, 38, 38)out = self.pool4(out) # (N, 512, 19, 19)out = F.relu(self.conv5_1(out)) # (N, 512, 19, 19)out = F.relu(self.conv5_2(out)) # (N, 512, 19, 19)out = F.relu(self.conv5_3(out)) # (N, 512, 19, 19)out = self.pool5(out) # (N, 512, 19, 19), pool5 不改变尺寸out = F.relu(self.conv6(out)) # (N, 1024, 19, 19)conv7_feats = F.relu(self.conv7(out)) # (N, 1024, 19, 19)return conv4_3_feats, conv7_feats #针对这两个特征层来进行后续的目标检测#加载预选连权重def load_pretrained_layers(self):# 新建的模型权重state_dict = self.state_dict()param_names = list(state_dict.keys())#预训练模型权重pretrained_state_dict = torchvision.models.vgg16(pretrained=True).state_dict()pretrained_param_names = list(pretrained_state_dict.keys())#加载修改之前的模型for i, param in enumerate(param_names[:-4]): state_dict[param] = pretrained_state_dict[pretrained_param_names[i]]# fc6conv_fc6_weight = pretrained_state_dict['classifier.0.weight'].view(4096, 512, 7, 7) # (4096, 512, 7, 7)conv_fc6_bias = pretrained_state_dict['classifier.0.bias'] # (4096)state_dict['conv6.weight'] = decimate(conv_fc6_weight, m=[4, None, 3, 3]) # (1024, 512, 3, 3)state_dict['conv6.bias'] = decimate(conv_fc6_bias, m=[4]) # (1024)# fc7conv_fc7_weight = pretrained_state_dict['classifier.3.weight'].view(4096, 4096, 1, 1) # (4096, 4096, 1, 1)conv_fc7_bias = pretrained_state_dict['classifier.3.bias'] # (4096)state_dict['conv7.weight'] = decimate(conv_fc7_weight, m=[4, 4, None, None]) # (1024, 1024, 1, 1)state_dict['conv7.bias'] = decimate(conv_fc7_bias, m=[4]) # (1024)self.load_state_dict(state_dict)print("\nLoaded base model.\n")class AuxiliaryConvolutions(nn.Module):#多尺度预测的网络部分def __init__(self):super(AuxiliaryConvolutions, self).__init__()self.conv8_1 = nn.Conv2d(1024, 256, kernel_size=1, padding=0)self.conv8_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1) self.conv9_1 = nn.Conv2d(512, 128, kernel_size=1, padding=0)self.conv9_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1) self.conv10_1 = nn.Conv2d(256, 128, kernel_size=1, padding=0)self.conv10_2 = nn.Conv2d(128, 256, kernel_size=3, padding=0) self.conv11_1 = nn.Conv2d(256, 128, kernel_size=1, padding=0)self.conv11_2 = nn.Conv2d(128, 256, kernel_size=3, padding=0) self.init_conv2d()def init_conv2d(self):for c in self.children():if isinstance(c, nn.Conv2d):nn.init.xavier_uniform_(c.weight)nn.init.constant_(c.bias, 0.)def forward(self, conv7_feats):out = F.relu(self.conv8_1(conv7_feats)) # (N, 256, 19, 19)out = F.relu(self.conv8_2(out)) # (N, 512, 10, 10)conv8_2_feats = out # (N, 512, 10, 10)out = F.relu(self.conv9_1(out)) # (N, 128, 10, 10)out = F.relu(self.conv9_2(out)) # (N, 256, 5, 5)conv9_2_feats = out # (N, 256, 5, 5)out = F.relu(self.conv10_1(out)) # (N, 128, 5, 5)out = F.relu(self.conv10_2(out)) # (N, 256, 3, 3)conv10_2_feats = out # (N, 256, 3, 3)out = F.relu(self.conv11_1(out)) # (N, 128, 3, 3)conv11_2_feats = F.relu(self.conv11_2(out)) # (N, 256, 1, 1)return conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats#这是几个预测层class PredictionConvolutions(nn.Module):#根据之前的得到的特征层进行预测def __init__(self, n_classes):super(PredictionConvolutions, self).__init__()self.n_classes = n_classesn_boxes = {'conv4_3': 4, #之前在特征图上预先产生的先验框的数量与卷积核对应'conv7': 6,'conv8_2': 6,'conv9_2': 6,'conv10_2': 4,'conv11_2': 4}#每个坐标有四个参数self.loc_conv4_3 = nn.Conv2d(512, n_boxes['conv4_3'] * 4, kernel_size=3, padding=1)self.loc_conv7 = nn.Conv2d(1024, n_boxes['conv7'] * 4, kernel_size=3, padding=1)self.loc_conv8_2 = nn.Conv2d(512, n_boxes['conv8_2'] * 4, kernel_size=3, padding=1)self.loc_conv9_2 = nn.Conv2d(256, n_boxes['conv9_2'] * 4, kernel_size=3, padding=1)self.loc_conv10_2 = nn.Conv2d(256, n_boxes['conv10_2'] * 4, kernel_size=3, padding=1)self.loc_conv11_2 = nn.Conv2d(256, n_boxes['conv11_2'] * 4, kernel_size=3, padding=1)#计算其中的分类卷积self.cl_conv4_3 = nn.Conv2d(512, n_boxes['conv4_3'] * n_classes, kernel_size=3, padding=1)self.cl_conv7 = nn.Conv2d(1024, n_boxes['conv7'] * n_classes, kernel_size=3, padding=1)self.cl_conv8_2 = nn.Conv2d(512, n_boxes['conv8_2'] * n_classes, kernel_size=3, padding=1)self.cl_conv9_2 = nn.Conv2d(256, n_boxes['conv9_2'] * n_classes, kernel_size=3, padding=1)self.cl_conv10_2 = nn.Conv2d(256, n_boxes['conv10_2'] * n_classes, kernel_size=3, padding=1)self.cl_conv11_2 = nn.Conv2d(256, n_boxes['conv11_2'] * n_classes, kernel_size=3, padding=1)# 初始化卷积核self.init_conv2d()def init_conv2d(self):for c in self.children():if isinstance(c, nn.Conv2d):nn.init.xavier_uniform_(c.weight)nn.init.constant_(c.bias, 0.)def forward(self, conv4_3_feats, conv7_feats, conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats):batch_size = conv4_3_feats.size(0)#计算其中的框的回归参数l_conv4_3 = self.loc_conv4_3(conv4_3_feats) # (N, 16, 38, 38)l_conv4_3 = l_conv4_3.permute(0, 2, 3,1).contiguous() # (N, 38, 38, 16), 把通道数放到后面去便于之后的计算l_conv4_3 = l_conv4_3.view(batch_size, -1, 4) # (N, 5776, 4)把每个框放到中间l_conv7 = self.loc_conv7(conv7_feats) # (N, 24, 19, 19)l_conv7 = l_conv7.permute(0, 2, 3, 1).contiguous() # (N, 19, 19, 24)l_conv7 = l_conv7.view(batch_size, -1, 4) # (N, 2166, 4), l_conv8_2 = self.loc_conv8_2(conv8_2_feats) # (N, 24, 10, 10)l_conv8_2 = l_conv8_2.permute(0, 2, 3, 1).contiguous() # (N, 10, 10, 24)l_conv8_2 = l_conv8_2.view(batch_size, -1, 4) # (N, 600, 4)l_conv9_2 = self.loc_conv9_2(conv9_2_feats) # (N, 24, 5, 5)l_conv9_2 = l_conv9_2.permute(0, 2, 3, 1).contiguous() # (N, 5, 5, 24)l_conv9_2 = l_conv9_2.view(batch_size, -1, 4) # (N, 150, 4)l_conv10_2 = self.loc_conv10_2(conv10_2_feats) # (N, 16, 3, 3)l_conv10_2 = l_conv10_2.permute(0, 2, 3, 1).contiguous() # (N, 3, 3, 16)l_conv10_2 = l_conv10_2.view(batch_size, -1, 4) # (N, 36, 4)l_conv11_2 = self.loc_conv11_2(conv11_2_feats) # (N, 16, 1, 1)l_conv11_2 = l_conv11_2.permute(0, 2, 3, 1).contiguous() # (N, 1, 1, 16)l_conv11_2 = l_conv11_2.view(batch_size, -1, 4) # (N, 4, 4)#计算检测的分类c_conv4_3 = self.cl_conv4_3(conv4_3_feats) # (N, 4 * n_classes, 38, 38)c_conv4_3 = c_conv4_3.permute(0, 2, 3,1).contiguous() # (N, 38, 38, 4 * n_classes)把通道数放到后面去了c_conv4_3 = c_conv4_3.view(batch_size, -1,self.n_classes) # (N, 5776, n_classes)计算每个类别的概率c_conv7 = self.cl_conv7(conv7_feats) # (N, 6 * n_classes, 19, 19)c_conv7 = c_conv7.permute(0, 2, 3, 1).contiguous() # (N, 19, 19, 6 * n_classes)c_conv7 = c_conv7.view(batch_size, -1, self.n_classes) # (N, 2166, n_classes),c_conv8_2 = self.cl_conv8_2(conv8_2_feats) # (N, 6 * n_classes, 10, 10)c_conv8_2 = c_conv8_2.permute(0, 2, 3, 1).contiguous() # (N, 10, 10, 6 * n_classes)c_conv8_2 = c_conv8_2.view(batch_size, -1, self.n_classes) # (N, 600, n_classes)c_conv9_2 = self.cl_conv9_2(conv9_2_feats) # (N, 6 * n_classes, 5, 5)c_conv9_2 = c_conv9_2.permute(0, 2, 3, 1).contiguous() # (N, 5, 5, 6 * n_classes)c_conv9_2 = c_conv9_2.view(batch_size, -1, self.n_classes) # (N, 150, n_classes)c_conv10_2 = self.cl_conv10_2(conv10_2_feats) # (N, 4 * n_classes, 3, 3)c_conv10_2 = c_conv10_2.permute(0, 2, 3, 1).contiguous() # (N, 3, 3, 4 * n_classes)c_conv10_2 = c_conv10_2.view(batch_size, -1, self.n_classes) # (N, 36, n_classes)c_conv11_2 = self.cl_conv11_2(conv11_2_feats) # (N, 4 * n_classes, 1, 1)c_conv11_2 = c_conv11_2.permute(0, 2, 3, 1).contiguous() # (N, 1, 1, 4 * n_classes)c_conv11_2 = c_conv11_2.view(batch_size, -1, self.n_classes) # (N, 4, n_classes)#把所有的边框回归于分类集合在一起locs = torch.cat([l_conv4_3, l_conv7, l_conv8_2, l_conv9_2, l_conv10_2, l_conv11_2], dim=1) # (N, 8732, 4)classes_scores = torch.cat([c_conv4_3, c_conv7, c_conv8_2, c_conv9_2, c_conv10_2, c_conv11_2],dim=1) # (N, 8732, n_classes)return locs, classes_scores #返回得到的预测的框于分类的情况#构建SSD300网络
class SSD300(nn.Module):def __init__(self, n_classes):super(SSD300, self).__init__()self.n_classes = n_classesself.base = VGGBase() #基础特征提取层 前面两个尺度的网络self.aux_convs = AuxiliaryConvolutions()#多尺度辅助预测 后面几个尺度的网络self.pred_convs = PredictionConvolutions(n_classes)#预测物体的位置以及分类#L2正则化 self.rescale_factors = nn.Parameter(torch.FloatTensor(1, 512, 1, 1)) # 512通道数在conv4_3_featsnn.init.constant_(self.rescale_factors, 20)# 在特征图上画框框self.priors_cxcy = self.create_prior_boxes()def forward(self, image):#通过前面的vgg初步提取特征 返回两个尺度的特征层conv4_3_feats, conv7_feats = self.base(image) # (N, 512, 38, 38), (N, 1024, 19, 19)# L2正则化norm = conv4_3_feats.pow(2).sum(dim=1, keepdim=True).sqrt() # (N, 1, 38, 38)conv4_3_feats = conv4_3_feats / norm # (N, 512, 38, 38)conv4_3_feats = conv4_3_feats * self.rescale_factors # (N, 512, 38, 38)#后续的特征提取网络conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats = \self.aux_convs(conv7_feats) # (N, 512, 10, 10), (N, 256, 5, 5), (N, 256, 3, 3), (N, 256, 1, 1)# 利用网路预测几张特征图上的 图像类别以及位置locs, classes_scores = self.pred_convs(conv4_3_feats, conv7_feats, conv8_2_feats, conv9_2_feats, conv10_2_feats,conv11_2_feats) # (N, 8732, 4), (N, 8732, n_classes)return locs, classes_scoresdef create_prior_boxes(self):#在特征图上创建先眼眶fmap_dims = {'conv4_3': 38,'conv7': 19,'conv8_2': 10,'conv9_2': 5,'conv10_2': 3,'conv11_2': 1}obj_scales = {'conv4_3': 0.1,'conv7': 0.2,'conv8_2': 0.375,'conv9_2': 0.55,'conv10_2': 0.725,'conv11_2': 0.9}aspect_ratios = {'conv4_3': [1., 2., 0.5],'conv7': [1., 2., 3., 0.5, .333],'conv8_2': [1., 2., 3., 0.5, .333],'conv9_2': [1., 2., 3., 0.5, .333],'conv10_2': [1., 2., 0.5],'conv11_2': [1., 2., 0.5]}fmaps = list(fmap_dims.keys())prior_boxes = []for k, fmap in enumerate(fmaps):for i in range(fmap_dims[fmap]):for j in range(fmap_dims[fmap]):cx = (j + 0.5) / fmap_dims[fmap]cy = (i + 0.5) / fmap_dims[fmap]for ratio in aspect_ratios[fmap]:prior_boxes.append([cx, cy, obj_scales[fmap] * sqrt(ratio), obj_scales[fmap] / sqrt(ratio)])if ratio == 1.:try:additional_scale = sqrt(obj_scales[fmap] * obj_scales[fmaps[k + 1]])except IndexError:additional_scale = 1.prior_boxes.append([cx, cy, additional_scale, additional_scale])prior_boxes = torch.FloatTensor(prior_boxes).to(device) # (8732, 4)prior_boxes.clamp_(0, 1) # (8732, 4) 裁剪多出边界的框return prior_boxesdef detect_objects(self, predicted_locs, predicted_scores, min_score, max_overlap, top_k):#检测物体batch_size = predicted_locs.size(0)n_priors = self.priors_cxcy.size(0)predicted_scores = F.softmax(predicted_scores, dim=2) # (N, 8732, n_classes) 计算分类的分数# 将所有信息的载体先创建表格来存放all_images_boxes = list()all_images_labels = list()all_images_scores = list()assert n_priors == predicted_locs.size(1) == predicted_scores.size(1)#判断数量正确与否for i in range(batch_size):#从之前的计算的数值中解码decoded_locs = cxcy_to_xy(gcxgcy_to_cxcy(predicted_locs[i], self.priors_cxcy)) # (8732, 4), 计算出坐标# 创建列表存储信息image_boxes = list()image_labels = list()image_scores = list()max_scores, best_label = predicted_scores[i].max(dim=1) # (8732)for c in range(1, self.n_classes):class_scores = predicted_scores[i][:, c] # (8732)score_above_min_score = class_scores > min_score #大于分数的门限n_above_min_score = score_above_min_score.sum().item()if n_above_min_score == 0:continueclass_scores = class_scores[score_above_min_score]class_decoded_locs = decoded_locs[score_above_min_score] class_scores, sort_ind = class_scores.sort(dim=0, descending=True) class_decoded_locs = class_decoded_locs[sort_ind] overlap = find_jaccard_overlap(class_decoded_locs, class_decoded_locs) # (n_qualified, n_min_score)# (NMS)非极大抑制的过程suppress = torch.zeros((n_above_min_score), dtype=torch.uint8).to(device) # (n_qualified)for box in range(class_decoded_locs.size(0)):if suppress[box] == 1:continuesuppress = torch.max(suppress, overlap[box] > max_overlap)suppress[box] = 0image_boxes.append(class_decoded_locs[1 - suppress])image_labels.append(torch.LongTensor((1 - suppress).sum().item() * [c]).to(device))image_scores.append(class_scores[1 - suppress])if len(image_boxes) == 0:image_boxes.append(torch.FloatTensor([[0., 0., 1., 1.]]).to(device))image_labels.append(torch.LongTensor([0]).to(device))image_scores.append(torch.FloatTensor([0.]).to(device))image_boxes = torch.cat(image_boxes, dim=0) # (n_objects, 4)image_labels = torch.cat(image_labels, dim=0) # (n_objects)image_scores = torch.cat(image_scores, dim=0) # (n_objects)n_objects = image_scores.size(0)if n_objects > top_k:image_scores, sort_ind = image_scores.sort(dim=0, descending=True)image_scores = image_scores[:top_k] # (top_k)image_boxes = image_boxes[sort_ind][:top_k] # (top_k, 4)image_labels = image_labels[sort_ind][:top_k] # (top_k)all_images_boxes.append(image_boxes)all_images_labels.append(image_labels)all_images_scores.append(image_scores)return all_images_boxes, all_images_labels, all_images_scores #损失函数 来更新网络
class MultiBoxLoss(nn.Module):def __init__(self, priors_cxcy, threshold=0.5, neg_pos_ratio=3, alpha=1.):super(MultiBoxLoss, self).__init__()self.priors_cxcy = priors_cxcy self.priors_xy = cxcy_to_xy(priors_cxcy)self.threshold = threshold #判断门限self.neg_pos_ratio = neg_pos_ratioself.alpha = alphaself.smooth_l1 = nn.L1Loss()#光滑标签来计算框的回归损失self.cross_entropy = nn.CrossEntropyLoss(reduce=False)#交叉熵计算分类损失def forward(self, predicted_locs, predicted_scores, boxes, labels):batch_size = predicted_locs.size(0)n_priors = self.priors_cxcy.size(0)n_classes = predicted_scores.size(2)assert n_priors == predicted_locs.size(1) == predicted_scores.size(1)true_locs = torch.zeros((batch_size, n_priors, 4), dtype=torch.float).to(device) # (N, 8732, 4)true_classes = torch.zeros((batch_size, n_priors), dtype=torch.long).to(device) # (N, 8732)for i in range(batch_size):n_objects = boxes[i].size(0)overlap = find_jaccard_overlap(boxes[i],self.priors_xy) # (n_objects, 8732)# 根据真实框找到最大的iou的框overlap_for_each_prior, object_for_each_prior = overlap.max(dim=0) # (8732)_, prior_for_each_object = overlap.max(dim=1) # (N_o)object_for_each_prior[prior_for_each_object] = torch.LongTensor(range(n_objects)).to(device)overlap_for_each_prior[prior_for_each_object] = 1.label_for_each_prior = labels[i][object_for_each_prior] # (8732)label_for_each_prior[overlap_for_each_prior < self.threshold] = 0 # (8732)true_classes[i] = label_for_each_prior#对预测框进行编码true_locs[i] = cxcy_to_gcxgcy(xy_to_cxcy(boxes[i][object_for_each_prior]), self.priors_cxcy) # (8732, 4)# 正样本的判断positive_priors = true_classes != 0 # (N, 8732)#计算位置的回归损失loc_loss = self.smooth_l1(predicted_locs[positive_priors], true_locs[positive_priors]) # (), scalarn_positives = positive_priors.sum(dim=1) # (N)n_hard_negatives = self.neg_pos_ratio * n_positives # (N)# 计算全部的损失conf_loss_all = self.cross_entropy(predicted_scores.view(-1, n_classes), true_classes.view(-1)) # (N * 8732)conf_loss_all = conf_loss_all.view(batch_size, n_priors) # (N, 8732)conf_loss_pos = conf_loss_all[positive_priors] # (sum(n_positives))# 找难的负样本# 存储负样本conf_loss_neg = conf_loss_all.clone() # (N, 8732)conf_loss_neg[positive_priors] = 0. # (N, 8732), conf_loss_neg, _ = conf_loss_neg.sort(dim=1, descending=True) # (N, 8732)hardness_ranks = torch.LongTensor(range(n_priors)).unsqueeze(0).expand_as(conf_loss_neg).to(device) # (N, 8732)hard_negatives = hardness_ranks < n_hard_negatives.unsqueeze(1) # (N, 8732)conf_loss_hard_neg = conf_loss_neg[hard_negatives] # (sum(n_hard_negatives))conf_loss = (conf_loss_hard_neg.sum() + conf_loss_pos.sum()) / n_positives.sum().float() # (), scalarreturn conf_loss + self.alpha * loc_loss #最后的总损失
训练代码以及预测代码可以访问
代码地址github
由于我之前已经跑过一遍了,就演示到这里了
(跑完完整的训练时间太长了)这里我们仅仅验证能训练模型
而且loss一直在下降
你们可以自己去跑一跑锻炼一下自己代码能力
最后测试77.2 mAP左右(2007-TEST) train(2007+2012)
昨天晚上发现程序忘记关了运行了一晚上虽然没完全训练好不过 可以作为一个例子
{‘aeroplane’: 0.7419820427894592,
‘bicycle’: 0.8121547698974609,
‘bird’: 0.6819309592247009,
‘boat’: 0.6568214297294617,
‘bottle’: 0.35232362151145935,
‘bus’: 0.8188199400901794,
‘car’: 0.8169978857040405,
‘cat’: 0.858344554901123,
‘chair’: 0.5217813849449158,
‘cow’: 0.7868558764457703,
‘diningtable’: 0.6424081921577454,
‘dog’: 0.7965350151062012,
‘horse’: 0.832635223865509,
‘motorbike’: 0.784022331237793,
‘person’: 0.7415710091590881,
‘pottedplant’: 0.4017500579357147,
‘sheep’: 0.7026932835578918,
‘sofa’: 0.7310159802436829,
‘train’: 0.8186013102531433,
‘tvmonitor’: 0.7230682373046875}
Mean Average Precision (mAP): 0.711
评估代码
from utils import *
from datasets import PascalVOCDataset
from tqdm.notebook import tqdm
from pprint import PrettyPrinter# Good formatting when printing the APs for each class and mAP
pp = PrettyPrinter()# Parameters
data_folder = './'
keep_difficult = True # difficult ground truth objects must always be considered in mAP calculation, because these objects DO exist!
batch_size = 64
workers = 4
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = './checkpoint_ssd300.pth.tar'# Load model checkpoint that is to be evaluated
checkpoint = torch.load(checkpoint)
model = checkpoint['model']
model = model.to(device)# Switch to eval mode
model.eval()# Load test data
test_dataset = PascalVOCDataset(data_folder,split='test',keep_difficult=keep_difficult)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False,collate_fn=test_dataset.collate_fn, num_workers=workers, pin_memory=True)def evaluate(test_loader, model):"""Evaluate.:param test_loader: DataLoader for test data:param model: model"""# Make sure it's in eval modemodel.eval()# Lists to store detected and true boxes, labels, scoresdet_boxes = list()det_labels = list()det_scores = list()true_boxes = list()true_labels = list()true_difficulties = list() # it is necessary to know which objects are 'difficult', see 'calculate_mAP' in utils.pywith torch.no_grad():# Batchesfor i, (images, boxes, labels, difficulties) in enumerate(tqdm(test_loader, desc='Evaluating')):images = images.to(device) # (N, 3, 300, 300)# Forward prop.predicted_locs, predicted_scores = model(images)# Detect objects in SSD outputdet_boxes_batch, det_labels_batch, det_scores_batch = model.detect_objects(predicted_locs, predicted_scores,min_score=0.01, max_overlap=0.45,top_k=200)# Evaluation MUST be at min_score=0.01, max_overlap=0.45, top_k=200 for fair comparision with the paper's results and other repos# Store this batch's results for mAP calculationboxes = [b.to(device) for b in boxes]labels = [l.to(device) for l in labels]difficulties = [d.to(device) for d in difficulties]det_boxes.extend(det_boxes_batch)det_labels.extend(det_labels_batch)det_scores.extend(det_scores_batch)true_boxes.extend(boxes)true_labels.extend(labels)true_difficulties.extend(difficulties)# Calculate mAPAPs, mAP = calculate_mAP(det_boxes, det_labels, det_scores, true_boxes, true_labels, true_difficulties)# Print AP for each classpp.pprint(APs)print('\nMean Average Precision (mAP): %.3f' % mAP)if __name__ == '__main__':evaluate(test_loader, model)
再贴上几个预测的图片把