seeker
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
import torch
|
|
|
|
from pi0mini.model import Pi0Config, Pi0Mini
|
|
|
|
|
|
def test_pi0_three_block_attention_mask() -> None:
|
|
config = Pi0Config(
|
|
image_size=32,
|
|
patch_size=16,
|
|
max_text_tokens=4,
|
|
max_action_tokens=3,
|
|
d_model=32,
|
|
n_heads=4,
|
|
n_layers=1,
|
|
time_feature_dim=8,
|
|
)
|
|
model = Pi0Mini(config)
|
|
|
|
# Two image patches per axis -> 4 image tokens, plus 2 language tokens.
|
|
prefix_edges, suffix_edges = model._masks(
|
|
batch=1,
|
|
prefix_length=6,
|
|
suffix_length=4, # one state + three action tokens
|
|
image_tokens=4,
|
|
attention_mask=torch.tensor([[1, 1]]),
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
assert prefix_edges.shape == (1, 6, 6)
|
|
assert torch.all(prefix_edges)
|
|
|
|
# Suffix key order is [six prefix keys, state, action0, action1, action2].
|
|
assert suffix_edges.shape == (1, 4, 10)
|
|
assert torch.all(suffix_edges[0, 0, :7])
|
|
assert not torch.any(suffix_edges[0, 0, 7:])
|
|
assert torch.all(suffix_edges[0, 1:, :])
|
|
|
|
|
|
def test_pi0_mask_respects_padded_language_keys() -> None:
|
|
model = Pi0Mini(
|
|
Pi0Config(
|
|
image_size=16,
|
|
patch_size=16,
|
|
max_text_tokens=3,
|
|
max_action_tokens=1,
|
|
d_model=16,
|
|
n_heads=4,
|
|
n_layers=1,
|
|
time_feature_dim=8,
|
|
)
|
|
)
|
|
prefix_edges, suffix_edges = model._masks(
|
|
batch=1,
|
|
prefix_length=4,
|
|
suffix_length=2,
|
|
image_tokens=1,
|
|
attention_mask=torch.tensor([[1, 0, 0]]),
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
# Padded prefix keys are hidden from valid prefix, state, and action rows.
|
|
assert not prefix_edges[0, 0, 2]
|
|
assert not prefix_edges[0, 0, 3]
|
|
assert not suffix_edges[0, :, 2].any()
|
|
assert not suffix_edges[0, :, 3].any()
|