[Feature]: Added support for TensorDictSequence module subsampling #332
Conversation
facebook-github-bot
added
the
CLA Signed
|
@vmoens won't this be a problem to create new sub_seq to frequently. Typically for Dreamer, this would require to access the subsequence after every training step. So if we create a subsequence each time this will be bad for performance no? |
|
Can't you reuse it? What would change between two iterations? |
|
The weights of the main sequence will have been updated so we need to update the weights of the sub sequence |
|
Not if they're the same weights |
|
I made the proposed changed and works as expected thanks for the suggestion |
vmoens
changed the title
| out_keys = deepcopy(self.out_keys) | ||
| id_to_keep = set([i for i in range(len(self.module))]) | ||
| for i, module in enumerate(self.module): | ||
| if all(key in in_keys for key in module.in_keys): |
There was a problem hiding this comment.
What's the usage of selecting in_keys? I can understand why we want to restrict the outputs, but I don't really see when we want to restrict inputs.
There was a problem hiding this comment.
Also what happens if you say you want some out_keys but they conflict with the in_keys? Is the sequence going to be empty?
There was a problem hiding this comment.
We want to be able to select the in_keys to be able to directly input an intermediarry block. For example imagine you have a hidden layer that you want to inject from a precomputed tensordict, this allows to do so.
There was a problem hiding this comment.
Yes it would be empty if your out keys are before the in_keys
There was a problem hiding this comment.
Got it
We're doing something similar in #352: you can input an incomplete tensordict and only the relevant ops will be executed. I wonder if we need both ways of doing the same thing. The advantage of your implementation is that it is self-consistent though.
|
I see that the tests are failing, have you tried merging main into this branch? |
Description
This PR adds the possibility of retrieving a consistent sub-network from a TensorDictSequence that is restricted to a given set of inputs and outputs. Only the blocks that are involved in the computation between input and output will be called.
Motivation and Context
The tracing capabilities of TensorDictSequence make it possible to look at restricted sub-networks with little effort. This allows for faster code execution (compared to running full modules where only part of the output is needed) and greater readability.
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!