RNN estimator with target predictor function on top.
(x, y)
| 329 | """ |
| 330 | |
| 331 | def rnn_estimator(x, y): |
| 332 | """RNN estimator with target predictor function on top.""" |
| 333 | x = input_op_fn(x) |
| 334 | if cell_type == 'rnn': |
| 335 | cell_fn = contrib_rnn.BasicRNNCell |
| 336 | elif cell_type == 'gru': |
| 337 | cell_fn = contrib_rnn.GRUCell |
| 338 | elif cell_type == 'lstm': |
| 339 | cell_fn = functools.partial( |
| 340 | contrib_rnn.BasicLSTMCell, state_is_tuple=False) |
| 341 | else: |
| 342 | raise ValueError('cell_type {} is not supported. '.format(cell_type)) |
| 343 | # TODO(ipolosukhin): state_is_tuple=False is deprecated |
| 344 | if bidirectional: |
| 345 | # forward direction cell |
| 346 | fw_cell = lambda: cell_fn(rnn_size) |
| 347 | bw_cell = lambda: cell_fn(rnn_size) |
| 348 | # attach attention cells if specified |
| 349 | if attn_length is not None: |
| 350 | def attn_fw_cell(): |
| 351 | return contrib_rnn.AttentionCellWrapper( |
| 352 | fw_cell(), |
| 353 | attn_length=attn_length, |
| 354 | attn_size=attn_size, |
| 355 | attn_vec_size=attn_vec_size, |
| 356 | state_is_tuple=False) |
| 357 | |
| 358 | def attn_bw_cell(): |
| 359 | return contrib_rnn.AttentionCellWrapper( |
| 360 | bw_cell(), |
| 361 | attn_length=attn_length, |
| 362 | attn_size=attn_size, |
| 363 | attn_vec_size=attn_vec_size, |
| 364 | state_is_tuple=False) |
| 365 | else: |
| 366 | attn_fw_cell = fw_cell |
| 367 | attn_bw_cell = bw_cell |
| 368 | |
| 369 | rnn_fw_cell = contrib_rnn.MultiRNNCell( |
| 370 | [attn_fw_cell() for _ in range(num_layers)], state_is_tuple=False) |
| 371 | # backward direction cell |
| 372 | rnn_bw_cell = contrib_rnn.MultiRNNCell( |
| 373 | [attn_bw_cell() for _ in range(num_layers)], state_is_tuple=False) |
| 374 | # pylint: disable=unexpected-keyword-arg, no-value-for-parameter |
| 375 | _, encoding = bidirectional_rnn( |
| 376 | rnn_fw_cell, |
| 377 | rnn_bw_cell, |
| 378 | x, |
| 379 | dtype=dtypes.float32, |
| 380 | sequence_length=sequence_length, |
| 381 | initial_state_fw=initial_state, |
| 382 | initial_state_bw=initial_state) |
| 383 | else: |
| 384 | rnn_cell = lambda: cell_fn(rnn_size) |
| 385 | |
| 386 | if attn_length is not None: |
| 387 | def attn_rnn_cell(): |
| 388 | return contrib_rnn.AttentionCellWrapper( |
nothing calls this directly
no test coverage detected