[tracer] skip_unsupported torchvision ops (#365)

alibaba · Oct 1, 2024 · 7a4aee5 · 7a4aee5
1 parent 133c095
commit 7a4aee5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/tinynn/graph/configs/gen_modules_yml.py b/tinynn/graph/configs/gen_modules_yml.py
@@ -49,7 +49,13 @@
         print(k, c, issubclass(c, torch.nn.Module))
         print(c.__module__)
         # Skip container modules
-        if c.__name__ in ('FeaturePyramidNetwork', 'SqueezeExcitation'):
+        if c.__name__ in (
+            'FeaturePyramidNetwork',
+            'SqueezeExcitation',
+            'Conv2dNormActivation',
+            'Conv3dNormActivation',
+            'MLP',
+        ):
             continue
         final_dict.setdefault('torchvision.ops', [])
         final_dict['torchvision.ops'].append(k)

diff --git a/tinynn/graph/configs/torch_module_override.yml b/tinynn/graph/configs/torch_module_override.yml
@@ -159,10 +159,7 @@ torchvision.ops:
 - DeformConv2d
 - DropBlock2d
 - DropBlock3d
-- Conv2dNormActivation
-- Conv3dNormActivation
 - FrozenBatchNorm2d
-- MLP
 - Permute
 - MultiScaleRoIAlign
 - PSRoIAlign