Upload model

configuration_retinanet.py

@@ -1,6 +1,7 @@
 from transformers.configuration_utils import PretrainedConfig
 from optimum.exporters.onnx.model_configs import ViTOnnxConfig
-from typing import Optional, Dict
+from optimum.utils import DummyVisionInputGenerator
+from typing import Optional, Dict, OrderedDict
 
 class RetinaNetConfig(PretrainedConfig):
     model_type = 'retinanet'
@@ -19,23 +20,49 @@ class RetinaNetConfig(PretrainedConfig):
        
         super().__init__(**kwargs)
 
+class RetinaNetObjectDetectionInputGenerator(DummyVisionInputGenerator):
+
+    SUPPORTED_INPUT_NAMES = (
+        "pixel_values",
+        "image_sizes"
+    )
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        if input_name == "image_sizes":
+            return self.random_int_tensor(
+                shape=[self.batch_size, 2],
+                min_value=1,
+                max_value=max(self.height, self.width),
+                framework=framework,
+                dtype=int_dtype,
+            )
+        
+        elif input_name == "pixel_values":
+            return self.random_float_tensor(
+                shape=[self.batch_size, self.num_channels, self.height, self.width],
+                framework=framework,
+                dtype=float_dtype,
+            )
+
 class RetinaNetOnnxConfig(ViTOnnxConfig):
+    DUMMY_INPUT_GENERATOR_CLASSES = (RetinaNetObjectDetectionInputGenerator,)
+
     @property
     def inputs(self) -> Dict[str, Dict[int, str]]:
         return {
             "pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"},
-            "image_sizes": {0: "batch_size", 1: "image_size"}
         }
 
     @property
     def outputs(self) -> Dict[str, Dict[int, str]]:
-        common_outputs = super().outputs 
-
-        if self.task == "object-detection":
-            common_outputs["logits"] = {0: "batch_size", 1: "num_queries", 2: "num_classes"}
-            common_outputs["pred_boxes"] = {0: "batch_size", 1: "num_queries", 2: "coordinates"}
         
-        return common_outputs
+        return OrderedDict(
+            {
+                "boxes": {0: "batch_size", 1: "num_predictions", 2: "bbox_coordinates"},
+                "labels": {0: "batch_size", 1: "num_predictions"},
+                "scores": {0: "batch_size", 1: "num_predictions"},
+            }
+        )
     
 __all__ = [
     'RetinaNetConfig',

modeling_retinanet.py

@@ -1,104 +1,126 @@
-import torch
-from dataclasses import dataclass
-from torchvision.models import ResNet50_Weights
-from torchvision.models.detection import retinanet_resnet50_fpn, RetinaNet_ResNet50_FPN_Weights
-from torchvision.models.detection.anchor_utils import AnchorGenerator
-
-from transformers import PreTrainedModel
-from transformers.utils import ModelOutput
-from typing import OrderedDict, List, Tuple
-
-from .configuration_retinanet import RetinaNetConfig
-
-def _default_anchorgen():
-        anchor_sizes = tuple((x, int(x * 2 ** (1.0 / 3)), int(x * 2 ** (2.0 / 3))) for x in [32, 64, 128, 256, 512])
-        aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
-        anchor_generator = RetinaNetAnchorGenerator(anchor_sizes, aspect_ratios)
-        return anchor_generator
-
-@dataclass
-class RetinaNetObjectDetectionOutput(ModelOutput):
-    logits: torch.FloatTensor = None
-    pred_boxes: torch.FloatTensor = None
-    image_sizes: List[Tuple] = None
-    anchors: List[torch.Tensor] = None
-    features: List[torch.Tensor] = None
-
-class RetinaNetAnchorGenerator(AnchorGenerator):
-    def __init__(
-            self,
-            sizes=((128, 256, 512),),
-            aspect_ratios=((0.5, 1.0, 2.0),)
-        ):
-        super().__init__(sizes, aspect_ratios)
-
-    def forward(self, pixel_values: torch.Tensor, feature_maps: List[torch.Tensor]) -> List[torch.Tensor]:
-        grid_sizes = [feature_map.shape[-2:] for feature_map in feature_maps]
-        image_size = pixel_values.shape[-2:]
-        dtype, device = feature_maps[0].dtype, feature_maps[0].device
-        strides = [
-            [
-                torch.empty((), dtype=torch.int64, device=device).fill_(image_size[0] // g[0]),
-                torch.empty((), dtype=torch.int64, device=device).fill_(image_size[1] // g[1]),
-            ]
-            for g in grid_sizes
-        ]
-        self.set_cell_anchors(dtype, device)
-        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides)
-        anchors: List[List[torch.Tensor]] = []
-        for _ in range(pixel_values.shape[0]):
-            anchors_in_image = [anchors_per_feature_map for anchors_per_feature_map in anchors_over_all_feature_maps]
-            anchors.append(anchors_in_image)
-        anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors]
-        return anchors
-
-class RetinaNetModelForObjectDetection(PreTrainedModel):
-    config_class = RetinaNetConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.config = config
-        
-        model_config = {
-            'weights': None,
-            'weights_backbone': None,
-            'num_classes': None
-        }
-
-        if config.pretrained:
-            model_config['weights'] = RetinaNet_ResNet50_FPN_Weights.DEFAULT
-        else:
-            model_config['num_classes'] = config.num_classes
-            if config.pretrained_backbone:
-                model_config['weights_backbone'] = ResNet50_Weights.DEFAULT
-           
-        self.model = retinanet_resnet50_fpn(**model_config)
-        self.model.anchor_generator = _default_anchorgen()
-
-    def forward(self, pixel_values, image_sizes, labels=None):
-        if labels is not None:
-            raise NotImplementedError
-
-        features = self.model.backbone(pixel_values)
-        if isinstance(features, torch.Tensor):
-            features = OrderedDict([("0", features)])
-        features = list(features.values())
-
-        # compute the retinanet heads outputs using the features
-        head_outputs = self.model.head(features)
-
-        # create the set of anchors
-        anchors = self.model.anchor_generator(pixel_values, features)
-
-        return RetinaNetObjectDetectionOutput(
-            logits=head_outputs['cls_logits'],
-            pred_boxes=head_outputs['bbox_regression'],
-            image_sizes=image_sizes,
-            anchors=anchors, 
-            features=features
-        )
-    
-__all__ = [
-    "RetinaNetModelForObjectDetection"
+import torch
+from dataclasses import dataclass
+from torchvision.models import ResNet50_Weights
+from torchvision.models.detection import retinanet_resnet50_fpn, RetinaNet_ResNet50_FPN_Weights
+from torchvision.models.detection.anchor_utils import AnchorGenerator
+
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+from typing import OrderedDict, List, Union
+
+from configuration_retinanet import RetinaNetConfig
+
+def _default_anchorgen():
+        anchor_sizes = tuple((x, int(x * 2 ** (1.0 / 3)), int(x * 2 ** (2.0 / 3))) for x in [32, 64, 128, 256, 512])
+        aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
+        anchor_generator = RetinaNetAnchorGenerator(anchor_sizes, aspect_ratios)
+        return anchor_generator
+
+@dataclass
+class RetinaNetObjectDetectionOutput(ModelOutput):
+    logits: torch.FloatTensor = None
+    pred_boxes: torch.FloatTensor = None
+    image_sizes: torch.Tensor = None
+    anchors: torch.Tensor = None
+    num_anchors_per_level: torch.Tensor = None
+
+class RetinaNetAnchorGenerator(AnchorGenerator):
+    def __init__(
+            self,
+            sizes=((128, 256, 512),),
+            aspect_ratios=((0.5, 1.0, 2.0),)
+        ):
+        super().__init__(sizes, aspect_ratios)
+
+    def forward(self, pixel_values: torch.Tensor, feature_maps: List[torch.Tensor]) -> List[torch.Tensor]:
+        grid_sizes = [feature_map.shape[-2:] for feature_map in feature_maps]
+        image_size = pixel_values.shape[-2:]
+        dtype, device = feature_maps[0].dtype, feature_maps[0].device
+        strides = [
+            [
+                torch.empty((), dtype=torch.int64, device=device).fill_(image_size[0] // g[0]),
+                torch.empty((), dtype=torch.int64, device=device).fill_(image_size[1] // g[1]),
+            ]
+            for g in grid_sizes
+        ]
+        self.set_cell_anchors(dtype, device)
+        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides)
+        anchors: List[List[torch.Tensor]] = []
+        for _ in range(pixel_values.shape[0]):
+            anchors_in_image = [anchors_per_feature_map for anchors_per_feature_map in anchors_over_all_feature_maps]
+            anchors.append(anchors_in_image)
+        anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors]
+        return anchors
+
+class RetinaNetModelForObjectDetection(PreTrainedModel):
+    config_class = RetinaNetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        
+        model_config = {
+            'weights': None,
+            'weights_backbone': None,
+            'num_classes': None
+        }
+
+        if config.pretrained:
+            model_config['weights'] = RetinaNet_ResNet50_FPN_Weights.DEFAULT
+        else:
+            model_config['num_classes'] = config.num_classes
+            if config.pretrained_backbone:
+                model_config['weights_backbone'] = ResNet50_Weights.DEFAULT
+           
+        self.model = retinanet_resnet50_fpn(**model_config)
+
+
+    def forward_without_processing(self, pixel_values, image_sizes=None, labels=None):
+        if labels is not None:
+            raise NotImplementedError
+        else:
+            self.model.training = False
+        
+        if image_sizes is None:
+            # construct a tensor [batchsize,2] of value pixel_values.shape[-2:]
+            raise NotImplementedError
+
+        features = self.model.backbone(pixel_values)
+        if isinstance(features, torch.Tensor):
+            features = OrderedDict([("0", features)])
+        features = list(features.values())
+
+        # compute the retinanet heads outputs using the features
+        head_outputs = self.model.head(features)
+
+        # create the set of anchors
+        self.model.anchor_generator = _default_anchorgen()
+        anchors = self.model.anchor_generator(pixel_values, features)
+        num_anchors_per_level = [x.size(2) * x.size(3) for x in features]
+
+        return RetinaNetObjectDetectionOutput(
+            logits=head_outputs['cls_logits'],
+            pred_boxes=head_outputs['bbox_regression'],
+            image_sizes=image_sizes,
+            anchors=torch.stack(anchors, dim=0), 
+            num_anchors_per_level=torch.tensor(num_anchors_per_level)
+        )
+    
+
+    def forward(self, pixel_values: Union[torch.Tensor, List[torch.Tensor]], labels=None):
+        """
+        Don't use preprocessor for calling the main forward function.
+        """
+        if labels is not None:
+            raise NotImplementedError
+        else:
+            self.model.training = False
+        
+        detections = self.model(pixel_values, labels)
+
+        return detections
+    
+__all__ = [
+    "RetinaNetModelForObjectDetection"
 ]