Skip to content

Conversation

vmoens
Copy link
Collaborator

@vmoens vmoens commented Aug 5, 2022

Description

This PR introduces several changes aimed at making replay buffer sampling more efficient:

  • MemmapTensor.clone() now returns a regular tensor. This is a more intuitive behaviour than the previous one (where another MemmapTensor was returned). Also, to get another MemmapTensor, one can simply call MemmapTensor(memmap.clone()).
  • tensordict indexing now never returns a SubTensorDict. To get a SubTensorDict, one must built it explicitely by calling the appropriate dedicated method. This is to make sure that no overhead is introduced by working with this class.
  • Replay buffer storage is now cast to the appropriate device by the helper function: this makes data sampling faster and avoids the creation of tensors on cpu that are then cast on gpu (tensors are now created on gpu immediately).

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 5, 2022
@vmoens vmoens added the bug Something isn't working label Aug 6, 2022
@vmoens vmoens changed the title [BugFix] Clone memmap tensors on regular tensors [BugFix] Clone memmap tensors on regular tensors and other replay buffer improvements Aug 6, 2022
)
from torchrl.objectives.costs.common import LossModule
from torchrl.objectives.costs.redq import REDQLoss
from torchrl.objectives.costs.deprecated import REDQLoss_deprecated
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

REDQLoss is memory intensive when used with large models.
Some work is needed to make it less memory intensive.

@vmoens vmoens merged commit c61ae7b into main Aug 8, 2022
@vmoens vmoens deleted the clone_memmap branch August 8, 2022 06:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants