def process(b):
# create emission graph
g_emissions = gtn.linear_graph(T, C, inputs.requires_grad)
cpu_data = inputs[b].cpu().contiguous()
g_emissions.set_weights(cpu_data.data_ptr())
# create transition graph
g_transitions = ASGLossFunction.create_transitions_graph(
transitions, calc_trans_grad)
# create force align criterion graph
g_fal = ASGLossFunction.create_force_align_graph(targets[b])
# compose the graphs
g_fal_fwd = gtn.forward_score(
gtn.intersect(gtn.intersect(g_fal, g_transitions),
g_emissions))
g_fcc_fwd = gtn.forward_score(
gtn.intersect(g_emissions, g_transitions))
g_loss = gtn.subtract(g_fcc_fwd, g_fal_fwd)
scale = 1.0
if reduction == "mean":
L = len(targets[b])
scale = 1.0 / L if L > 0 else scale
elif reduction != "none":
raise ValueError("invalid value for reduction '" +
str(reduction) + "'")
# Save for backward:
losses[b] = g_loss
scales[b] = scale
emissions_graphs[b] = g_emissions
transitions_graphs[b] = g_transitions
def process(b):
# Create emissions graph:
emissions = gtn.linear_graph(T, C, inputs.requires_grad)
cpu_data = inputs[b].cpu().contiguous()
emissions.set_weights(cpu_data.data_ptr())
target = make_chain_graph(targets[b])
target.arc_sort(True)
# Create token to grapheme decomposition graph
tokens_target = gtn.remove(gtn.project_output(gtn.compose(target, lexicon)))
tokens_target.arc_sort()
# Create alignment graph:
alignments = gtn.project_input(
gtn.remove(gtn.compose(tokens, tokens_target))
)
alignments.arc_sort()
# Add transition scores:
if transitions is not None:
alignments = gtn.intersect(transitions, alignments)
alignments.arc_sort()
loss = gtn.forward_score(gtn.intersect(emissions, alignments))
# Normalize if needed:
if transitions is not None:
norm = gtn.forward_score(gtn.intersect(emissions, transitions))
loss = gtn.subtract(loss, norm)
losses[b] = gtn.negate(loss)
# Save for backward:
if emissions.calc_grad:
emissions_graphs[b] = emissions
def process(b):
# create emission graph
g_emissions = gtn.linear_graph(T, C, log_probs.requires_grad)
cpu_data = log_probs[b].cpu().contiguous()
g_emissions.set_weights(cpu_data.data_ptr())
# create criterion graph
g_criterion = CTCLossFunction.create_ctc_graph(
targets[b], blank_idx)
# compose the graphs
g_loss = gtn.negate(
gtn.forward_score(gtn.intersect(g_emissions, g_criterion)))
scale = 1.0
if reduction == "mean":
L = len(targets[b])
scale = 1.0 / L if L > 0 else scale
elif reduction != "none":
raise ValueError("invalid value for reduction '" +
str(reduction) + "'")
# Save for backward:
losses[b] = g_loss
scales[b] = scale
emissions_graphs[b] = g_emissions
def test_simple_decomposition(self):
T = 5
tokens = ["a", "b", "ab", "ba", "aba"]
scores = torch.randn((1, T, len(tokens)), requires_grad=True)
labels = [[0, 1, 0]]
transducer = Transducer(tokens=tokens,
graphemes_to_idx={
"a": 0,
"b": 1
})
# Hand construct the alignment graph with all of the decompositions
alignments = gtn.Graph(False)
alignments.add_node(True)
# Add the path ['a', 'b', 'a']
alignments.add_node()
alignments.add_arc(0, 1, 0)
alignments.add_arc(1, 1, 0)
alignments.add_node()
alignments.add_arc(1, 2, 1)
alignments.add_arc(2, 2, 1)
alignments.add_node(False, True)
alignments.add_arc(2, 3, 0)
alignments.add_arc(3, 3, 0)
# Add the path ['a', 'ba']
alignments.add_node(False, True)
alignments.add_arc(1, 4, 3)
alignments.add_arc(4, 4, 3)
# Add the path ['ab', 'a']
alignments.add_node()
alignments.add_arc(0, 5, 2)
alignments.add_arc(5, 5, 2)
alignments.add_arc(5, 3, 0)
# Add the path ['aba']
alignments.add_node(False, True)
alignments.add_arc(0, 6, 4)
alignments.add_arc(6, 6, 4)
emissions = gtn.linear_graph(T, len(tokens), True)
emissions.set_weights(scores.data_ptr())
expected_loss = gtn.subtract(
gtn.forward_score(emissions),
gtn.forward_score(gtn.intersect(emissions, alignments)),
)
loss = transducer(scores, labels)
self.assertAlmostEqual(loss.item(), expected_loss.item(), places=5)
loss.backward()
gtn.backward(expected_loss)
expected_grad = torch.tensor(emissions.grad().weights_to_numpy())
expected_grad = expected_grad.view((1, T, len(tokens)))
self.assertTrue(
torch.allclose(scores.grad, expected_grad, rtol=1e-4, atol=1e-5))
def emissions_graph(emissions_vec, T, N, logprobs=False):
if not logprobs:
emissions_vec = [
-math.inf if e == 0 else math.log(e) for e in emissions_vec
]
g = gtn.linear_graph(T, N)
g.set_weights(emissions_vec)
return g
def time_indexed_func():
N1 = 100
N2 = 50
A1 = 20
A2 = 500
graphs1 = [gtn.linear_graph(N1, A1) for _ in range(B)]
graphs2 = [gtn.linear_graph(N2, A2) for _ in range(B)]
for g in graphs2:
for i in range(N2):
for j in range(A2):
g.add_arc(i, i, j)
out = [None] * B
def process(b):
out[b] = gtn.forward_score(gtn.compose(graphs1[b], graphs2[b]))
def indexed_func():
gtn.parallel_for(process, range(B))
time_func(indexed_func, 100, "parallel indexed python func")
def linearFstFromArray(weights):
seq_len, vocab_size = weights.shape
fst = gtn.linear_graph(seq_len,
vocab_size,
calc_grad=weights.requires_grad)
# Set FST weights
data = weights.contiguous()
fst.set_weights(data.data_ptr())
fst.arc_sort(olabel=True)
return fst
def time_forward_score():
graphs = [gtn.linear_graph(1000, 100) for _ in range(B)]
def fwd():
gtn.forward_score(graphs)
time_func(fwd, 100, "parallel forward_score Fwd")
out = gtn.forward_score(graphs)
def bwd():
gtn.backward(out, [True])
time_func(bwd, 100, "parallel forward_score bwd")
def time_compose():
N1 = 100
N2 = 50
A1 = 20
A2 = 500
graphs1 = [gtn.linear_graph(N1, A1) for _ in range(B)]
graphs2 = [gtn.linear_graph(N2, A2) for _ in range(B)]
for g in graphs2:
for i in range(N2):
for j in range(A2):
g.add_arc(i, i, j)
def fwd():
gtn.compose(graphs1, graphs2)
time_func(fwd, 20, "parallel compose Fwd")
out = gtn.compose(graphs1, graphs2)
def bwd():
gtn.backward(out, [True])
time_func(bwd, 20, "parallel compose bwd")
def test_linear_creation(self):
M = 5
N = 10
arr = [random.random() for _ in range(M * N)]
g = gtn.linear_graph(M, N)
g.set_weights(arr)
self.assertEqual(g.num_nodes(), M + 1)
self.assertEqual(g.num_arcs(), M * N)
self.assertEqual(g.labels_to_list(),
[j for _ in range(M) for j in range(N)])
weights = g.weights_to_list()
for i, w in enumerate(weights):
self.assertAlmostEqual(w, arr[i], places=5)
def forward_single(b):
emissions = gtn.linear_graph(T, C, inputs.requires_grad)
data = inputs[b].contiguous()
emissions.set_weights(data.data_ptr())
target = GTNLossFunction.make_target_graph(targets[b])
# Score the target:
target_score = gtn.forward_score(gtn.intersect(target, emissions))
# Normalization term:
norm = gtn.forward_score(emissions)
# Compute the loss:
loss = gtn.subtract(norm, target_score)
# Save state for backward:
losses[b] = loss
emissions_graphs[b] = emissions
def process(b):
for t in range(0, T - kernel_size + 1, stride):
input_graph = gtn.linear_graph(kernel_size, C, inputs.requires_grad)
window = cpu_inputs[b, t : t + kernel_size, :].contiguous()
input_graph.set_weights(window.data_ptr())
if viterbi:
window_outputs = [
gtn.viterbi_score(gtn.intersect(input_graph, kernel))
for kernel in kernels
]
else:
window_outputs = [
gtn.forward_score(gtn.intersect(input_graph, kernel))
for kernel in kernels
]
output_graphs[b].append(window_outputs)
# Save for backward:
if input_graph.calc_grad:
input_graphs[b].append(input_graph)
def process(b):
emissions = gtn.linear_graph(T, C, False)
cpu_data = outputs[b].cpu().contiguous()
emissions.set_weights(cpu_data.data_ptr())
if self.transitions is not None:
full_graph = gtn.intersect(emissions, self.transitions)
else:
full_graph = emissions
# Find the best path and remove back-off arcs:
path = gtn.remove(gtn.viterbi_path(full_graph))
# Left compose the viterbi path with the "alignment to token"
# transducer to get the outputs:
path = gtn.compose(path, self.tokens)
# When there are ambiguous paths (allow_repeats is true), we take
# the shortest:
path = gtn.viterbi_path(path)
path = gtn.remove(gtn.project_output(path))
paths[b] = path.labels_to_list()
def process(b):
# create emission graph
g_emissions = gtn.linear_graph(T, C, False)
cpu_data = outputs[b].cpu().contiguous()
g_emissions.set_weights(cpu_data.data_ptr())
# create transition graph
g_transitions = utils.ASGLossFunction.create_transitions_graph(
self.transitions)
g_path = gtn.viterbi_path(gtn.intersect(g_emissions,
g_transitions))
prediction = g_path.labels_to_list()
collapsed_prediction = [p for p, _ in groupby(prediction)]
if self.garbage_idx is not None:
# remove garbage tokens
collapsed_prediction = [
p for p in collapsed_prediction if p != self.garbage_idx
]
predictions[b] = utils.unpack_replabels(collapsed_prediction,
self.num_replabels)
tokens = token_graph(word_pieces)
gtn.draw(tokens, "tokens.pdf", idx_to_wp, idx_to_wp)
# Recognizes "abc":
abc = gtn.Graph(False)
abc.add_node(True)
abc.add_node()
abc.add_node()
abc.add_node(False, True)
abc.add_arc(0, 1, let_to_idx["a"])
abc.add_arc(1, 2, let_to_idx["b"])
abc.add_arc(2, 3, let_to_idx["c"])
gtn.draw(abc, "abc.pdf", idx_to_let)
# Compute the decomposition graph for "abc":
abc_decomps = gtn.remove(gtn.project_output(gtn.compose(abc, lex)))
gtn.draw(abc_decomps, "abc_decomps.pdf", idx_to_wp, idx_to_wp)
# Compute the alignment graph for "abc":
abc_alignments = gtn.project_input(
gtn.remove(gtn.compose(tokens, abc_decomps)))
gtn.draw(abc_alignments, "abc_alignments.pdf", idx_to_wp)
# From here we can use the alignment graph with an emissions graph and
# transitions graphs to compute the sequence level criterion:
emissions = gtn.linear_graph(10, len(word_pieces), True)
loss = gtn.subtract(
gtn.forward_score(emissions),
gtn.forward_score(gtn.intersect(emissions, abc_alignments)))
print(f"Loss is {loss.item():.2f}")
