-
Notifications
You must be signed in to change notification settings - Fork 6.9k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.
Already on GitHub? Sign in to your account
RegionProposalNetwork can't be AOTInductor compiled with dynamic batch size #8285
Comments
Thanks for the report @rbavery
No, we haven't been looking at the RPN's support for torch.compile yet.
Just note that a lot of this code is public and changing the behaviour e.g. the expected intput/output would technically be breaking backward compatibility. So adding support for AOT while still preserving BC may not be a trivial task. Beyond the RPN, what model specifically are you interested in tracing? |
Got it, I initially went with supporting TorchScript scripting since it seemed easier and would only require adding type annotations. I've made edits to this model which uses a SWIN Transformer backbone, an FPN, and a Faster RCNN head: https://github.com/allenai/satlas/blob/main/configs/satlas_explorer_marine_infrastructure.txt So far I addressed torchscript scripting issues with type annotations in the Satlas model source. the first issue I hit was with torchvision is here:
and I've had some trouble addressing this with typing, since the class attribute is already typed, I'm not sure how to enable Torschript scripting to understand this attribute can be either a List[float] or None. I might need to make code modifications. I'll try to do so in a way that preserves backwards compat and leaves passing test and PR if it is helpful. |
I was able to get torch scripting to work by refactoring the Satlas source code, mostly by adding typing, removing control flow in some spots, and replacing the use of complex python data structures containing tensors with plain tensors. inference on dynamic batches appears to work without error. No changes to torchvision were needed. but not AOTInductor unfortunately. I made some progress forking torchvision and trying to remove the use of Both methods for exporting were relatively painful. I'm hoping that AOTInductor comes up with a solution for handling data-dependent shapes, or making the process to write code that handles data-dependent shapes easier. I realize it is early days for AOTInductor still, but documentation would go a long way. I'd be happy to contribute but still feel fairly new to the process of handing data dependent shapes. |
馃悰 Describe the bug
this is a cross post of pytorch/pytorch#121036
Just raising it here to notify the maintainers that I'm going to take a crack at fixing the RegionProposalNetwork and potentially other modules to be either traceable, AOTInductor compileable, or both. Are there any current efforts in this direction I should be aware of?
For AOTInductor I think this will at least involve changing the AnchorGenerator, which has a method that mutates an anchor attribute to instead return anchor values.
To support tracing, my plan is to address each TracerWarning (see below). First I'll be looking to remove the iteration over tensors in ImageList that prevent the model from generalziing after tracing.
Versions
I'm using the nightlies, see pytorch/pytorch#121036
The text was updated successfully, but these errors were encountered: