Skip to content
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’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Enhancement] Change the order of condition to make fx wok #2883

Merged
merged 3 commits into from
Aug 3, 2023

Conversation

youkaichao
Copy link
Contributor

Motivation

mmcv.cnn.bricks.wrappers has many forward functions in the form of if x.numel() == 0 and obsolete_torch_version(TORCH_VERSION, (1, 5)), where the first condition x.numel() == 0 would fail torch.fx. We can take advantage of shortcut evaluation, by changing the order of condition to if obsolete_torch_version(TORCH_VERSION, (1, 5)) and x.numel() == 0. This way, the torch_version compares first, and will not fail torch.fx in modern pytorch versions.

Modification

Change the order of condition, and add tests for it.

BC-breaking (Optional)

This PR does not break any compatibility.

Comment on lines 388 to 396
print(gm_module.code)
for Net in (Linear, ):
net = Net(1, 1)
gm_module = fx.symbolic_trace(net)
print(gm_module.code)
for Net in (Conv2d, ConvTranspose2d, Conv3d, ConvTranspose3d):
net = Net(1, 1, 1)
gm_module = fx.symbolic_trace(net)
print(gm_module.code)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it necessary to print the code?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh, that's used to supress the unused variable error.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, you can replace gm_module = fx.symbolic_trace(net) with fx.symbolic_trace(net).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in fc883ec.

@zhouzaida zhouzaida changed the title change the order of condition, so that torch.fx can trace these modules [Enhancement] Change the order of condition to make fx wok Aug 3, 2023
@zhouzaida zhouzaida merged commit 9036241 into open-mmlab:main Aug 3, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants