56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import torch
|
|
|
|
from data.schema import FACT_TYPE, RULE_TYPE
|
|
from losses.alignment_loss import cross_modal_alignment_loss
|
|
from losses.matching_loss import hungarian_matching_loss
|
|
from losses.prototype_loss import canonical_prototype_loss
|
|
from models.fact_prototype import CanonicalFactPrototype
|
|
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,
|
|
num_state_types=3,
|
|
)
|
|
tokens = torch.randn(2, 10, 32)
|
|
mask = torch.ones(2, 10, dtype=torch.bool)
|
|
predictions = extractor(tokens, mask)
|
|
targets = [
|
|
torch.tensor(
|
|
[
|
|
[FACT_TYPE, 0, 0, 4, -1, -1],
|
|
[RULE_TYPE, 1, 1, 5, 6, 7],
|
|
]
|
|
),
|
|
torch.tensor([[FACT_TYPE, 1, 0, 6, -1, -1]]),
|
|
]
|
|
result = hungarian_matching_loss(predictions, targets)
|
|
assert torch.isfinite(result.loss)
|
|
result.loss.backward()
|
|
assert extractor.slot_queries.grad is not None
|
|
|
|
|
|
def test_cross_modal_alignment_and_prototype_losses():
|
|
assignments = [(torch.tensor([0]), torch.tensor([0]))]
|
|
targets = [torch.tensor([[FACT_TYPE, 0, 0, 4, -1, -1]])]
|
|
visual = torch.randn(1, 2, 16, requires_grad=True)
|
|
text = torch.randn(1, 2, 16, requires_grad=True)
|
|
alignment = cross_modal_alignment_loss(
|
|
visual, targets, assignments, text, targets, assignments
|
|
)
|
|
prototype = CanonicalFactPrototype(16, 3, 12, 7, 12)
|
|
prototype_value = canonical_prototype_loss(visual, targets, assignments, prototype)
|
|
loss = alignment + prototype_value
|
|
assert torch.isfinite(loss)
|
|
loss.backward()
|
|
assert visual.grad is not None
|