-
Notifications
You must be signed in to change notification settings - Fork 3.3k
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
Add trace functionality to the function to_torchscript #4142
Merged
Merged
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here you assume that
example_input_array
is a tensor, but this is not true.If forward takes
*args
,example_input_array
is a tuple, and if forward takes**kwargs
,example_input_array
must be a dict.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Given your comment on the issue that
.trace()
accepts either atuple
or atorch.Tensor
(that is automatically converted to a tuple), it means that the input should be:example_input_array: Optional[Union[torch.Tensor, Tuple[torch.Tensor]]]
?However, when the forward function accepts
**kwargs
,self.example_input_array
could be a dict, in which case.trace(example_inputs=example_inputs)
will fail?What would be the best way to approach this? Does this mean that
.trace()
cannot be used ifforward
expects a dict?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here is an example. tracing supports single input and tuple, which gets unrolled to multiple positional args. In these two cases, you can use the Lightning self.example_input_array. However, dicts will not be passed as kwargs, and instead as a single input. In Lightning however, a dict would mean **kwargs.
I see several ways to handle it:
self.example_input_array
is a dictself.example_input_array
, and require the user to give inputs to the method directlyThen there is a second issue. You should use the
pytorch_lightning.utilities.apply_func.move_data_to_device
to move the example input to the device, since it could be a tuple.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
cc @ananthsub
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
1 & 2 could be combined by raising a warning instead of an error. From PL's side throw a warning similar to:
Then output the actual error produced by
.trace()
.If in the future
.trace()
would be updated to support a dict, there is no need for a change (except removing the warning) on PL's side.Personally, PL is for me about removing boilerplate code. Since
self.example_input_array
is already a thing in PL, it's better to use it. Therefore, I would advise against option 3.I haven't used
self.example_input_array
personally yet, but in how many projects would this be a dict?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
yes, makes sense.
Would you like to follow up on this with a PR? Would greatly appreciate this. For me the main concern is to properly move the input to the device with the function I referenced. For the way inputs are passed in, I don't have a strong opinon.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
IDK how much future support there will be for tracing vs scripting (scripting is strongly recommended). Rather than adding more trace support at the top-level of the PL module, why not override
to_torchscript
in your lightning module to determine how you want to export? then you have way more flexibility with tracingThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@awaelchli Ok, I'll follow up with another pull request using the
move_data_to_device
function.@ananthsub edit moved my comment to the feature request, as it is a more relevant place for this discussion: #4140
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@awaelchli I addressed your issues in a follow-up pull request (could not be added to this one due to it already being merged):
#4360