GRU with CuDNN implementation which is only available for GPU.
(inputs, init_h, kernel, recurrent_kernel, bias, mask, time_major,
go_backwards)
| 485 | |
| 486 | |
| 487 | def cudnn_gru(inputs, init_h, kernel, recurrent_kernel, bias, mask, time_major, |
| 488 | go_backwards): |
| 489 | """GRU with CuDNN implementation which is only available for GPU.""" |
| 490 | if not time_major and mask is None: |
| 491 | inputs = array_ops.transpose(inputs, perm=(1, 0, 2)) |
| 492 | seq_axis, batch_axis = (0, 1) |
| 493 | else: |
| 494 | seq_axis, batch_axis = (0, 1) if time_major else (1, 0) |
| 495 | # For init_h, cuDNN expects one more dim of num_layers before or after batch |
| 496 | # dim for time major or batch major inputs respectively |
| 497 | init_h = array_ops.expand_dims(init_h, axis=seq_axis) |
| 498 | |
| 499 | weights = array_ops.split(kernel, 3, axis=1) |
| 500 | weights += array_ops.split(recurrent_kernel, 3, axis=1) |
| 501 | # Note that the bias was initialized as shape (2, 3 * units), flat it into |
| 502 | # (6 * units) |
| 503 | bias = array_ops.split(K.flatten(bias), 6) |
| 504 | # Note that the gate order for CuDNN is different from the canonical format. |
| 505 | # canonical format is [z, r, h], whereas CuDNN is [r, z, h]. The swap need to |
| 506 | # be done for kernel, recurrent_kernel, input_bias, recurrent_bias. |
| 507 | # z is update gate weights. |
| 508 | # r is reset gate weights. |
| 509 | # h is output gate weights. |
| 510 | weights[0], weights[1] = weights[1], weights[0] |
| 511 | weights[3], weights[4] = weights[4], weights[3] |
| 512 | bias[0], bias[1] = bias[1], bias[0] |
| 513 | bias[3], bias[4] = bias[4], bias[3] |
| 514 | |
| 515 | params = _canonical_to_params( |
| 516 | weights=weights, |
| 517 | biases=bias, |
| 518 | shape=constant_op.constant([-1]), |
| 519 | transpose_weights=True) |
| 520 | |
| 521 | if mask is not None: |
| 522 | sequence_length = calculate_sequence_by_mask(mask, time_major) |
| 523 | if go_backwards: |
| 524 | # Three reversals are required. E.g., |
| 525 | # normal input = [1, 2, 3, 0, 0] # where 0 need to be masked |
| 526 | # reversed_input_to_cudnn = [3, 2, 1, 0, 0] |
| 527 | # output_from_cudnn = [6, 5, 4, 0, 0] |
| 528 | # expected_output = [0, 0, 6, 5 ,4] |
| 529 | inputs = array_ops.reverse_sequence_v2( |
| 530 | inputs, sequence_length, seq_axis=seq_axis, batch_axis=batch_axis) |
| 531 | outputs, h, _, _, _ = gen_cudnn_rnn_ops.cudnn_rnnv3( |
| 532 | inputs, |
| 533 | input_h=init_h, |
| 534 | input_c=0, |
| 535 | params=params, |
| 536 | is_training=True, |
| 537 | rnn_mode='gru', |
| 538 | sequence_lengths=sequence_length, |
| 539 | time_major=time_major) |
| 540 | if go_backwards: |
| 541 | outputs = array_ops.reverse_sequence_v2( |
| 542 | outputs, sequence_length, seq_axis=seq_axis, batch_axis=batch_axis) |
| 543 | outputs = array_ops.reverse(outputs, axis=[seq_axis]) |
| 544 | else: |
no test coverage detected