A bit of cleanup, testing w/ PyTorch 2.0 nightlies, initial torchcompile support

Author: rwightman

README.md

@@ -16,6 +16,11 @@ Aside from the default model configs, there is a lot of flexibility to facilitat
 
 ## Updates
 
+### 2023-02-09
+* Testing with PyTorch 2.0 (nightlies), add --torchcompile support to train and validate scripts
+* A small code cleanup pass, support bwd/fwd compat across timm 0.8.x and previous releases
+* Use `timm` convert_sync_batchnorm function as it handles updated models w/ BatchNormAct2d layers
+
 ### 2022-01-06
 * New `efficientnetv2_ds` weights 50.1 mAP @ 1024x0124, using AGC clipping and `timm`'s `efficientnetv2_rw_s` backbone. Memory use comparable to D3, speed faster than D4. Smaller than optimal training batch size so can probably do better... 
 

effdet/bench.py

@@ -80,8 +80,16 @@ def _batch_detection(
         img_scale_i = None if img_scale is None else img_scale[i]
         img_size_i = None if img_size is None else img_size[i]
         detections = generate_detections(
-            class_out[i], box_out[i], anchor_boxes, indices[i], classes[i],
-            img_scale_i, img_size_i, max_det_per_image=max_det_per_image, soft_nms=soft_nms)
+            class_out[i],
+            box_out[i],
+            anchor_boxes,
+            indices[i],
+            classes[i],
+            img_scale_i,
+            img_size_i,
+            max_det_per_image=max_det_per_image,
+            soft_nms=soft_nms,
+        )
         batch_detections.append(detections)
     return torch.stack(batch_detections, dim=0)
 
@@ -101,15 +109,27 @@ def __init__(self, model):
     def forward(self, x, img_info: Optional[Dict[str, torch.Tensor]] = None):
         class_out, box_out = self.model(x)
         class_out, box_out, indices, classes = _post_process(
-            class_out, box_out, num_levels=self.num_levels, num_classes=self.num_classes,
-            max_detection_points=self.max_detection_points)
+            class_out,
+            box_out,
+            num_levels=self.num_levels,
+            num_classes=self.num_classes,
+            max_detection_points=self.max_detection_points,
+        )
         if img_info is None:
             img_scale, img_size = None, None
         else:
             img_scale, img_size = img_info['img_scale'], img_info['img_size']
         return _batch_detection(
-            x.shape[0], class_out, box_out, self.anchors.boxes, indices, classes,
-            img_scale, img_size, max_det_per_image=self.max_det_per_image, soft_nms=self.soft_nms
+            x.shape[0],
+            class_out,
+            box_out,
+            self.anchors.boxes,
+            indices,
+            classes,
+            img_scale,
+            img_size,
+            max_det_per_image=self.max_det_per_image,
+            soft_nms=self.soft_nms,
         )
 
 
@@ -126,7 +146,11 @@ def __init__(self, model, create_labeler=True):
         self.soft_nms = model.config.soft_nms
         self.anchor_labeler = None
         if create_labeler:
-            self.anchor_labeler = AnchorLabeler(self.anchors, self.num_classes, match_threshold=0.5)
+            self.anchor_labeler = AnchorLabeler(
+                self.anchors,
+                self.num_classes,
+                match_threshold=0.5,
+            )
         self.loss_fn = DetectionLoss(model.config)
 
     def forward(self, x, target: Dict[str, torch.Tensor]):
@@ -139,19 +163,39 @@ def forward(self, x, target: Dict[str, torch.Tensor]):
             num_positives = target['label_num_positives']
         else:
             cls_targets, box_targets, num_positives = self.anchor_labeler.batch_label_anchors(
-                target['bbox'], target['cls'])
-
-        loss, class_loss, box_loss = self.loss_fn(class_out, box_out, cls_targets, box_targets, num_positives)
+                target['bbox'],
+                target['cls'],
+            )
+
+        loss, class_loss, box_loss = self.loss_fn(
+            class_out,
+            box_out,
+            cls_targets,
+            box_targets,
+            num_positives,
+        )
         output = {'loss': loss, 'class_loss': class_loss, 'box_loss': box_loss}
         if not self.training:
             # if eval mode, output detections for evaluation
             class_out_pp, box_out_pp, indices, classes = _post_process(
-                class_out, box_out, num_levels=self.num_levels, num_classes=self.num_classes,
-                max_detection_points=self.max_detection_points)
+                class_out,
+                box_out,
+                num_levels=self.num_levels,
+                num_classes=self.num_classes,
+                max_detection_points=self.max_detection_points,
+            )
             output['detections'] = _batch_detection(
-                x.shape[0], class_out_pp, box_out_pp, self.anchors.boxes, indices, classes,
-                target['img_scale'], target['img_size'],
-                max_det_per_image=self.max_det_per_image, soft_nms=self.soft_nms)
+                x.shape[0],
+                class_out_pp,
+                box_out_pp,
+                self.anchors.boxes,
+                indices,
+                classes,
+                target['img_scale'],
+                target['img_size'],
+                max_det_per_image=self.max_det_per_image,
+                soft_nms=self.soft_nms,
+            )
         return output