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

Jump to bottom

[Tests] Adding tensordict __repr__ tests #435

Merged
merged 18 commits into from
Sep 21, 2022
Merged

[Tests] Adding tensordict __repr__ tests #435

Changes from 1 commit
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
36526d5
Adding tensordict repr tests
Sep 12, 2022
c42a140
empty commit
vmoens Sep 13, 2022
056e40d
Merge branch 'facebookresearch:main' into tensordict_repr_tests
sladebot Sep 13, 2022
8342872
Adding tests for indexed tensordicts (nested)
Sep 13, 2022
2357963
Merge branch 'main' into tensordict_repr_tests
vmoens Sep 15, 2022
1b4e642
Merge branch 'tensordict_repr_tests' of https://github.com/sladebot/r…
vmoens Sep 15, 2022
194c5cd
refactor tests
vmoens Sep 15, 2022
0cfc283
refactor tests
vmoens Sep 15, 2022
9fab2de
Refactoring indexed tensordict tests
Sep 16, 2022
f82201a
proposed changes to shared tensors
vmoens Sep 19, 2022
7ea6476
amend previous commit
vmoens Sep 19, 2022
0ad21b8
Merge branch 'main' into tensordict_repr_tests
vmoens Sep 19, 2022
7dc9638
amend previous commit
vmoens Sep 19, 2022
98fd65f
always display "tensor" unless shared tensor on cpu
vmoens Sep 20, 2022
ef82d21
fix test_repr_share_memory
vmoens Sep 20, 2022
e5e471d
fix naming after sharing mem of MetaTensor
vmoens Sep 20, 2022
312a74e
amend
vmoens Sep 20, 2022
1a8e5cd
amend
vmoens Sep 20, 2022
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
proposed changes to shared tensors
  • Loading branch information
@vmoens
vmoens committed Sep 19, 2022
commit f82201ab0274b7eeeb57b1bec068fe61f9b36fca
59 changes: 35 additions & 24 deletions test/test_tensordict.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1681,7 +1681,7 @@ def share_memory_td(self, device, dtype):

def test_repr_plain(self, device, dtype):
tensordict = self.td(device, dtype)
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
else:
is_shared = False
Expand All @@ -1695,7 +1695,7 @@ def test_repr_plain(self, device, dtype):

def test_repr_memmap(self, device, dtype):
tensordict = self.memmap_td(device, dtype)
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
else:
is_shared = False
Expand All @@ -1710,26 +1710,29 @@ def test_repr_memmap(self, device, dtype):
def test_repr_share_memory(self, device, dtype):
tensordict = self.share_memory_td(device, dtype)
is_shared = True
tensor_class = "SharedTensor"
expected = f"""TensorDict(
fields={{
a: SharedTensor(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})"""
assert repr(tensordict) == expected

def test_repr_nested(self, device, dtype):
nested_td = self.nested_td(device, dtype)
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_class = "SharedTensor"
else:
is_shared = False
tensor_class = "Tensor"
expected = f"""TensorDict(
fields={{
b: Tensor(torch.Size([4, 3, 2, 1, 5]), dtype={dtype}),
b: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype}),
my_nested_td: TensorDict(
fields={{
a: Tensor(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})}},
Expand All @@ -1740,13 +1743,15 @@ def test_repr_nested(self, device, dtype):

def test_repr_stacked(self, device, dtype):
stacked_td = self.stacked_td(device, dtype)
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_class = "SharedTensor"
else:
is_shared = False
tensor_class = "Tensor"
expected = f"""LazyStackedTensorDict(
fields={{
a: Tensor(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})"""
Expand All @@ -1755,21 +1760,23 @@ def test_repr_stacked(self, device, dtype):
@pytest.mark.parametrize("index", [None, (slice(None), 0)])
def test_repr_indexed_tensordict(self, device, dtype, index):
tensordict = self.td(device, dtype)[index]
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_class = "SharedTensor"
else:
is_shared = False
tensor_class = "Tensor"
if index is None:
expected = f"""TensorDict(
fields={{
a: Tensor(torch.Size([1, 4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([1, 4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([1, 4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})"""
else:
expected = f"""TensorDict(
fields={{
a: Tensor(torch.Size([4, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 2, 1]),
device={str(device)},
is_shared={is_shared})"""
Expand All @@ -1779,17 +1786,19 @@ def test_repr_indexed_tensordict(self, device, dtype, index):
@pytest.mark.parametrize("index", [None, (slice(None), 0)])
def test_repr_indexed_nested_tensordict(self, device, dtype, index):
nested_tensordict = self.nested_td(device, dtype)[index]
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_class = "SharedTensor"
else:
is_shared = False
tensor_class = "Tensor"
if index is None:
expected = f"""TensorDict(
fields={{
b: Tensor(torch.Size([1, 4, 3, 2, 1, 5]), dtype={dtype}),
b: {tensor_class}(torch.Size([1, 4, 3, 2, 1, 5]), dtype={dtype}),
my_nested_td: TensorDict(
fields={{
a: Tensor(torch.Size([1, 4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([1, 4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([1, 4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})}},
Expand All @@ -1799,10 +1808,10 @@ def test_repr_indexed_nested_tensordict(self, device, dtype, index):
else:
expected = f"""TensorDict(
fields={{
b: Tensor(torch.Size([4, 2, 1, 5]), dtype={dtype}),
b: {tensor_class}(torch.Size([4, 2, 1, 5]), dtype={dtype}),
my_nested_td: TensorDict(
fields={{
a: Tensor(torch.Size([4, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 2, 1]),
device={str(device)},
is_shared={is_shared})}},
Expand All @@ -1814,21 +1823,23 @@ def test_repr_indexed_nested_tensordict(self, device, dtype, index):
@pytest.mark.parametrize("index", [None, (slice(None), 0)])
def test_repr_indexed_stacked_tensordict(self, device, dtype, index):
stacked_tensordict = self.stacked_td(device, dtype)
if device is not None and device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_class = "SharedTensor"
else:
is_shared = False
tensor_class = "Tensor"
if index is None:
expected = f"""LazyStackedTensorDict(
fields={{
a: Tensor(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})"""
else:
expected = f"""LazyStackedTensorDict(
fields={{
a: Tensor(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 3, 2, 1]),
device={str(device)},
is_shared={is_shared})"""
Expand All @@ -1838,16 +1849,16 @@ def test_repr_indexed_stacked_tensordict(self, device, dtype, index):
@pytest.mark.parametrize("device_cast", get_available_devices())
def test_repr_device_to_device(self, device, dtype, device_cast):
td = self.td(device, dtype)
if device_cast.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_klass = "SharedTensor"
tensor_class = "SharedTensor"
else:
is_shared = False
tensor_klass = "Tensor"
tensor_class = "Tensor"
td2 = td.to(device_cast)
expected = f"""TensorDict(
fields={{
a: {tensor_klass}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
a: {tensor_class}(torch.Size([4, 3, 2, 1, 5]), dtype={dtype})}},
batch_size=torch.Size([4, 3, 2, 1]),
device={str(device_cast)},
is_shared={is_shared})"""
Expand All @@ -1859,7 +1870,7 @@ def test_repr_batch_size_update(self, device, dtype):
td.batch_size = torch.Size([4, 3, 2])
is_shared = False
tensor_class = "Tensor"
if device is None or device.type == "cuda":
if device.type == "cuda" or (device is None and torch.cuda.device_count() > 0):
is_shared = True
tensor_class = "SharedTensor"
expected = f"""TensorDict(
Expand Down