Sum state log probs to ilabel log probs using unique label indices.
(labels, num_labels, states, unique)
| 543 | |
| 544 | # pylint: disable=redefined-outer-name |
| 545 | def _state_to_olabel_unique(labels, num_labels, states, unique): |
| 546 | """Sum state log probs to ilabel log probs using unique label indices.""" |
| 547 | |
| 548 | num_label_states = _get_dim(labels, 1) + 1 |
| 549 | label_states = states[:, :, 1:num_label_states] |
| 550 | blank_states = states[:, :, num_label_states:] |
| 551 | |
| 552 | unique_y, unique_idx = unique |
| 553 | mul_reduce = _sum_states(unique_idx, label_states) |
| 554 | |
| 555 | num_frames = states.shape[0] |
| 556 | batch_size = states.shape[1] |
| 557 | num_states = num_label_states - 1 |
| 558 | batch_state_major = array_ops.transpose(mul_reduce, perm=[1, 2, 0]) |
| 559 | batch_state_major = array_ops.reshape(batch_state_major, |
| 560 | [batch_size * num_states, num_frames]) |
| 561 | batch_offset = math_ops.range(batch_size, dtype=unique_y.dtype) * num_labels |
| 562 | indices = unique_y + array_ops.expand_dims(batch_offset, axis=-1) |
| 563 | indices = array_ops.reshape(indices, [-1, 1]) |
| 564 | scatter = array_ops.scatter_nd( |
| 565 | indices=indices, |
| 566 | updates=batch_state_major, |
| 567 | shape=[batch_size * num_labels, num_frames]) |
| 568 | scatter = array_ops.reshape(scatter, [batch_size, num_labels, num_frames]) |
| 569 | scatter = array_ops.where( |
| 570 | math_ops.equal(scatter, 0.0), |
| 571 | array_ops.fill(array_ops.shape(scatter), math_ops.log(0.0)), scatter) |
| 572 | label_olabels = array_ops.transpose(scatter, [2, 0, 1]) |
| 573 | label_olabels = label_olabels[:, :, 1:] |
| 574 | |
| 575 | blank_olabels = math_ops.reduce_logsumexp(blank_states, axis=2, keepdims=True) |
| 576 | |
| 577 | return array_ops.concat([blank_olabels, label_olabels], axis=-1) |
| 578 | |
| 579 | |
| 580 | def ctc_loss_and_grad(logits, labels, label_length, logit_length, unique=None): |
no test coverage detected