LSTM with CuDNN implementation which is only available for GPU. Note that currently only right padded data is supported, or the result will be polluted by the unmasked data which should be filtered. Args: inputs: Input tensor of LSTM layer. init_h: Initial state tensor for the cell o
(inputs, init_h, init_c, kernel, recurrent_kernel, bias, mask,
time_major, go_backwards)
| 1110 | |
| 1111 | |
| 1112 | def cudnn_lstm(inputs, init_h, init_c, kernel, recurrent_kernel, bias, mask, |
| 1113 | time_major, go_backwards): |
| 1114 | """LSTM with CuDNN implementation which is only available for GPU. |
| 1115 | |
| 1116 | Note that currently only right padded data is supported, or the result will be |
| 1117 | polluted by the unmasked data which should be filtered. |
| 1118 | |
| 1119 | Args: |
| 1120 | inputs: Input tensor of LSTM layer. |
| 1121 | init_h: Initial state tensor for the cell output. |
| 1122 | init_c: Initial state tensor for the cell hidden state. |
| 1123 | kernel: Weights for cell kernel. |
| 1124 | recurrent_kernel: Weights for cell recurrent kernel. |
| 1125 | bias: Weights for cell kernel bias and recurrent bias. Only recurrent bias |
| 1126 | is used in this case. |
| 1127 | mask: Boolean tensor for mask out the steps within sequence. |
| 1128 | time_major: Boolean, whether the inputs are in the format of |
| 1129 | [time, batch, feature] or [batch, time, feature]. |
| 1130 | go_backwards: Boolean (default False). If True, process the input sequence |
| 1131 | backwards and return the reversed sequence. |
| 1132 | |
| 1133 | Returns: |
| 1134 | last_output: Output tensor for the last timestep, which has shape |
| 1135 | [batch, units]. |
| 1136 | outputs: Output tensor for all timesteps, which has shape |
| 1137 | [batch, time, units]. |
| 1138 | state_0: The cell output, which has same shape as init_h. |
| 1139 | state_1: The cell hidden state, which has same shape as init_c. |
| 1140 | runtime: Constant string tensor which indicate real runtime hardware. This |
| 1141 | value is for testing purpose and should not be used by user. |
| 1142 | """ |
| 1143 | if not time_major and mask is None: |
| 1144 | inputs = array_ops.transpose(inputs, perm=(1, 0, 2)) |
| 1145 | seq_axis, batch_axis = (0, 1) |
| 1146 | else: |
| 1147 | seq_axis, batch_axis = (0, 1) if time_major else (1, 0) |
| 1148 | # For init_h and init_c, cuDNN expects one more dim of num_layers before or |
| 1149 | # after batch dim for time major or batch major inputs respectively |
| 1150 | init_h = array_ops.expand_dims(init_h, axis=seq_axis) |
| 1151 | init_c = array_ops.expand_dims(init_c, axis=seq_axis) |
| 1152 | |
| 1153 | weights = array_ops.split(kernel, 4, axis=1) |
| 1154 | weights += array_ops.split(recurrent_kernel, 4, axis=1) |
| 1155 | # CuDNN has an extra set of bias for inputs, we disable them (setting to 0), |
| 1156 | # so that mathematically it is same as the canonical LSTM implementation. |
| 1157 | full_bias = array_ops.concat((array_ops.zeros_like(bias), bias), 0) |
| 1158 | |
| 1159 | params = _canonical_to_params( |
| 1160 | weights=weights, |
| 1161 | biases=array_ops.split(full_bias, 8), |
| 1162 | shape=constant_op.constant([-1]), |
| 1163 | transpose_weights=True) |
| 1164 | |
| 1165 | if mask is not None: |
| 1166 | sequence_length = calculate_sequence_by_mask(mask, time_major) |
| 1167 | if go_backwards: |
| 1168 | # Three reversals are required. E.g., |
| 1169 | # normal input = [1, 2, 3, 0, 0] # where 0 need to be masked |
no test coverage detected