SSD单发多框检测——模型及其损失函数

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import cv2
import numpy as np

class VGGBase(nn.Module):
    def __init__(self):
        super(VGGBase, self).__init__()
        vgg = torchvision.models.vgg16(pretrained=True)
        self.features = nn.Sequential(*list(vgg.features)[:-2])  # 移除最后两个池化层

    def forward(self, x):
    	#只使用特征层进行预测，没有使用它的线性层
        x = self.features(x)
        return x

#在可以可以添加一些额外的卷积层，增强提取特征的能力

class PredictionLayers(nn.Module):
    def __init__(self, num_classes):
        super(PredictionLayers, self).__init__()
        self.num_classes = num_classes
        self.loc_layers = nn.ModuleList()
        self.conf_layers = nn.ModuleList()

        self.loc_layers.append(nn.Conv2d(512, 4 * 4, kernel_size=3, padding=1))
        self.conf_layers.append(nn.Conv2d(512, 4 * num_classes, kernel_size=3, padding=1))

        self.loc_layers.append(nn.Conv2d(1024, 6 * 4, kernel_size=3, padding=1))
        self.conf_layers.append(nn.Conv2d(1024, 6 * num_classes, kernel_size=3, padding=1))

        self.loc_layers.append(nn.Conv2d(512, 6 * 4, kernel_size=3, padding=1))
        self.conf_layers.append(nn.Conv