pytroch复现ssd

### 使用 PyTorch 实现 SSD 模型 #### 构建 VGG 基础网络 SSD 模型采用 VGG16 作为基础模型，并在此基础上进行了修改。为了适应不同尺度的目标检测需求，SSD 输入图像尺寸通常设置为 300×300 或者 512×512[^1]。 ```python import torch.nn as nn from torchvision import models class BaseVGG(nn.Module): def __init__(self): super(BaseVGG, self).__init__() vgg = models.vgg16(pretrained=True).features # 修改vgg最后几层以适配SSD的需求 self.base_layers = list(vgg.children())[:-8] def forward(self, x): for layer in self.base_layers: x = layer(x) return x ``` #### 添加额外卷积层 除了继承自 VGG 的卷积层外，还需要增加多个新的卷积层来提取更丰富的特征图。这些附加的卷积层有助于捕捉多尺度的信息，从而提高小物体检测的效果。 ```python def make_extra_layers(): layers = [] # 定义各层参数 cfg = [ (256, 'S', 512), (128, 'S', 256), (128, '', 256), (128, '', 256)] padding = {'': 1, 'S': 0} stride = {'': 1, 'S': 2} input_channels = 1024 for k, v in enumerate(cfg): out_channels = v[-1] conv_layer = nn.Conv2d(input_channels, out_channels, kernel_size=(1 if isinstance(k % 2 == 0) else 3), stride=stride[v[1]], padding=padding[v[1]]) layers += [conv_layer, nn.ReLU(inplace=True)] input_channels = out_channels return nn.Sequential(*layers) extra_conv_layers = make_extra_layers() ``` #### 创建预测头（Prediction Heads） 对于每一个特征映射位置，都需要为其分配一定数量的默认边界框(anchor boxes)，并由相应的分类器和回归器来进行最终预测。这里我们定义了一个简单的函数用于创建这些头部组件[^2]。 ```python def create_multibox(base_model, extra_convs): loc_layers = [] conf_layers = [] mbox = [4, 6, 6, 6, 4, 4] # 各层次feature map对应的default box数目 base_output_channels = 512 extras_output_channels = [out_ch for _,_,out_ch in extra_convs][:-1] sources = [*base_model, *extras_output_channels] for idx, num_boxes in enumerate(mbox): loc_layers.append( nn.Conv2d(sources[idx], num_boxes*4, kernel_size=3, padding=1)) conf_layers.append( nn.Conv2d(sources[idx], num_boxes*num_classes, kernel_size=3, padding=1)) return (nn.ModuleList(loc_layers), nn.ModuleList(conf_layers)) ``` #### 组合所有模块完成整个SSD架构构建 通过组合上述各个部分可以得到完整的SSD模型实例化对象net。当给定输入张量`x`时，可以通过调用该对象的方法获取到两个列表形式的结果——分别是定位信息(`featuremap_locs`)以及类别概率分布(`featuremap_classes`)。 ```python class SSD(nn.Module): def __init__(self, phase='train'): super(SSD, self).__init__() self.phase = phase self.num_classes = 21 # VOC数据集共有20类目标加背景共21类 self.vgg = BaseVGG() self.extras = make_extra_layers() self.loc, self.conf = create_multibox([self.vgg]+list(self.extras), [(256,'S',512),(128,'S',256)]) def forward(self, x): sources = [self.vgg(x)] for k,v in enumerate(self.extras): x = F.relu(v(x), inplace=True) if k%2==1: sources.append(x) featuremap_locs = [] featuremap_classes = [] for (x,l,c) in zip(sources,self.loc,self.conf): featuremap_locs.append(l(x).permute(0,2,3,1).contiguous()) featuremap_classes.append(c(x).permute(0,2,3,1).contiguous()) output_loc = torch.cat([o.view(o.size(0),-1) for o in featuremap_locs],dim=1) output_conf = torch.cat([o.view(o.size(0),-1) for o in featuremap_classes],dim=1) if self.phase=='test': from data.config import voc_config as config priors = Variable(torch.from_numpy(config['priors']).type(type(x.data)), volatile=True) decoded_boxes = decode(output_loc,priors,config['variance']) scores = F.softmax(output_conf,dim=-1) return detected_objects(decoded_boxes,scores,num_classes=self.num_classes,top_k=config['top_k'],conf_thresh=config['confidence_threshold']) elif self.phase=='train': return output_loc.view(output_loc.size(0), -1, 4), \ output_conf.view(output_conf.size(0), -1, self.num_classes) if __name__ == '__main__': net = SSD(phase='eval') print(net.vgg) x = torch.randn(2, 3, 300, 300) featuremap_locs, featuremap_classes = net(x) for i in featuremap_classes: print(i.shape) print('..................') for j in featuremap_locs: print(j.shape) ```