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}")