Unlocking the Power of DETR for Custom Image-to-Set Prediction Tasks

The DEtection TRansformer (DETR) has gained significant attention for its ability to reformulate object detection as a direct set prediction problem. While DETR is widely recognized for its performance in object detection, its underlying architecture makes it suitable for a wide range of set prediction tasks. In this blog post, we’ll explore how to adapt DETR for custom tasks beyond object detection.

Why DETR for Set Prediction?

DETR’s strength lies in its transformer-based design, which enables it to model global relationships in the input data. Its set-based loss eliminates the need for heuristics like non-maximum suppression (NMS), making it a versatile tool for predicting unordered sets of outputs.

DETR Architecture Components

In their original paper End-to-End Object Detection with Transformers, Carion et al. provided easy-to-use PyTorch code for implementing DETR. We will use this code as a starting point for adapting DETR to custom tasks.

class DETR(nn.Module):
    def __init__(self, num_classes: int, hidden_dim: int, nheads: int,
                 num_encoder_layers: int, num_decoder_layers: int):
        super().__init__()
        ...

Backbone

In the provided DETR implementation, the backbone is a ResNet-50 model with weights pretrained on the ImageNet dataset:

self.backbone = nn.Sequential(*list(resnet50(weights=ResNet50_Weights.DEFAULT).children())[:-2])

While the ResNet model family is a good starting point, you may replace it with a different backbone architecture (e.g., Vision Transformer (ViT)) or use a pretrained model better suited to your task.

To reduce the computational cost of the transformer, a 1x1 convolution is applied to the output of the backbone to match the transformer's hidden dimension:

self.conv = nn.Conv2d(2048, hidden_dim, 1)

If you choose a different backbone, ensure the input and output channel dimensions of this convolution are compatible with the backbone.

Transformer

The core of DETR is its transformer model. In the implementation, a standard PyTorch transformer is used:

self.transformer = nn.Transformer(hidden_dim, nheads, num_encoder_layers, num_decoder_layers, batch_first=True)

You can experiment with alternative transformer architectures, such as deformable attention, which may improve performance on specific tasks.

Query Vectors

The number of query vectors used by the transformer decoder should exceed the maximum number of elements in the output set (e.g., the maximum number of objects in an image). In DETR, the query vectors are learned parameters:

self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))

Positional Encodings

The provided PyTorch implementation of DETR uses learned embeddings for the positional encodings of the input image. These embeddings are computed separately for rows and columns. Later, in the forward pass, the row and column embeddings are concatenated, resulting in a consistent vector of dimensionality hidden_dim. The number of positional embeddings (in this case 50) acts as an upper bound on the number of patches per row and column in the input image:

self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))

You can also experiment with fixed 2D sinusoidal positional encodings, such as those provided here.

Prediction Heads

The prediction head of the model operates on each transformed query vector and predicts an element of the resulting set. In the case of object detection, two linear layers are used to predict the bounding box and class label of the object. A <no class> label is added to the label set to account for the absence of objects:

self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
self.linear_bbox = nn.Linear(hidden_dim, 4)

For custom tasks, you may need to adapt the prediction heads. Depending on the elements of your set, you might require only a single prediction head or even more than two. You can also increase complexity by using multi-layer perceptrons (MLPs) or other architectures.

Forward Pass Implementation

Once the components are defined, implementing the forward pass is straightforward. First, the input image is passed through the backbone and the convolutional layer to obtain feature maps with a fixed channel dimension:

def forward(self, x: torch.Tensor) -> torch.Tensor:
    x = self.backbone(x)
    h = self.conv(x)

Next, the 2D positional embeddings are constructed by concatenating the row and column embeddings:

H, W = h.shape[-2:]
pos = torch.cat([
        self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1),
        self.row_embed[:H].unsqueeze(1).repeat(1, W, 1)
    ], dim=-1)

The feature maps and positional embeddings are then flattened across the spatial dimensions and added together. Along with the query vectors, they are passed through the transformer:

h = h.flatten(2).permute(0, 2, 1)
pos = pos.flatten(0, 1).unsqueeze(0)
query = self.query_pos.unsqueeze(0)
h = self.transformer(pos + h, query)

Finally, the output of the transformer is passed through the prediction heads. A sigmoid function is applied after the bounding box prediction head to ensure the bounding box coordinates are in the range [0, 1]:

return self.linear_class(h), self.linear_bbox(h).sigmoid()

Acknowledgment

This blog post was supported by the Konrad Zuse School of Excellence in Learning and Intelligent Systems (ELIZA). I am grateful for their support in advancing my research and science communication efforts.