Vectorization refactor #205
Merged
Conversation
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 was referenced Jun 6, 2022
fantaosha
reviewed
Jun 7, 2022
…a cost function and its weight.
…Layer to create them at init time.
…rator, and added it to TheseusLayer.
… from vectorization. Otherwise it works just like the cost function it wraps.
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
mhmukadam
reviewed
Jun 8, 2022
There was a problem hiding this comment.
Nice work on the refactor @luisenp and of course nice work on the original features @fantaosha! Reviewed everything so far and left comments.
mhmukadam
reviewed
Jun 9, 2022
rtqichen
approved these changes
Jun 9, 2022
…-to-end tests now run with vectorize=True.
mhmukadam
reviewed
Jun 9, 2022
theseus/core/theseus_function.py
Outdated
| @@ -64,6 +64,16 @@ def register_aux_vars(self, aux_var_names: Sequence[str]): | |||
| for name in aux_var_names: | |||
| self.register_aux_var(name) | |||
|
|
|||
| def register_vars(self, vars: Iterable[Variable], is_optim_vars: bool = False): | |||
| for var_ in vars: | |||
There was a problem hiding this comment.
Nit: trailing underscore.
There was a problem hiding this comment.
Good catch. Changed.
mhmukadam
reviewed
Jun 9, 2022
| @@ -80,7 +80,7 @@ def ee_pose_err_fn(optim_vars, aux_vars): | |||
| max_iterations=15, | |||
| step_size=0.5, | |||
| ) | |||
| theseus_optim = th.TheseusLayer(optimizer) | |||
| theseus_optim = th.TheseusLayer(optimizer, vectorize=False) | |||
There was a problem hiding this comment.
Oh, forgot about these two. Added.
mhmukadam
reviewed
Jun 9, 2022
mhmukadam
approved these changes
Jun 9, 2022
bamos
approved these changes
Jun 9, 2022
suddhu
pushed a commit
to suddhu/theseus
that referenced
this pull request
Jan 21, 2023
* Created a wrapper cost function class that combines the aux vars for a cost function and its weight. * Disabled support for optimization variables in cost weights. * Changed Objective to iterate over CFWrapper if available, and TheseusLayer to create them at init time. * Added a Vectorizer class and moved CFWrappers there. * Renamed vectorizer as Vectorize, added logic to replace Objective iterator, and added it to TheseusLayer. * Added a CostFunctionSchema -> List[CostFunction] to use for vectorization grouping. * _CostFunctionWrapper is now meant to just store a cached value coming from vectorization. Otherwise it works just like the cost function it wraps. * Added code to automatically compute shared vars in Vectorize. * Changed vectorized costs construction to ensure that their weight is also a copy. * Implemented main cost function vectorization logic. * Updated bug that was causing detached gradients. * Fixed invalid check in theseus end-to-end unit tests. * Added unit test for schema and shared var computation. * Added a test to check that computed vectorized errors are correct. * Moved vectorization update call to base linearization class. * Changed code to allow batch_size > 1 in shared variables. * Fixed unit test and added call to Objective.update() in update_vectorization() if batch_size is None. * Added new private iterator for vectorized costs. * Replaced _register_vars_in_list with TheseusFunction.register_vars. * Renamed vectorize_cost_fns kwarg as vectorize. * Added license headers. * Small refactor. * Fixed bug that was preventing vectorized costs to work with to(). End-to-end tests now run with vectorize=True. * Renamed the private Objective cost function iterator to _get_iterator(). * Renamed kwarg in register_vars. * Set vectorize=True for inverse kinematics and backward tests. * Remove lingering comments. Co-authored-by: Taosha Fan <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Refactored the vectorization logic. I did some quick tests on motion planning and tactile pose estimation, and seems to be working.
Some TODOs left:
retract_and_update_variablesin that one as well)