2373

pending completion

Build # 2373

Build Type

Pull #810

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/google/trax/commit/<a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/281b1f93bc3b51fdd542e2546fbf7961726b4a1d">281b1f93b">&lt;a href=&quot;https://github.com/google/trax/commit/</a><a class="double-link" href="https://github.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/281b1f93bc3b51fdd542e2546fbf7961726b4a1d">281b1f93b</a><a href="https://github.com/google/trax/commit/281b1f93bc3b51fdd542e2546fbf7961726b4a1d">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/google/trax/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/281b1f93bc3b51fdd542e2546fbf7961726b4a1d&quot;&gt;281b1f93b&lt;/a&gt;&lt;a href=&quot;https://github.com/google/trax/commit/281b1f93bc3b51fdd542e2546fbf7961726b4a1d&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/google/trax/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/google/trax/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp... (continued)

Pull Request Pull Request #810: Change order of arguments for the Embedding layer to match Keras and pytorch.

Run Details

3 of 3 new or added lines in 2 files covered. (100.0%)

2623 of 9402 relevant lines covered (27.9%)

0.28 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

23.08

/trax/models/research/bert.py

# coding=utf-8
# Copyright 2020 The Trax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""BERT."""

import jax
import tensorflow as tf

from trax import fastmath
from trax import layers as tl
from trax.fastmath import numpy as np

# pylint: disable=invalid-name


class AddBias(tl.Layer):

  def forward(self, inputs):
    x = inputs
    return x + self.weights

  def init_weights_and_state(self, input_signature):
    self.weights = np.zeros(input_signature.shape[-1])


def BERTClassifierHead(n_classes):
  return tl.Serial([
      tl.Select([0], n_in=2),
      tl.Dense(n_classes,
               kernel_initializer=tl.RandomNormalInitializer(0.02),
               bias_initializer=tl.RandomNormalInitializer(1e-6),
              ),
      tl.LogSoftmax(),
  ])


def BERTRegressionHead():
  return tl.Serial([
      tl.Select([0], n_in=2),
      tl.Dense(1,
               kernel_initializer=tl.RandomNormalInitializer(0.02),
               bias_initializer=tl.RandomNormalInitializer(1e-6),
              ),
  ])


# TODO(kitaev): regression head, masked LM head


def BERT(d_model=768,
         vocab_size=30522,
         max_len=512,
         type_vocab_size=2,
         n_heads=12,
         d_ff=3072,
         n_layers=12,
         head=None,
         init_checkpoint=None,
         mode='eval',
        ):
  """BERT (default hparams are for bert-base-uncased)."""
  layer_norm_eps = 1e-12
  d_head = d_model // n_heads

  word_embeddings = tl.Embedding(vocab_size, d_model)
  type_embeddings = tl.Embedding(type_vocab_size, d_model)
  position_embeddings = tl.PositionalEncoding(max_len, mode=mode)
  embeddings = [
      tl.Select([0, 1, 0], n_in=3),  # Drops 'idx' input.
      tl.Parallel(
          word_embeddings,
          type_embeddings,
          [tl.PaddingMask(),
           tl.Fn('Squeeze', lambda x: np.squeeze(x, (1, 2)), n_out=1)]
      ),
      tl.Add(),
      position_embeddings,
      tl.LayerNorm(epsilon=layer_norm_eps),
  ]

  encoder = []
  for _ in range(n_layers):
    attn = tl.SelfAttention(n_heads=n_heads, d_qk=d_head, d_v=d_head,
                            bias=True, masked=True, mode=mode)
    feed_forward = [
        tl.Dense(d_ff),
        tl.Gelu(),
        tl.Dense(d_model)
    ]
    encoder += [
        tl.Select([0, 1, 1]),  # Save a copy of the mask
        tl.Residual(attn, AddBias()),  # pylint: disable=no-value-for-parameter
        tl.LayerNorm(epsilon=layer_norm_eps),
        tl.Residual(*feed_forward),
        tl.LayerNorm(epsilon=layer_norm_eps),
    ]

  encoder += [tl.Select([0], n_in=2)]  # Drop the mask

  pooler = [
      tl.Fn('', lambda x: (x[:, 0, :], x), n_out=2),
      tl.Dense(d_model),
      tl.Tanh(),
  ]

  init_checkpoint = init_checkpoint if mode == 'train' else None
  bert = PretrainedBERT(
      embeddings + encoder + pooler, init_checkpoint=init_checkpoint)

  if head is not None:
    bert = tl.Serial(bert, head())

  return bert


class PretrainedBERT(tl.Serial):
  """Wrapper that always initializes weights from a pre-trained checkpoint."""

  def __init__(self, *sublayers, init_checkpoint=None):
    super().__init__(*sublayers)

    # TODO(kitaev): Support shorthand model names in the trax OSS release
    if init_checkpoint == 'bert-base-uncased':
      raise NotImplementedError(
          'Please manually specify the path to bert_model.ckpt')
    self.init_checkpoint = init_checkpoint

  def new_weights(self, input_signature):
    weights = super().new_weights(input_signature)
    if self.init_checkpoint is None:
      return weights

    print('Loading pre-trained weights from', self.init_checkpoint)
    ckpt = tf.train.load_checkpoint(self.init_checkpoint)

    def reshape_qkv(name):
      x = ckpt.get_tensor(name)
      return x.reshape((x.shape[0], -1, 64)).swapaxes(0, 1)
    def reshape_o(name):
      x = ckpt.get_tensor(name)
      return x.reshape((-1, 64, x.shape[-1]))
    def reshape_bias(name):
      x = ckpt.get_tensor(name)
      return x.reshape((-1, 64))

    new_w = [
        ckpt.get_tensor('bert/embeddings/word_embeddings'),
        ckpt.get_tensor('bert/embeddings/token_type_embeddings'),
        ckpt.get_tensor('bert/embeddings/position_embeddings')[None, ...],
        ckpt.get_tensor('bert/embeddings/LayerNorm/gamma'),
        ckpt.get_tensor('bert/embeddings/LayerNorm/beta'),
    ]

    for i in range(12):  # 12 layers
      new_w += [
          reshape_qkv(f'bert/encoder/layer_{i}/attention/self/query/kernel'),
          reshape_qkv(f'bert/encoder/layer_{i}/attention/self/key/kernel'),
          reshape_qkv(f'bert/encoder/layer_{i}/attention/self/value/kernel'),
          reshape_o(f'bert/encoder/layer_{i}/attention/output/dense/kernel'),
          reshape_bias(f'bert/encoder/layer_{i}/attention/self/query/bias'),
          reshape_bias(f'bert/encoder/layer_{i}/attention/self/key/bias'),
          reshape_bias(f'bert/encoder/layer_{i}/attention/self/value/bias'),
          ckpt.get_tensor(
              f'bert/encoder/layer_{i}/attention/output/dense/bias'),
          ckpt.get_tensor(
              f'bert/encoder/layer_{i}/attention/output/LayerNorm/gamma'),
          ckpt.get_tensor(
              f'bert/encoder/layer_{i}/attention/output/LayerNorm/beta'),
          ckpt.get_tensor(f'bert/encoder/layer_{i}/intermediate/dense/kernel'),
          ckpt.get_tensor(f'bert/encoder/layer_{i}/intermediate/dense/bias'),
          ckpt.get_tensor(f'bert/encoder/layer_{i}/output/dense/kernel'),
          ckpt.get_tensor(f'bert/encoder/layer_{i}/output/dense/bias'),
          ckpt.get_tensor(f'bert/encoder/layer_{i}/output/LayerNorm/gamma'),
          ckpt.get_tensor(f'bert/encoder/layer_{i}/output/LayerNorm/beta'),
      ]

    new_w += [
        ckpt.get_tensor('bert/pooler/dense/kernel'),
        ckpt.get_tensor('bert/pooler/dense/bias'),
    ]

    for a, b in zip(fastmath.tree_leaves(weights), new_w):
      assert a.shape == b.shape, (
          f'Expected shape {a.shape}, got shape {b.shape}')
    weights = jax.tree_unflatten(jax.tree_structure(weights), new_w)
    move_to_device = jax.jit(lambda x: x)
    weights = jax.tree_map(move_to_device, weights)
    return weights

1	# coding=utf-8
2	# Copyright 2020 The Trax Authors.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	# http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	# Lint as: python3
17	"""BERT."""	1✔
18
19	import jax	1✔
20	import tensorflow as tf	1✔
21
22	from trax import fastmath	1✔
23	from trax import layers as tl	1✔
24	from trax.fastmath import numpy as np	1✔
25
26	# pylint: disable=invalid-name
27
28
29	class AddBias(tl.Layer):	1✔
30
31	def forward(self, inputs):	1✔
32	x = inputs	×
33	return x + self.weights	×
34
35	def init_weights_and_state(self, input_signature):	1✔
36	self.weights = np.zeros(input_signature.shape[-1])	×
37
38
39	def BERTClassifierHead(n_classes):	1✔
40	return tl.Serial([	×
41	tl.Select([0], n_in=2),
42	tl.Dense(n_classes,
43	kernel_initializer=tl.RandomNormalInitializer(0.02),
44	bias_initializer=tl.RandomNormalInitializer(1e-6),
45	),
46	tl.LogSoftmax(),
47	])
48
49
50	def BERTRegressionHead():	1✔
51	return tl.Serial([	×
52	tl.Select([0], n_in=2),
53	tl.Dense(1,
54	kernel_initializer=tl.RandomNormalInitializer(0.02),
55	bias_initializer=tl.RandomNormalInitializer(1e-6),
56	),
57	])
58
59
60	# TODO(kitaev): regression head, masked LM head
61
62
63	def BERT(d_model=768,	1✔
64	vocab_size=30522,
65	max_len=512,
66	type_vocab_size=2,
67	n_heads=12,
68	d_ff=3072,
69	n_layers=12,
70	head=None,
71	init_checkpoint=None,
72	mode='eval',
73	):
74	"""BERT (default hparams are for bert-base-uncased)."""
75	layer_norm_eps = 1e-12	×
76	d_head = d_model // n_heads	×
77
78	word_embeddings = tl.Embedding(vocab_size, d_model)	×
79	type_embeddings = tl.Embedding(type_vocab_size, d_model)	×
80	position_embeddings = tl.PositionalEncoding(max_len, mode=mode)	×
81	embeddings = [	×
82	tl.Select([0, 1, 0], n_in=3), # Drops 'idx' input.
83	tl.Parallel(
84	word_embeddings,
85	type_embeddings,
86	[tl.PaddingMask(),
87	tl.Fn('Squeeze', lambda x: np.squeeze(x, (1, 2)), n_out=1)]
88	),
89	tl.Add(),
90	position_embeddings,
91	tl.LayerNorm(epsilon=layer_norm_eps),
92	]
93
94	encoder = []	×
95	for _ in range(n_layers):	×
96	attn = tl.SelfAttention(n_heads=n_heads, d_qk=d_head, d_v=d_head,	×
97	bias=True, masked=True, mode=mode)
98	feed_forward = [	×
99	tl.Dense(d_ff),
100	tl.Gelu(),
101	tl.Dense(d_model)
102	]
103	encoder += [	×
104	tl.Select([0, 1, 1]), # Save a copy of the mask
105	tl.Residual(attn, AddBias()), # pylint: disable=no-value-for-parameter
106	tl.LayerNorm(epsilon=layer_norm_eps),
107	tl.Residual(*feed_forward),
108	tl.LayerNorm(epsilon=layer_norm_eps),
109	]
110
111	encoder += [tl.Select([0], n_in=2)] # Drop the mask	×
112
113	pooler = [	×
114	tl.Fn('', lambda x: (x[:, 0, :], x), n_out=2),
115	tl.Dense(d_model),
116	tl.Tanh(),
117	]
118
119	init_checkpoint = init_checkpoint if mode == 'train' else None	×
120	bert = PretrainedBERT(	×
121	embeddings + encoder + pooler, init_checkpoint=init_checkpoint)
122
123	if head is not None:	×
124	bert = tl.Serial(bert, head())	×
125
126	return bert	×
127
128
129	class PretrainedBERT(tl.Serial):	1✔
130	"""Wrapper that always initializes weights from a pre-trained checkpoint."""
131
132	def __init__(self, *sublayers, init_checkpoint=None):	1✔
133	super().__init__(*sublayers)	×
134
135	# TODO(kitaev): Support shorthand model names in the trax OSS release
136	if init_checkpoint == 'bert-base-uncased':	×
137	raise NotImplementedError(
138	'Please manually specify the path to bert_model.ckpt')
139	self.init_checkpoint = init_checkpoint	×
140
141	def new_weights(self, input_signature):	1✔
142	weights = super().new_weights(input_signature)	×
143	if self.init_checkpoint is None:	×
144	return weights	×
145
146	print('Loading pre-trained weights from', self.init_checkpoint)	×
147	ckpt = tf.train.load_checkpoint(self.init_checkpoint)	×
148
149	def reshape_qkv(name):	×
150	x = ckpt.get_tensor(name)	×
151	return x.reshape((x.shape[0], -1, 64)).swapaxes(0, 1)	×
152	def reshape_o(name):	×
153	x = ckpt.get_tensor(name)	×
154	return x.reshape((-1, 64, x.shape[-1]))	×
155	def reshape_bias(name):	×
156	x = ckpt.get_tensor(name)	×
157	return x.reshape((-1, 64))	×
158
159	new_w = [	×
160	ckpt.get_tensor('bert/embeddings/word_embeddings'),
161	ckpt.get_tensor('bert/embeddings/token_type_embeddings'),
162	ckpt.get_tensor('bert/embeddings/position_embeddings')[None, ...],
163	ckpt.get_tensor('bert/embeddings/LayerNorm/gamma'),
164	ckpt.get_tensor('bert/embeddings/LayerNorm/beta'),
165	]
166
167	for i in range(12): # 12 layers	×
168	new_w += [	×
169	reshape_qkv(f'bert/encoder/layer_{i}/attention/self/query/kernel'),
170	reshape_qkv(f'bert/encoder/layer_{i}/attention/self/key/kernel'),
171	reshape_qkv(f'bert/encoder/layer_{i}/attention/self/value/kernel'),
172	reshape_o(f'bert/encoder/layer_{i}/attention/output/dense/kernel'),
173	reshape_bias(f'bert/encoder/layer_{i}/attention/self/query/bias'),
174	reshape_bias(f'bert/encoder/layer_{i}/attention/self/key/bias'),
175	reshape_bias(f'bert/encoder/layer_{i}/attention/self/value/bias'),
176	ckpt.get_tensor(
177	f'bert/encoder/layer_{i}/attention/output/dense/bias'),
178	ckpt.get_tensor(
179	f'bert/encoder/layer_{i}/attention/output/LayerNorm/gamma'),
180	ckpt.get_tensor(
181	f'bert/encoder/layer_{i}/attention/output/LayerNorm/beta'),
182	ckpt.get_tensor(f'bert/encoder/layer_{i}/intermediate/dense/kernel'),
183	ckpt.get_tensor(f'bert/encoder/layer_{i}/intermediate/dense/bias'),
184	ckpt.get_tensor(f'bert/encoder/layer_{i}/output/dense/kernel'),
185	ckpt.get_tensor(f'bert/encoder/layer_{i}/output/dense/bias'),
186	ckpt.get_tensor(f'bert/encoder/layer_{i}/output/LayerNorm/gamma'),
187	ckpt.get_tensor(f'bert/encoder/layer_{i}/output/LayerNorm/beta'),
188	]
189
190	new_w += [	×
191	ckpt.get_tensor('bert/pooler/dense/kernel'),
192	ckpt.get_tensor('bert/pooler/dense/bias'),
193	]
194
195	for a, b in zip(fastmath.tree_leaves(weights), new_w):	×
196	assert a.shape == b.shape, (	×
197	f'Expected shape {a.shape}, got shape {b.shape}')
198	weights = jax.tree_unflatten(jax.tree_structure(weights), new_w)	×
199	move_to_device = jax.jit(lambda x: x)	×
200	weights = jax.tree_map(move_to_device, weights)	×
201	return weights	×

google / trax / 2373

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous