Repository: pytorch
Issue: Export torchvision detection model retinanet_resnet50_fpn
Labels: triaged,oncall: pt2,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

This model export was requested by user at https://github.com/pytorch/pytorch/issues/108337 . It is fairly similar to maskrcnn which is one of the priority models that internal export. While investigating the user report I hacked up a bunch of stuff in torchvision which isn't easily landable, I want to durably record it here.

Repro script:

```
import torch
import torchvision

model = torchvision.models.detection.retinanet_resnet50_fpn(
        weights=torchvision.models.detection.RetinaNet_ResNet50_FPN_Weights.DEFAULT)
model = torch.compile(model, mode="default")

torch.save(model.state_dict(), "retina.pt")
print("model.state_dict:",model.state_dict().keys())

x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
model.eval()
exported_model = torch._dynamo.export(model, x)
torch.save(exported_model, "retina_retina.pt")
```

Tested on 621463a3e6b488b2bff04e355a1abd9a4c5bb2cd

Last time I did this for maskrcnn: https://docs.google.com/document/d/159NTQQhz8ovIBxbQvGQ-fZ10pF9e2RPXm1JZYqdEzt4/edit#heading=h.jw6vkqei769s

**f-string problem.** Same as in maskrcnn. Dynamo still chokes on f-string messages (known bug for quite a long time, https://github.com/pytorch/pytorch/issues/103602 )

```
diff --git a/torchvision/models/detection/retinanet.py b/torchvision/models/detection/retinanet.py
index 3a9cf80d1d..3fe67f6440 100644
--- a/torchvision/models/detection/retinanet.py
+++ b/torchvision/models/detection/retinanet.py
@@ -596,10 +596,7 @@ class RetinaNet(nn.Module):
         original_image_sizes: List[Tuple[int, int]] = []
         for img in images:
             val = img.shape[-2:]
-            torch._assert(
-                len(val) == 2,
-                f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}",
-            )
+            assert len(val) == 2
             original_image_sizes.append((val[0], val[1]))
 
         # transform the input
```

**Scale factor tire fire.** Same as in maskrcnn. I actually half fixed this but one last bit I forgot to do last time https://github.com/pytorch/vision/pull/7942

**Data-dependent orig_kval.**

```
diff --git a/torchvision/models/detection/_utils.py b/torchvision/models/detection/_utils.py
index 559db858ac..bb13d8b72c 100644
--- a/torchvision/models/detection/_utils.py
+++ b/torchvision/models/detection/_utils.py
@@ -506,7 +506,7 @@ def _topk_min(input: Tensor, orig_kval: int, axis: int) -> int:
         min_kval (int): Appropriately selected k-value.
     """
     if not torch.jit.is_tracing():
-        return min(orig_kval, input.size(axis))
+        return torch.sym_min(orig_kval, input.size(axis))
     axis_dim_val = torch._shape_as_tensor(input)[axis].unsqueeze(0)
     min_kval = torch.min(torch.cat((torch.tensor([orig_kval], dtype=axis_dim_val.dtype), axis_dim_val), 0))
     return _fake_cast_onnx(min_kval)
```

`orig_kval` is unbacked so we cannot do a stock min on it which will attempt to test if `orig_kval < input.size(axis)` which we won't know. Dynamo ought to be able to translate this to sym_min automatically.

**NN module setattr** Same as in maskrcnn

```
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
index 9db4fb24bb5..56b5c5e3be0 100644
--- a/torch/_dynamo/symbolic_convert.py
+++ b/torch/_dynamo/symbolic_convert.py
@@ -1201,13 +1201,6 @@ class InstructionTranslatorBase(Checkpointable[InstructionTranslatorGraphState])
         prior = self.copy_graphstate()
         val, obj = self.popn(2)
 
-        if isinstance(obj, NNModuleVariable):
-            # We don't allow side effects during export
-            # https://github.com/pytorch/torchdynamo/issues/1475
-            assert (
-                not self.export
-            ), f"Mutating module attribute {inst.argval} during export."
-
         try:
             self.output.guards.update(
                 BuiltinVariable(setattr)
```

**Batched nms threshold trick.** We can either torch.cond this, or taking a page from torchvision._is_tracing test here, just always do the coordinate trick.

```
diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py
index a541f8d880..b60f6e6254 100644
--- a/torchvision/ops/boxes.py
+++ b/torchvision/ops/boxes.py
@@ -69,7 +69,7 @@ def batched_nms(
         _log_api_usage_once(batched_nms)
     # Benchmarks that drove the following thresholds are at
     # https://github.com/pytorch/vision/issues/1311#issuecomment-781329339
-    if boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing():
+    if False and boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing():
         return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
     else:
         return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)
```

**nms support.** I could have sworn that I meta-fied this but apparently not. NMS is data dependent so it needs @zou3519 impl_abstract for out of tree registration.

```
torch._dynamo.exc.Unsupported: unsupported operator: torchvision.nms.default (see https://docs.google.com/document/d/1GgvOe7C8_NVOMLOCwDaYV1mXXyHMXY7ExoewHqooxrs/edit#heading=h.64r4npvq0w0 for how to fix)
```

### Versions

main

cc @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519