[学习笔记] SSD代码笔记 + EifficientNet backbone 练习

SSD代码笔记 + EifficientNet backbone 练习

ssd代码完全ok了，然后用最近性能和速度都非常牛的Eifficient Net做backbone设计了自己的TinySSD网络，没有去调参，所以网络并没有很好的收敛，之后我会调一调，实际去应用。

torch.clamp

torch.clamp(input, min, max, out=None) → Tensor

就是clip的功能

eg:

>>> a = torch.randn(4)
>>> a
tensor([-1.7120,  0.1734, -0.0478, -0.0922])
>>> torch.clamp(a, min=-0.5, max=0.5)
tensor([-0.5000,  0.1734, -0.0478, -0.0922])

计算iou

交集除以并集，首先要算面积，面积就是给定两个点坐标求宽高乘积即可。

交集面积就是两个框离原点最远的左上角点与离原点最近的右下角点组成区域的面积。

def area_of(left_top,right_bottom):# (num_boxes,2),(num_boxes,2)
	hw = torch.clamp(right_bottom-left_top,0.0) # (num_boxes,2)
    # 这里做clip的原因是如果框不重叠的化，如果不clip算出来就是负值，有了clip就是0
    return hw[...,0] * hw[...,1]

def iou_of(boxes0,boxes1,eps = 1e-5):# (N,4) and (1,4) or (N,4)
    # 注意这里其实boxes0和boxes1和size其实是不一样的，所以size较少的那个会broadcast到较大的那个然后在做max和min操作。
    overlap_left_top = torch.max(boxes0[...,:2],boxes1[...:2])# 左上角点最远的
    overlap_right_bottom = torch.min(boxes0[...,2:],boxes1[...,2:]) # 右下角点最近的
    area0 = area_of(boxes0[...,:2],boxes0[...,2:]) # 左上角和右下角点计算面积
    area1 = area_of(boxes1[...,:2],boxes1[...,2:]) # predict box的面积
    overlap_area = area_of(overlap_left_top,over_right_bottom) # 并集面积
    return overlap_area / (area0 + area1 - overlap_area + eps)

生成priorbox

这里生成坐标的方法是真学到了，product构造映射。实际上featur map上每个点都是加了0.5作为中心，然后除以ratio，这里ratio一般和feature map 的size是不一样的，取决于网络设计，我这里用的刚好一样。ratio就是对应到原图有多少个滑窗。后面考虑了其他尺度的anchor。

#有的注释是原来代码里的，有的英文注释是我加的。
class PriorBox(nn.Module):
    def __init__(self):
        super(PriorBox, self).__init__()
        self.image_size = 512
        self.feature_maps = [16,8,4,2,1]
        self.min_sizes = [30,60,111,162,213]
        self.max_sizes = [60,111,162,213,512]
        self.strides = [32,64,128,256,512]
        self.aspect_ratios = [[2], [2, 3], [2, 3], [2], [2]]
        self.clip = True

    def forward(self):
        """Generate SSD Prior Boxes.
            It returns the center, height and width of the priors. The values are relative to the image size
            Returns:
                priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values
                    are relative to the image size.
        """
        priors = []
        for k, f in enumerate(self.feature_maps): # every size of feature map
            scale = self.image_size / self.strides[k] # how many boxes (not anchor) in a row in raw img
            # 512 / 32 = 16
            for i, j in product(range(f), repeat=2): # xy generator in feature map
                # unit center x,y
                cx = (j + 0.5) / scale # see as blocks and xy in center of it
                cy = (i + 0.5) / scale # 15,15 -> 15.5,15.5 -> 15.5/16,15.5/16 which means the xy in center of feature map

                # small sized square box
                size = self.min_sizes[k] # min size
                h = w = size / self.image_size # small size
                priors.append([cx, cy, w, h]) # the small size one

                # big sized square box
                size = sqrt(self.min_sizes[k] * self.max_sizes[k]) # the same as small one
                h = w = size / self.image_size
                priors.append([cx, cy, w, h])

                # change h/w ratio of the small sized box
                # considering the w/ratio , w*ratio , h/ratio and h * ratio
                size = self.min_sizes[k]
                h = w = size / self.image_size
                for ratio in self.aspect_ratios[k]:
                    ratio = sqrt(ratio)
                    priors.append([cx, cy, w * ratio, h / ratio])
                    priors.append([cx, cy, w / ratio, h * ratio])

        priors = torch.Tensor(priors)
        if self.clip:
            priors.clamp_(max=1, min=0)
        return priors

priorbox的分配

很好的利用了broadcast机制，计算每个iou，然后得到target与所有prior重叠度最高的匹配，以及prior与target重叠度最高的匹配，然后通过阈值滤去。

def assign_priors(gt_boxes, gt_labels, corner_form_priors,
                  iou_threshold):
    """Assign ground truth boxes and targets to priors.

    Args:
        gt_boxes (num_targets, 4): ground truth boxes.
        gt_labels (num_targets): labels of targets.
        priors (num_priors, 4): corner form priors
    Returns:
        boxes (num_priors, 4): real values for priors.
        labels (num_priros): labels for priors.
    """
    # size: num_priors x num_targets
    ious = iou_of(gt_boxes.unsqueeze(0), corner_form_priors.unsqueeze(1))
    # size: num_priors
    best_target_per_prior, best_target_per_prior_index = ious.max(1) # 每个prior的iou最大的值以及在target里的索引
    # size: num_targets
    best_prior_per_target, best_prior_per_target_index = ious.max(0) # 每个target与所有prior的iou最大值以及在priors里的索引

    for target_index, prior_index in enumerate(best_prior_per_target_index):
        best_target_per_prior_index[prior_index] = target_index # 让每个Prior对应iou最大的target (0,0,1,2,3)
    # 2.0 is used to make sure every target has a prior assigned
    best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2) # dim = 0 ,value = 2,只要重叠的iou最大，就认为其重叠度是2
    # size: num_priors
    labels = gt_labels[best_target_per_prior_index] # num_priors，先按照iou最大分
    labels[best_target_per_prior < iou_threshold] = 0  # the backgournd id，小于阈值的认为是背景，有的iou尽管最大但是其iou还是很小，所以也需要滤去
    boxes = gt_boxes[best_target_per_prior_index] # 直接给box
    return boxes, labels

hard_negative_mining

通过给出mask考虑算哪些loss不算哪些loss，因为负样本实在太多了，所以这是一个方法。

def hard_negative_mining(loss, labels, neg_pos_ratio):
    """
    It used to suppress the presence of a large number of negative prediction.
    It works on image level not batch level.
    For any example/image, it keeps all the positive predictions and
     cut the number of negative predictions to make sure the ratio
     between the negative examples and positive examples is no more
     the given ratio for an image.

    Args:
        loss (N, num_priors): the loss for each example.
        labels (N, num_priors): the labels.
        neg_pos_ratio:  the ratio between the negative examples and positive examples.
    """
    pos_mask = labels > 0
    num_pos = pos_mask.long().sum(dim=1, keepdim=True)
    num_neg = num_pos * neg_pos_ratio

    loss[pos_mask] = -math.inf
    _, indexes = loss.sort(dim=1, descending=True)
    _, orders = indexes.sort(dim=1)
    neg_mask = orders < num_neg
    return pos_mask | neg_mask

Loss Function

bbox用smotth L1 loss，交叉熵分类loss。

class MultiBoxLoss(nn.Module):
    def __init__(self, neg_pos_ratio):
        """Implement SSD MultiBox Loss.

        Basically, MultiBox loss combines classification loss
         and Smooth L1 regression loss.
        """
        super(MultiBoxLoss, self).__init__()
        self.neg_pos_ratio = neg_pos_ratio

    def forward(self, confidence, predicted_locations, labels, gt_locations):
        """Compute classification loss and smooth l1 loss.

        Args:
            confidence (batch_size, num_priors, num_classes): class predictions.
            predicted_locations (batch_size, num_priors, 4): predicted locations.
            labels (batch_size, num_priors): real labels of all the priors.
            gt_locations (batch_size, num_priors, 4): real boxes corresponding all the priors.
        """
        num_classes = confidence.size(2)
        with torch.no_grad():
            # derived from cross_entropy=sum(log(p))
            loss = -F.log_softmax(confidence, dim=2)[:, :, 0]
            mask = box_utils.hard_negative_mining(loss, labels, self.neg_pos_ratio)

        confidence = confidence[mask, :]
        #print(confidence.view(-1, num_classes))
        #print(labels[mask])
        classification_loss = F.cross_entropy(confidence.view(-1, num_classes), labels[mask], reduction='sum')

        pos_mask = labels > 0
        predicted_locations = predicted_locations[pos_mask, :].view(-1, 4)
        gt_locations = gt_locations[pos_mask, :].view(-1, 4)
        smooth_l1_loss = F.smooth_l1_loss(predicted_locations, gt_locations, reduction='sum')
        num_pos = gt_locations.size(0)
        return smooth_l1_loss / num_pos, classification_loss / num_pos

Model

我看网络模型搜索得到的Eifficient Net性能和速度都是最优，直接拿来做backbone，但是调参还没调好，只是直接用其输出然后再加5层卷积层分别做特征金字塔，感觉感受野可能太大了，网络收敛性能不是很好，后面会调好参的，但是还是可以跑的。

使用EFnet作为后端的训练效果：

Eifficient Net Model

import torch
from torch import nn
from torch.nn import functional as F

from .utils import (
    relu_fn,
    round_filters,
    round_repeats,
    drop_connect,
    Conv2dSamePadding,
    get_model_params,
    efficientnet_params,
    load_pretrained_weights,
)

class MBConvBlock(nn.Module):
    """
    Mobile Inverted Residual Bottleneck Block

    Args:
        block_args (namedtuple): BlockArgs, see above
        global_params (namedtuple): GlobalParam, see above

    Attributes:
        has_se (bool): Whether the block contains a Squeeze and Excitation layer.
    """

    def __init__(self, block_args, global_params):
        super().__init__()
        self._block_args = block_args
        self._bn_mom = 1 - global_params.batch_norm_momentum
        self._bn_eps = global_params.batch_norm_epsilon
        self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)
        self.id_skip = block_args.id_skip  # skip connection and drop connect

        # Expansion phase
        inp = self._block_args.input_filters  # number of input channels
        oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels
        if self._block_args.expand_ratio != 1:
            self._expand_conv = Conv2dSamePadding(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
            self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)

        # Depthwise convolution phase
        k = self._block_args.kernel_size
        s = self._block_args.stride
        self._depthwise_conv = Conv2dSamePadding(
            in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise
            kernel_size=k, stride=s, bias=False)
        self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)

        # Squeeze and Excitation layer, if desired
        if self.has_se:
            num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
            self._se_reduce = Conv2dSamePadding(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
            self._se_expand = Conv2dSamePadding(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)

        # Output phase
        final_oup = self._block_args.output_filters
        self._project_conv = Conv2dSamePadding(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
        self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)

    def forward(self, inputs, drop_connect_rate=None):
        """
        :param inputs: input tensor
        :param drop_connect_rate: drop connect rate (float, between 0 and 1)
        :return: output of block
        """

        # Expansion and Depthwise Convolution
        x = inputs
        if self._block_args.expand_ratio != 1:
            x = relu_fn(self._bn0(self._expand_conv(inputs)))
        x = relu_fn(self._bn1(self._depthwise_conv(x)))

        # Squeeze and Excitation
        if self.has_se:
            x_squeezed = F.adaptive_avg_pool2d(x, 1)
            x_squeezed = self._se_expand(relu_fn(self._se_reduce(x_squeezed)))
            x = torch.sigmoid(x_squeezed) * x

        x = self._bn2(self._project_conv(x))

        # Skip connection and drop connect
        input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters
        if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters:
            if drop_connect_rate:
                x = drop_connect(x, p=drop_connect_rate, training=self.training)
            x = x + inputs  # skip connection
        return x

class EfficientNet(nn.Module):
    """
    An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods

    Args:
        blocks_args (list): A list of BlockArgs to construct blocks
        global_params (namedtuple): A set of GlobalParams shared between blocks

    Example:
        model = EfficientNet.from_pretrained('efficientnet-b0')

    """

    def __init__(self, blocks_args=None, global_params=None):
        super().__init__()
        assert isinstance(blocks_args, list), 'blocks_args should be a list'
        assert len(blocks_args) > 0, 'block args must be greater than 0'
        self._global_params = global_params
        self._blocks_args = blocks_args

        # Batch norm parameters
        bn_mom = 1 - self._global_params.batch_norm_momentum
        bn_eps = self._global_params.batch_norm_epsilon

        # Stem
        in_channels = 3  # rgb
        out_channels = round_filters(32, self._global_params)  # number of output channels
        self._conv_stem = Conv2dSamePadding(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
        self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)

        # Build blocks
        self._blocks = nn.ModuleList([])
        for block_args in self._blocks_args:

            # Update block input and output filters based on depth multiplier.
            block_args = block_args._replace(
                input_filters=round_filters(block_args.input_filters, self._global_params),
                output_filters=round_filters(block_args.output_filters, self._global_params),
                num_repeat=round_repeats(block_args.num_repeat, self._global_params)
            )

            # The first block needs to take care of stride and filter size increase.
            self._blocks.append(MBConvBlock(block_args, self._global_params))
            if block_args.num_repeat > 1:
                block_args = block_args._replace(input_filters=block_args.output_filters, stride=1)
            for _ in range(block_args.num_repeat - 1):
                self._blocks.append(MBConvBlock(block_args, self._global_params))

        # Head
        in_channels = block_args.output_filters  # output of final block
        out_channels = round_filters(1280, self._global_params)
        self._conv_head = Conv2dSamePadding(in_channels, out_channels, kernel_size=1, bias=False)
        self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)

        # Final linear layer
        self._dropout = self._global_params.dropout_rate
        self._fc = nn.Linear(out_channels, self._global_params.num_classes)

    def extract_features(self, inputs):
        """ Returns output of the final convolution layer """

        # Stem
        x = relu_fn(self._bn0(self._conv_stem(inputs)))

        # Blocks
        for idx, block in enumerate(self._blocks):
            drop_connect_rate = self._global_params.drop_connect_rate
            if drop_connect_rate:
                drop_connect_rate *= float(idx) / len(self._blocks)
            x = block(x) # , drop_connect_rate) # see https://github.com/tensorflow/tpu/issues/381

        return x

    def forward(self, inputs):
        """ Calls extract_features to extract features, applies final linear layer, and returns logits. """

        # Convolution layers
        x = self.extract_features(inputs)

        # Head
        x = relu_fn(self._bn1(self._conv_head(x)))
        x = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1)
        if self._dropout:
            x = F.dropout(x, p=self._dropout, training=self.training)
        x = self._fc(x)
        return x

    @classmethod
    def from_name(cls, model_name, override_params=None):
        cls._check_model_name_is_valid(model_name)
        blocks_args, global_params = get_model_params(model_name, override_params)
        return EfficientNet(blocks_args, global_params)

    @classmethod
    def from_pretrained(cls, model_name):
        model = EfficientNet.from_name(model_name)
        load_pretrained_weights(model, model_name)
        return model

    @classmethod
    def get_image_size(cls, model_name):
        cls._check_model_name_is_valid(model_name)
        _, _, res, _ = efficientnet_params(model_name)
        return res

    @classmethod
    def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=False):
        """ Validates model name. None that pretrained weights are only available for
        the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """
        num_models = 4 if also_need_pretrained_weights else 8
        valid_models = ['efficientnet_b'+str(i) for i in range(num_models)]
        if model_name.replace('-','_') not in valid_models:
            raise ValueError('model_name should be one of: ' + ', '.join(valid_models))

My TinySSD Model

'''
@Descripttion: This is Aoru Xue's demo,which is only for reference
@version:
@Author: Aoru Xue
@Date: 2019-06-14 00:42:10
@LastEditors: Aoru Xue
@LastEditTime: 2019-09-02 17:04:26
'''

import torch
from torch import nn
from efficientnet_pytorch import EfficientNet
from prior_box import PriorBox
from torchsummary import summary
import torch.nn.functional as F
from box_utils import *
from PIL import Image
class TinySSD(nn.Module):
    def __init__(self,training = True):
        super(TinySSD,self).__init__()
        self.basenet = EfficientNet.from_name('efficientnet-b0')
        self.training = training
        for idx,num_anchors in enumerate([4, 6, 6, 4, 4]):
            setattr(self,"predict_bbox_{}".format(idx + 1),nn.Conv2d(
                320,num_anchors * 4,kernel_size = 3,padding = 1
            ))
            setattr(self,"predict_class_{}".format(idx + 1),nn.Conv2d( # 这里3 是 2 + 1
                320,3 * num_anchors,kernel_size = 3,padding = 1
            ))
        self.priors = None
        for idx,k in enumerate([[320,320],[320,320],[320,320]]):
            setattr(self,"feature_{}".format(idx + 2),nn.Sequential(
                nn.Conv2d(k[0],k[1],kernel_size = 3,padding =1),
                nn.BatchNorm2d(k[1]),
                nn.ReLU(),
                nn.Conv2d(k[1],k[1],kernel_size = 3,padding =1),
                nn.BatchNorm2d(k[1]),
                nn.ReLU(),
                nn.MaxPool2d(2)
            ))
    def forward(self,x):
        x = self.basenet.extract_features(x)
        feature_1 = x
        feature_2 = self.feature_2(x)
        feature_3 = self.feature_3(feature_2)
        feature_4 = self.feature_4(feature_3)
        feature_5 = F.max_pool2d(feature_4,kernel_size = 2)

        '''
        (2,4*4,16,16)
        (2,4*6,8,8)
        (2,4*6,4,4),
        (2,4*4,2,2),
        (2,4*4,1,1)

        -> 每个 anchor 中心，连续4个值代表x y w h
        '''
        confidences = []
        locations = []
        locations.append(self.predict_bbox_1(feature_1).permute(0,2,3,1).contiguous())
        locations.append(self.predict_bbox_2(feature_2).permute(0,2,3,1).contiguous())
        locations.append(self.predict_bbox_3(feature_3).permute(0,2,3,1).contiguous())
        locations.append(self.predict_bbox_4(feature_4).permute(0,2,3,1).contiguous())
        locations.append(self.predict_bbox_5(feature_5).permute(0,2,3,1).contiguous())
        locations = torch.cat([o.view(o.size(0), -1) for o in locations], 1) #(batch_size,total_anchor_num*4)
        locations = locations.view(locations.size(0), -1, 4) # (batch_size,total_anchor_num,4)

        confidences.append(self.predict_class_1(feature_1).permute(0,2,3,1).contiguous())
        confidences.append(self.predict_class_2(feature_2).permute(0,2,3,1).contiguous())
        confidences.append(self.predict_class_3(feature_3).permute(0,2,3,1).contiguous())
        confidences.append(self.predict_class_4(feature_4).permute(0,2,3,1).contiguous())
        confidences.append(self.predict_class_5(feature_5).permute(0,2,3,1).contiguous())
        confidences = torch.cat([o.view(o.size(0), -1) for o in confidences], 1) #(batch_size,total_anchor_num*4)
        confidences = confidences.view(confidences.size(0), -1, 3) # (batch_size,total_anchor_num,4)
        if not self.training:
            if self.priors is None:
                self.priors = PriorBox()()
                self.priors = self.priors.cuda()
            boxes = convert_locations_to_boxes(
                locations, self.priors, 0.1, 0.2
            )
            confidences = F.softmax(confidences, dim=2)
            return confidences, boxes
        else:
            #print(confidences.size(),locations.size())
            return (confidences, locations) #  (2,1111,3) (2,1111,4)

if __name__ == "__main__":
    net = TinySSD()
    net.cuda()
    #prior = PriorBox()
    #print(len(prior()))
    #gt_prior = assign_priors(torch.Tensor([[0,0,10/512,10/512],[55/512,55/512,30/512,30/512]]),torch.Tensor([1,2,5]),prior(),0.5)
    #print(gt_prior[1])
    #x = torch.randn(1,3,512,512)
    #out = net(x.cuda())
    #print(out[0].size())
    #print(out[1].size())
    #print(prior()[:200,:])
    #print(out[0][0])
    #print(out[1][0])
    summary(net,(3,512,512),device="cuda")

dataset

'''
@Descripttion: This is Aoru Xue's demo,which is only for reference
@version:
@Author: Aoru Xue
@Date: 2019-06-15 12:48:09
@LastEditors: Aoru Xue
@LastEditTime: 2019-09-13 10:43:34
'''
import torch
import torch.nn
from torch.utils.data import Dataset
from PIL import Image
from prior_box import PriorBox
from box_utils import *
import cv2 as cv
import random
import numpy as np
import glob
import xml.etree.ElementTree as ET

class Mydataset(Dataset):
    def __init__(self,img_path = "./dataset",transform = None,center_variance = 0.1,size_variance = 0.2):
        self.center_variance = center_variance
        self.size_variance = size_variance
        self.img_paths = glob.glob(img_path + "/images/*.jpg")
        self.labels = [label.replace(".jpg",".xml").replace("images","labels") for label in self.img_paths]
        self.class_names = ("__background__","basketball","volleyball")
        prior = PriorBox()
        self.center_form_priors = prior() # center form
        self.imgW,self.imgH = 512,512
        self.corner_form_priors = center_form_to_corner_form(self.center_form_priors)
        #print(self.center_form_priors.size(),self.corner_form_priors.size())
        self.transform = transform
    def __len__(self):
        return len(self.img_paths)
    def __getitem__(self,idx):
        img = Image.open(self.img_paths[idx]).convert("RGB")
        label_file = self.labels[idx]
        gt_bboxes,gt_classes = self._get_annotation(idx)

        if self.transform:
            img = self.transform(img)

        gt_bboxes,gt_classes = assign_priors(gt_bboxes,gt_classes,self.corner_form_priors,0.5) # corner form
        #imH,imW = cv_img.shape[:2]

        gt_bboxes = corner_form_to_center_form(gt_bboxes) # (1524, 4) center form
        locations = convert_boxes_to_locations(gt_bboxes, self.center_form_priors, self.center_variance, self.size_variance) # 相当于归一化
        # 拟合距离而不是直接拟合，这样更容易拟合。
        return [img,locations,gt_classes]
    def _get_annotation(self,idx):
        annotation_file = self.labels[idx]
        objects = ET.parse(annotation_file).findall("object")
        boxes = []
        labels = []
        #is_difficult = []
        for obj in objects:
            class_name = obj.find('name').text.lower().strip()
            bbox = obj.find('bndbox')
            # VOC dataset format follows Matlab, in which indexes start from 0
            x1 = float(bbox.find('xmin').text) - 1
            y1 = float(bbox.find('ymin').text) - 1
            x2 = float(bbox.find('xmax').text) - 1
            y2 = float(bbox.find('ymax').text) - 1

            boxes.append([x1/self.imgW,y1/self.imgH,x2/self.imgW,y2/self.imgH])
            labels.append(self.class_names.index(class_name))
        return (torch.tensor(boxes, dtype=torch.float),
                torch.tensor(labels, dtype=torch.long))
if __name__ == '__main__':
    datset = Mydataset()
    import cv2 as cv
    img,gt_loc,gt_labels = datset[0]
    cv_img = np.array(img)
    cv_img = cv.cvtColor(cv_img,cv.COLOR_RGB2BGR)
    idx = gt_labels > 0
    #print(gt_loc.size(),dataset.priors.size())
    loc = convert_locations_to_boxes(gt_loc,datset.center_form_priors,0.1,0.2)
    loc = loc[idx]
    label = gt_labels[idx]
    for i in range(loc.size(0)):
        print(loc.size())
        x1,y1,w,h = loc[i,:]
        #print(x,y,r)
        x1 = x1.item() * 512.
        y1 = y1.item() * 512.
        w= w.item() * 512.
        h = h.item() * 512.
        #cv.circle(cv_img,(int(x),int(y)),int(r),(255,0,0),2)
        cv.rectangle(cv_img,(int(x1 - w/2),int(y1-h/2)),(int(x1 + w/2),int(y1 + h/2)),(255,0,0),2)
    cv.imshow("cv",cv_img)
    cv.waitKey(0)

训练

'''
@Descripttion: This is Aoru Xue's demo,which is only for reference
@version:
@Author: Aoru Xue
@Date: 2019-06-15 12:56:39
@LastEditors: Aoru Xue
@LastEditTime: 2019-09-10 20:46:54
'''
import torch
import torchvision
from TinySSD import TinySSD
#from vgg_ssd import build_ssd_model
from dataset import Mydataset
from torchvision import transforms
#from transforms import *
from torch.utils.data import DataLoader
from multibox_loss import MultiBoxLoss
import torch.optim as optim
from tqdm import tqdm
def train(dataloader,net,loss_fn,optimizer,epochs = 200):
    for epoch in range(epochs):
        running_loss_bbox = 0.
        running_loss_class = 0.
        for img,gt_bbox,gt_class in tqdm(dataloader):
            img = img.cuda()
            gt_bbox = gt_bbox.cuda()
            gt_class = gt_class.cuda()
            optimizer.zero_grad()
            pred_class,pred_locations = net(img)
            """Compute classification loss and smooth l1 loss.

                Args:
                    confidence (batch_size, num_priors, num_classes): class predictions.
                    predicted_locations (batch_size, num_priors, 4): predicted locations.
                    labels (batch_size, num_priors): real labels of all the priors.
                    gt_locations (batch_size, num_priors, 4): real boxes corresponding all the priors.
            """
            regression_loss, classification_loss = loss_fn(pred_class ,pred_locations,gt_class,gt_bbox)
            loss = regression_loss + classification_loss
            loss.backward()
            running_loss_bbox += regression_loss.item()
            running_loss_class += classification_loss.item()
            optimizer.step()
            #print(pred_bbox.size(),pred_class.size())

            #print("epoch: {},bbox loss:{:.8f} , class loss:{:.8f}".format(epoch + 1,loss[0].cpu().item(),loss[1].cpu().item()))
        print("*" * 20)
        print("average bbox loss: {:.8f}; average class loss: {:.8f}".format(running_loss_bbox/len(dataloader),running_loss_class/len(dataloader)))
        if epoch % 5 == 0:
            torch.save(net.state_dict(),"./ckpt/{}.pkl".format(epoch))
if __name__ == "__main__":
    net = TinySSD()
    net.cuda()
    loss_fn = MultiBoxLoss(3.)
    transform = transforms.Compose([
        transforms.Resize((512,512)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]
    )
    # transform = Compose([
    #     ConvertFromInts(),
    #     PhotometricDistort(),
    #     Expand([123, 117, 104]),
    #     RandomSampleCrop(),
    #     RandomMirror(),
    #     ToPercentCoords(),
    #     Resize(300),
    #     SubtractMeans([123, 117, 104]),
    #     ToTensor(),
    # ])

    optm = optim.Adam(net.parameters(),lr = 1e-3)
    dtset = Mydataset(img_path = "./dataset",transform = transform)
    dataloader = DataLoader(dtset,batch_size = 8,shuffle = True)

    train(dataloader,net,loss_fn,optm)