Get ROI feature vector from Faster R-CNN using Pytorch

Question

I am training a faster R-CNN model in pytorch and I want to extract feature vector from roi-heads layer.

I am using the following code:

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 9  # 1 class (wheat) + background

# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features

# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

features = []
def save_features(model, inp, out):
    features.append(out[0].data)

# you can also hook layers inside the roi_heads
layer_to_hook = 'roi_heads'
for name, layer in model.named_modules():
    if name == layer_to_hook:
        layer.register_forward_hook(save_features)

But during training I am getting the following error:

AttributeError                            Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_11556/1891582782.py in <module>
     14 for epoch in range(num_epochs):
     15 
---> 16     train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq = 10)
     17     lr_scheduler.step()
     18     # Evaluate on the test dataset

~\pytorch_custom_object_detection\engine.py in train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq)
     28         targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
     29 
---> 30         loss_dict = model(images, targets)
     31 
     32         losses = sum(loss for loss in loss_dict.values())

~\anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~\anaconda3\lib\site-packages\torchvision\models\detection\generalized_rcnn.py in forward(self, images, targets)
     95             features = OrderedDict([('0', features)])
     96         proposals, proposal_losses = self.rpn(images, features, targets)
---> 97         detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
     98         detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)
     99 

~\anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1121         if _global_forward_hooks or self._forward_hooks:
   1122             for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):
-> 1123                 hook_result = hook(self, input, result)
   1124                 if hook_result is not None:
   1125                     result = hook_result

~\AppData\Local\Temp/ipykernel_11556/181024080.py in save_features(model, inp, out)
     53     def save_features(model, inp, out):
     54         print(out)
---> 55         features.append(out.data)
     56 
     57     # you can also hook layers inside the roi_heads

AttributeError: 'tuple' object has no attribute 'data'

I want to know what is the going on and is there any other way I can do this. Thanks in advance.