def _add_seq_rnn(self, cell_type: str):
"""Add the clinical picture inference module; implemented in as an RNN. """
with tf.variable_scope('sequence'):
# Add dropout on deltas
if self.dropout > 0:
self.delta_inputs = tf.keras.layers.Dropout(rate=self.dropout)(self.delta_inputs,
training=self.training)
# Concat observation_t and delta_t (deltas are already shifted by one)
if self.delta_combine == 'concat':
self.x = tf.concat([self.snapshot_encodings, self.delta_inputs], axis=-1, name='rnn_input_concat')
elif self.delta_combine == 'add':
self.x = self.snapshot_encodings + self.delta_inputs
else:
raise ValueError("Invalid delta combination method: %s" % self.delta_combine)
# Add dropout on concatenated inputs
if self.dropout > 0:
self.x = tf.keras.layers.Dropout(rate=self.dropout)(self.x, training=self.training)
_cell_types = {
# Original RAN from https://arxiv.org/abs/1705.07393
'RAN': rnn_cell.RANCell,
'RAN-LN': lambda units: rnn_cell.RANCell(units, normalize=True),
'VHRAN': lambda units: rnn_cell.VHRANCell(units, self.x.shape[-1], depth=self.rnn_highway_depth),
'VHRAN-LN': lambda units: rnn_cell.VHRANCell(units, self.x.shape[-1], depth=self.rnn_highway_depth,
normalize=True),
'RHN': lambda units: rnn_cell.RHNCell(units, self.x.shape[-1],
depth=self.rnn_highway_depth,
is_training=self.training),
'RHN-LN': lambda units: rnn_cell.RHNCell(units, self.x.shape[-1],
depth=self.rnn_highway_depth,
is_training=self.training,
normalize=True),
# Super secret simplified RAN variant from Eq. group (2) in https://arxiv.org/abs/1705.07393
# 'LRAN': lambda num_cells: rnn_cell.SimpleRANCell(self.x.shape[-1]),
# 'LRAN-LN': lambda num_cells: rnn_cell.SimpleRANCell(self.x.shape[-1], normalize=True),
'LSTM': tf.nn.rnn_cell.BasicLSTMCell,
'LSTM-LN': tf.contrib.rnn.LayerNormBasicLSTMCell,
'GRU': tf.nn.rnn_cell.GRUCell,
'GRU-LN': rnn_cell.LayerNormGRUCell
}
if cell_type not in _cell_types:
raise ValueError('unsupported cell type %s', cell_type)
self.cell_fn = _cell_types[cell_type]
if self.rnn_direction == 'bidirectional':
self.seq_final_output = layers.bidirectional_rnn_layer(self.cell_fn,
self.num_hidden, self.x, self.seq_lengths)
else:
self.seq_final_output = layers.rnn_layer(self.cell_fn, self.num_hidden, self.x, self.seq_lengths)
print('Final output:', self.seq_final_output)
# Even more fun dropout
if self.dropout > 0:
self.seq_final_output = \
tf.keras.layers.Dropout(rate=self.dropout)(self.seq_final_output, training=self.training)
def _rnn_encoder(model):
"""
:type model: modeling.BERTModel
"""
with tf.variable_scope('rnn_encoder'):
# Embed clinical observations
embedded_observations = layers.embedding_layer(model.observations, model.vocabulary_size,
model.embedding_size,
model.vocab_dropout,
training=model.training)
# Reshape to (batch * seq_len) x doc_len x embedding
flattened_embedded_obs = tf.reshape(embedded_observations,
[model.batch_size * model.max_seq_len,
model.max_snapshot_size,
model.embedding_size],
name='flat_emb_obs')
flattened_snapshot_sizes = tf.reshape(model.snapshot_sizes, [model.batch_size * model.max_seq_len],
name='flat_snapshot_sizes')
# Apply RNN to all documents in all batches
flattened_snapshot_encodings = layers.rnn_layer(cell_fn=cell_fn,
num_hidden=num_hidden,
inputs=flattened_embedded_obs,
lengths=flattened_snapshot_sizes,
return_interpretable_weights=False)
# Reshape back to (batch x seq_len x encoding_size)
return tf.reshape(flattened_snapshot_encodings,
[model.batch_size, model.max_seq_len, flattened_snapshot_encodings.shape[-1]],
name='rnn_snapshot_encoding')
def _add_seq_rnn(self, cell_type: str):
"""Add the clinical picture inference module; implemented in as an RNN. """
with tf.variable_scope('sequence'):
# Add dropout on deltas
if self.dropout > 0:
self.deltas = tf.layers.dropout(self.deltas,
rate=self.dropout,
training=self.training)
# Concat observation_t and delta_t (deltas are already shifted by one)
self.x = tf.concat([self.snapshot_encodings, self.deltas],
axis=-1,
name='rnn_input_concat')
# Add dropout on concatenated inputs
if self.dropout > 0:
self.x = tf.layers.dropout(self.x,
rate=self.dropout,
training=self.training)
_cell_types = {
# Original RAN from https://arxiv.org/abs/1705.07393
'RAN':
rnn_cell.RANCell,
'RAN-LN':
lambda num_cells: rnn_cell.RANCell(num_cells, normalize=True),
# Super secret simplified RAN variant from Eq. group (2) in https://arxiv.org/abs/1705.07393
'LRAN':
lambda num_cells: rnn_cell.SimpleRANCell(self.x.shape[-1]),
'LRAN-LN':
lambda num_cells: rnn_cell.SimpleRANCell(self.x.shape[-1],
normalize=True),
'LSTM':
tf.nn.rnn_cell.BasicLSTMCell,
'LSTM-LN':
tf.contrib.rnn.LayerNormBasicLSTMCell,
'GRU':
tf.nn.rnn_cell.GRUCell,
'GRU-LN':
rnn_cell.LayerNormGRUCell
}
if cell_type not in _cell_types:
raise ValueError('unsupported cell type %s', cell_type)
self.cell_fn = _cell_types[cell_type]
# Compute weights AND final RNN output if looking at RANv2 variants
if cell_type.startswith('IRAN'):
self.seq_final_output, self.rnn_weights = layers.rnn_layer(
self.cell_fn,
self.num_hidden,
self.x,
self.seq_lengths,
return_interpretable_weights=True)
else:
self.seq_final_output = layers.rnn_layer(
self.cell_fn,
self.num_hidden,
self.x,
self.seq_lengths,
return_interpretable_weights=False)
# Even more fun dropout
if self.dropout > 0:
self.seq_final_output = tf.layers.dropout(
self.seq_final_output,
rate=self.dropout,
training=self.training)
# Convert to sexy logits
self.logits = tf.layers.dense(self.seq_final_output,
units=self.num_classes,
activation=None,
name='class_logits')