Skip to content

GitLab

Commit 9f9b4307 authored by Le Hou's avatar Le Hou Committed by A. Unique TensorFlower
Browse files

Save layers for access in EncoderScaffold. 

PiperOrigin-RevId: 339345258
parent 284d7783
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 0 deletions
official/nlp/modeling/networks/encoder_scaffold.py
View file @ 9f9b4307
...@@ -114,6 +114,9 @@ class EncoderScaffold(tf.keras.Model): ...@@ -114,6 +114,9 @@ class EncoderScaffold(tf.keras.Model):
inputs = embedding_network.inputs inputs = embedding_network.inputs
embeddings, attention_mask = embedding_network(inputs) embeddings, attention_mask = embedding_network(inputs)
embedding_layer = None embedding_layer = None
position_embedding_layer = None
type_embedding_layer = None
embedding_norm_layer = None
else: else:
embedding_network = None embedding_network = None
seq_length = embedding_cfg.get('seq_length', None) seq_length = embedding_cfg.get('seq_length', None)
...@@ -224,6 +227,10 @@ class EncoderScaffold(tf.keras.Model): ...@@ -224,6 +227,10 @@ class EncoderScaffold(tf.keras.Model):
self._embedding_layer = embedding_layer self._embedding_layer = embedding_layer
self._embedding_network = embedding_network self._embedding_network = embedding_network
self._position_embedding_layer = position_embedding_layer
self._type_embedding_layer = type_embedding_layer
self._embedding_norm_layer = embedding_norm_layer
self._embedding_network = embedding_network
self._hidden_layers = hidden_layers self._hidden_layers = hidden_layers
self._pooler_layer = pooler_layer self._pooler_layer = pooler_layer
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment