Keypoint detection with r-cnn feature extraction backnone

nicholasprimiano: <p>I’m training a keypoint detection model using the builtin pytorch r-cnn class. It requires a backbone feature extraction network. I got decent results using efficientnet and convnext backbones but would like to try other architectures like one of the bulitin vision transformers.  The model works when I access the efficientnet or convnext “.features” attribute. If I understand it correctly this attribute accesses the network without the top/classification layer. I manged to access this layer of the transfromer using pytorch’s feature_extraction method but I am not able to use it as a backbone (I get inccret dimension errors that change each time I run it). I realize this post is light in details. Please let me know what additional information would be helpful. Any thoughts greatly appreciated.</p>

ptrblck: <p>Check the shape of the output features for the feature extractor models in the working configs and then compare it to the feature shape of your failing use case.<br>
Based on your description it seems that the current model outputs features which cause shape mismatches and which also change. If so, you might need to use adaptive pooling layers to make sure the features have a static shape.</p>

nicholasprimiano: <p>Here’s the structure of the model I’m using with vit_b_16 as the backbone to the rcnn:</p>
<p>def get_model(num_keypoints):</p>
<pre><code>#backbone = torchvision.models.convnext_large(weights='DEFAULT').features
#backbone.out_channels = 1536 # 1536 for convnext_large

model = torchvision.models.vit_b_16(weights='DEFAULT')
backbone = nn.Sequential(*list(model.children())[:-1])
backbone.out_channels = 768 

anchor_generator = AnchorGenerator(sizes=((64, 128, 256),), aspect_ratios=((0.83, 1.7, 1.2),)) # these need to be tweaked for the very large dataset

roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
                                                output_size=7,
                                                sampling_ratio=2)
keypoint_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
                                                         output_size=14,
                                                         sampling_ratio=2)

model = KeypointRCNN(backbone,
                      num_classes=2,
                      rpn_anchor_generator=anchor_generator,
                      box_roi_pool=roi_pooler,
                      keypoint_roi_pool=keypoint_roi_pooler,
                      num_keypoints=num_keypoints)
       
    
return model
</code></pre>
<p>When I train the model I get this output. Input images are size (224,224,3) batch size 2.</p>
<p>AssertionError: Expected (batch_size, seq_length, hidden_dim) got torch.Size([2, 768, 46, 46])</p>

ptrblck: <p>Thanks for the code snippet.<br>
Based on your code I see this error:</p>
<pre><code class="lang-python">model.eval()
x = torch.randn(2, 3, 224, 224)
out = model(x)
# AssertionError: Expected (batch_size, seq_length, hidden_dim) got torch.Size([2, 768, 50, 50])
</code></pre>
<p>which I guess might come from rewrapping the <code>VisionTransformer</code> model into an <code>nn.Sequential</code> container.<br>
This step would remove all functional API calls from the original <code>forward</code> method which is defined <a href="https://github.com/pytorch/vision/blob/c206a471617e41ba04a0f3cc5d926a4b7c391afe/torchvision/models/vision_transformer.py#L289-L305">here</a>.<br>
The <code>assert</code> is raised in the <code>Encoder</code> class <a href="https://github.com/pytorch/vision/blob/c206a471617e41ba04a0f3cc5d926a4b7c391afe/torchvision/models/vision_transformer.py#L155">here</a> which seems to be missing the <code>reshape</code> (and <code>permute</code>) from the <code>_process_input</code> call as well as the <code>torch.cat</code> call from <a href="https://github.com/pytorch/vision/blob/c206a471617e41ba04a0f3cc5d926a4b7c391afe/torchvision/models/vision_transformer.py#L268-L296">here</a>.</p>

nicholasprimiano: <p>How can I access the feature extraction or last layer (which KeypointRCNN() requires as a backbone) backbone without wrapping the layers in nn.Sequential?</p>

ptrblck: <p>Maybe you could use a forward hook or <code>torchvision</code>’s <code>create_feature_extrator</code> method as described in <a href="https://pytorch.org/vision/stable/feature_extraction.html">this tutorial</a>.</p>

nicholasprimiano: <p>Thanks for the response. I tried to configure it with the feature extractor. Here is the traceback I get:<br>
AssertionError  Cell In[65], line 8<br>
5 lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.3)<br>
7 for epoch in range(num_epochs):<br>
----&gt; 8     train_one_epoch(model, optimizer, data_loader_train, device, epoch, print_freq=100)<br>
9     lr_scheduler.step()<br>
10     evaluate(model, data_loader_test, device)</p>
<p>File /mnt/f/patella/engine.py:31, in train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler)<br>
29 targets = [{k: v.to(device) for k, v in t.items()} for t in targets]<br>
30 with torch.cuda.amp.autocast(enabled=scaler is not None):<br>
—&gt; 31     loss_dict = model(images, targets)<br>
32     losses = sum(loss for loss in loss_dict.values())<br>
34 # reduce losses over all GPUs for logging purposes</p>
<p>File ~/anaconda3/envs/insall_salvati/lib/python3.10/site-packages/torch/nn/modules/module.py:1488, in Module._call_impl(self, *args, **kwargs)<br>
1483 # If we don’t have any hooks, we want to skip the rest of the logic in<br>
1484 # this function, and just call forward.<br>
1485 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks<br>
1486         or _global_backward_pre_hooks or _global_backward_hooks<br>
1487         or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1488     return forward_call(*args, **kwargs)<br>
1489 # Do not call functions when jit is used<br>
1490 full_backward_hooks, non_full_backward_hooks = [], []<br>
…<br>
1126 if type(condition) is not torch.Tensor and has_torch_function((condition,)):<br>
1127     return handle_torch_function(_assert, (condition,), condition, message)<br>
 → 1128 assert condition, message</p>
<p>AssertionError: Wrong image height! Expected 224 but got Proxy(getitem_2)!</p>
<p>I get the same AssertionError if create the feature extractor with all nodes (create_feature_extractor(network, eval_return_nodes=eval_nodes[:], train_return_nodes=train_nodes[:]) or with only the last layer (create_feature_extractor(network, return_nodes =[‘getitem_5’]).</p>