28 lines
798 B
Python
28 lines
798 B
Python
import torch
|
|
|
|
from losses.matching_loss import hungarian_matching_loss
|
|
from models.slot_extractor import SlotExtractor
|
|
|
|
|
|
def test_slot_extractor_and_matching_loss_backward():
|
|
extractor = SlotExtractor(
|
|
input_dim=32,
|
|
state_dim=64,
|
|
num_slots=4,
|
|
num_layers=2,
|
|
num_heads=4,
|
|
dropout=0.0,
|
|
num_subjects=12,
|
|
num_predicates=7,
|
|
num_objects=12,
|
|
)
|
|
tokens = torch.randn(2, 10, 32)
|
|
mask = torch.ones(2, 10, dtype=torch.bool)
|
|
predictions = extractor(tokens, mask)
|
|
targets = [torch.tensor([[0, 0, 4], [1, 1, 5]]), torch.tensor([[1, 0, 6]])]
|
|
result = hungarian_matching_loss(predictions, targets)
|
|
assert torch.isfinite(result.loss)
|
|
result.loss.backward()
|
|
assert extractor.slot_queries.grad is not None
|
|
|