【浅显易懂系列】DETR详细解释，一看就懂也！

2 minute read

Published: September 01, 2024

DETR（Detection Transformer）是一种基于Transformer的目标检测模型，它通过端到端的训练来直接预测目标的位置和类别。下面，我将选取DETR模型中的核心部分，使用PyTorch进行逐行解释，并注释其维度。

1. Transformer Encoder 部分

Transformer Encoder通常用于处理输入的图像特征（通常是CNN特征图经过flatten和线性变换后的序列）。

class TransformerEncoder(nn.Module):
    def __init__(self, d_model, nhead, num_encoder_layers, dim_feedforward=2048):
        super(TransformerEncoder, self).__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)

    def forward(self, src):
        # src: [batch_size, src_seq_len, d_model]
        # 输出: 经过Transformer Encoder处理的特征
        output = self.transformer_encoder(src)
        return output

2. Transformer Decoder 部分

Transformer Decoder用于生成目标检测结果，通常结合位置编码和目标查询（object queries）。

class TransformerDecoder(nn.Module):
    def __init__(self, d_model, nhead, num_decoder_layers, dim_feedforward=2048):
        super(TransformerDecoder, self).__init__()
        decoder_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward)
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
        # tgt: [batch_size, tgt_seq_len, d_model] 目标查询
        # memory: [batch_size, src_seq_len, d_model] Encoder的输出
        # 输出: 经过Transformer Decoder处理的目标查询
        output = self.transformer_decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
                                          tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
        return output

3. DETR 主体部分

DETR模型将Encoder和Decoder结合起来，并添加预测头以输出检测结果。

class DETR(nn.Module):
    def __init__(self, num_classes, hidden_dim, nheads, num_encoder_layers, num_decoder_layers, num_queries):
        super(DETR, self).__init__()
        # Encoder
        self.transformer = TransformerEncoder(d_model=hidden_dim, nhead=nheads, num_encoder_layers=num_encoder_layers)
        
        # Decoder
        self.transformer_decoder = TransformerDecoder(d_model=hidden_dim, nhead=nheads, num_decoder_layers=num_decoder_layers)
        
        # Object queries
        self.query_embed = nn.Embedding(num_queries, hidden_dim)
        
        # Prediction heads
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)  # +1 for background
        self.bbox_embed = MLP(hidden_dim, hidden_dim * 4, 4, 3)

    def forward(self, images):
        # 假设images已经通过CNN处理成[batch_size, num_channels, H, W]
        # 这里略过CNN部分，直接模拟CNN输出特征
        src = ...  # [batch_size, src_seq_len, hidden_dim]，其中src_seq_len是特征图的序列长度
        
        # 目标查询
        hs = self.query_embed.weight.unsqueeze(0).repeat(images.shape[0], 1, 1)  # [batch_size, num_queries, hidden_dim]
        
        # Encoder
        memory = self.transformer(src)
        
        # Decoder
        tgt = torch.zeros_like(hs)  # 初始化为零的目标查询
        outputs = self.transformer_decoder(tgt, memory, tgt_mask=None, memory_key_padding_mask=None)
        
        # 输出预测
        outputs_class = self.class_embed(outputs)
        outputs_coord = self.bbox_embed(outputs).sigmoid()  # 假设使用sigmoid来限制坐标范围
        
        return outputs_class, outputs_coord

当然，为了更全面地理解DETR（Detection Transformer）模型，我们需要包括CNN特征提取器和位置编码等关键组件的细节。以下是DETR模型的一个更详细的解释，包括这些组件：

4. CNN特征提取器

DETR模型通常使用一个预训练的卷积神经网络（CNN）作为特征提取器，将输入图像转换为特征图。这个CNN可以是ResNet、VGG或其他任何适用于图像识别的网络结构。

class Backbone(nn.Module):
    def __init__(self, name='resnet50', pretrained=True):
        super(Backbone, self).__init__()
        # 这里以ResNet50为例，使用torchvision中的预训练模型
        if name == 'resnet50':
            self.resnet = torchvision.models.resnet50(pretrained=pretrained)
            # 通常我们只使用到resnet的某个层级的特征图，例如layer4的输出
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # 调整为1x1以获取全局特征

        # 移除全连接层和其他不必要的层
        # ...

    def forward(self, x):
        # x: [batch_size, 3, height, width] 输入图像
        # 假设我们只使用resnet的最后一个stage的特征
        features = self.resnet.layer4(self.resnet.relu(self.resnet.bn1(self.resnet.conv1(x))))
        # ... 可能还有其他层需要处理
        features = self.avgpool(features)  # 池化到1x1
        features = features.view(features.shape[0], -1, features.shape[-1])  # 展平为序列
        # 可能还需要添加一个线性层来调整特征维度以匹配Transformer的输入
        features = self.some_linear_layer(features)  # 假设的线性层
        return features

5. 位置编码

由于Transformer模型本身不包含关于序列中元素位置的信息，因此在处理图像或序列数据时，通常需要添加位置编码（Positional Encoding）。在DETR中，位置编码通常与CNN特征图相结合，为模型提供空间位置信息。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        # 初始化位置编码
        # ...

    def forward(self, x):
        # x: [batch_size, seq_len, d_model] 输入特征
        # 应用位置编码
        # 通常这里会计算一个与x维度相同的矩阵，并将其加到x上
        # ...
        return x + self.positional_encodings[:, :x.size(1)]

然而，在DETR的原始论文和许多实现中，位置编码不是显式添加到CNN特征上的，而是通过目标查询（object queries）或Transformer解码器的自注意力机制隐式地引入位置信息。这是因为目标查询在解码器中通过自注意力与其他查询和编码器输出交互，从而学习到目标的相对位置。

6. DETR 主体部分（包含CNN和位置编码的假设）

class DETR(nn.Module):
    def __init__(self, ...):
        # ... 之前的初始化代码
        self.backbone = Backbone('resnet50')  # 使用ResNet50作为特征提取器
        # 可能不需要显式的位置编码，因为DETR通过其他方式处理位置

    def forward(self, images):
        # 提取CNN特征
        features = self.backbone(images)  # [batch_size, num_channels * spatial_dim, hidden_dim]
        # 假设features已经被调整为适合Transformer的序列长度

        # 这里可以假设位置编码已经通过目标查询或其他方式隐式处理

        # Encoder
        memory = self.transformer(features)  # 可能需要调整features的维度以匹配Transformer的输入

        # Decoder
        # ... 与之前的Decoder部分相同

        # 输出预测
        # ...

        return outputs_class, outputs_coord

请注意，上述代码中的PositionalEncoding类在实际DETR实现中可能不是必需的，因为DETR通过其他机制（如目标查询和自注意力）来隐式地处理位置信息。此外，Backbone类的实现也高度依赖于所使用的具体CNN架构和DETR模型的特定要求。

如此，你看懂了吗？已经够详细了吧？？？

Share on

Twitter Facebook LinkedIn

Yifei Cao