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.

Sign up for GitHub

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

[Refactor] Deprecate direct usage of memmap tensors #1684

Merged
merged 11 commits into from
Nov 15, 2023
Merged

[Refactor] Deprecate direct usage of memmap tensors #1684

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
eeefb18
empty
vmoens Nov 8, 2023
3832c37
amend
vmoens Nov 8, 2023
dd525cf
amend
vmoens Nov 14, 2023
24a4ed3
amend
vmoens Nov 14, 2023
e362a75
Merge branch 'main' into deprecate_direct_memmap
vmoens Nov 14, 2023
0f9dec3
fix rb_dist tests
vmoens Nov 14, 2023
382e15e
Update run_all.sh
vmoens Nov 14, 2023
0df1bed
lint
vmoens Nov 14, 2023
8e03df4
empty
vmoens Nov 15, 2023
4d1c67c
amend
vmoens Nov 15, 2023
07cbbfe
amend
vmoens Nov 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
amend
  • Loading branch information
@vmoens
vmoens committed Nov 14, 2023
commit dd525cfa469a510a53f5b26f2cc6a6f44a9161b7
2 changes: 1 addition & 1 deletion .github/unittest/linux/scripts/run_all.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ python -c "import functorch"
pip3 install git+https://github.com/pytorch/torchsnapshot

# install tensordict
pip3 install git+https://github.com/pytorch/tensordict.git
pip3 install git+https://github.com/pytorch/tensordict.git@memmap_tensor_refact

printf "* Installing torchrl\n"
python setup.py develop
Expand Down
19 changes: 3 additions & 16 deletions torchrl/data/replay_buffers/storages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

import torch
from tensordict import is_tensorclass
from tensordict.memmap import MemmapTensor
from tensordict.memmap import MemoryMappedTensor
from tensordict.tensordict import is_tensor_collection, TensorDict, TensorDictBase
from tensordict.utils import expand_right

Expand Down Expand Up @@ -760,20 +760,7 @@ def _get_default_collate(storage, _is_tensordict=False):
f"Could not find a default collate_fn for storage {type(storage)}."
)

@implement_for("torch", None, "2.2.0")
def make_memmap(tensor, path):
return MemmapTensor.from_tensor(tensor, filename=path)

@implement_for("torch", "2.2.0")
def make_memmap(tensor, path):
from tensordict._memory_map import from_tensor
return from_tensor(tensor, filename=path)

@implement_for("torch", None, "2.2.0")
def make_empty_memmap(shape, dtype, path):
return MemmapTensor(shape, dtype=dtype, filename=path)

@implement_for("torch", "2.2.0")
return MemoryMappedTensor.from_tensor(tensor, filename=path)
def make_empty_memmap(shape, dtype, path):
from tensordict._memory_map import empty_like
return empty_like(torch.zeros((), dtype=dtype).expand(shape), filename=path)
return MemoryMappedTensor.from_tensor(tensor, filename=path)
Loading