Skip to content

Commit 8e11893

Browse files
wwkongtensorflower-gardener
wwkong
authored and
tensorflower-gardener
committed
Implement and test a registry function for tfm.nlp.layers.EinsumDense + small formatting fixes.
PiperOrigin-RevId: 568854078
1 parent 8b52ba2 commit 8e11893

File tree

6 files changed

+442
-5
lines changed

6 files changed

+442
-5
lines changed

‎tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ py_library(
1313
name = "einsum_utils",
1414
srcs = ["einsum_utils.py"],
1515
srcs_version = "PY3",
16+
deps = ["//tensorflow_privacy/privacy/fast_gradient_clipping:common_manip_utils"],
1617
)
1718

1819
py_test(
@@ -24,6 +25,33 @@ py_test(
2425
deps = [":einsum_utils"],
2526
)
2627

28+
py_library(
29+
name = "einsum_dense",
30+
srcs = ["einsum_dense.py"],
31+
srcs_version = "PY3",
32+
deps = [
33+
":einsum_utils",
34+
"//tensorflow_privacy/privacy/fast_gradient_clipping:common_manip_utils",
35+
"//tensorflow_privacy/privacy/fast_gradient_clipping:type_aliases",
36+
],
37+
)
38+
39+
py_test(
40+
name = "einsum_dense_test",
41+
size = "large",
42+
srcs = ["einsum_dense_test.py"],
43+
python_version = "PY3",
44+
shard_count = 12,
45+
srcs_version = "PY3",
46+
deps = [
47+
":dense",
48+
":einsum_dense",
49+
"//tensorflow_privacy/privacy/fast_gradient_clipping:clip_grads",
50+
"//tensorflow_privacy/privacy/fast_gradient_clipping:common_test_utils",
51+
"//tensorflow_privacy/privacy/fast_gradient_clipping:layer_registry",
52+
],
53+
)
54+
2755
py_library(
2856
name = "dense",
2957
srcs = ["dense.py"],

‎tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright 2023, The TensorFlow Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""Fast clipping function for `tfm.nlp.layers.EinsumDense`."""
15+
16+
from collections.abc import Mapping, Sequence
17+
from typing import Any, Optional
18+
import tensorflow as tf
19+
from tensorflow_privacy.privacy.fast_gradient_clipping import type_aliases
20+
from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import einsum_utils
21+
22+
23+
def einsum_layer_computation(
24+
layer_instance: tf.keras.layers.EinsumDense,
25+
input_args: Sequence[Any],
26+
input_kwargs: Mapping[str, Any],
27+
tape: tf.GradientTape,
28+
num_microbatches: Optional[tf.Tensor] = None,
29+
) -> type_aliases.RegistryFunctionOutput:
30+
"""Registry function for `tf.keras.layers.EinsumDense`.
31+
32+
For the technical details, see the documentation of
33+
`einsum_utils.compute_fast_einsum_gradient_norm()`.
34+
35+
Args:
36+
layer_instance: A `tf.keras.layers.EinsumDense` instance.
37+
input_args: See `dense_layer_computation()` in `dense.py`.
38+
input_kwargs: See `dense_layer_computation()` in `dense.py`.
39+
tape: See `dense_layer_computation()` in `dense.py`.
40+
num_microbatches: See `dense_layer_computation()` in `dense.py`.
41+
42+
Returns:
43+
See `dense_layer_computation()` in `dense.py`.
44+
"""
45+
if input_kwargs:
46+
raise ValueError("EinsumDense layer calls should not receive kwargs.")
47+
del input_kwargs
48+
if len(input_args) != 1:
49+
raise ValueError("Only layer inputs of length 1 are permitted.")
50+
orig_activation = layer_instance.activation
51+
# Some activation functions may not apply a transform to the elements of the
52+
# output individually (which is needed for the fast clipping trick to work).
53+
# To avoid this case, we watch the variables that are only generated by the
54+
# linear transformation of the `EinsumDense` layer instance.
55+
layer_instance.activation = None
56+
base_vars = layer_instance(*input_args)
57+
tape.watch(base_vars)
58+
layer_instance.activation = orig_activation
59+
outputs = orig_activation(base_vars) if orig_activation else base_vars
60+
61+
def sqr_norm_fn(grads):
62+
return einsum_utils.compute_fast_einsum_squared_gradient_norm(
63+
layer_instance.equation,
64+
input_args[0],
65+
grads,
66+
layer_instance.bias_axes,
67+
num_microbatches,
68+
)
69+
70+
return base_vars, outputs, sqr_norm_fn

‎tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense_test.py

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
# Copyright 2023, The TensorFlow Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from absl.testing import parameterized
16+
import tensorflow as tf
17+
import tensorflow_models as tfm
18+
from tensorflow_privacy.privacy.fast_gradient_clipping import common_test_utils
19+
from tensorflow_privacy.privacy.fast_gradient_clipping import layer_registry
20+
from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import einsum_dense
21+
22+
23+
def get_einsum_layer_generators():
24+
def pure_einsum_layer(equation, output_dims, bias_axes):
25+
return tf.keras.layers.EinsumDense(
26+
equation, output_dims, bias_axes=bias_axes
27+
)
28+
29+
def sigmoid_einsum_layer(equation, output_dims, bias_axes):
30+
return tf.keras.layers.EinsumDense(
31+
equation, output_dims, bias_axes=bias_axes, activation='sigmoid'
32+
)
33+
34+
return {
35+
'pure_einsum': pure_einsum_layer,
36+
'sigmoid_einsum': sigmoid_einsum_layer,
37+
}
38+
39+
40+
def get_einsum_parameter_tuples():
41+
# (equation, input_dims, output_dims, bias_axes)
42+
return [
43+
# Case C1.
44+
('ab,bc->ac', [2], [3], None),
45+
('ab,bc->ac', [2], [3], 'c'),
46+
('abc,cd->abd', [2, 3], [2, 4], None),
47+
('abc,cd->abd', [2, 3], [2, 4], 'b'),
48+
('abc,cd->abd', [2, 3], [2, 4], 'd'),
49+
('abc,cd->abd', [2, 3], [2, 4], 'bd'),
50+
('abc,cef->abef', [2, 3], [2, 4, 5], None),
51+
('abc,cef->abef', [2, 3], [2, 4, 5], 'bf'),
52+
# Case C2.
53+
('...b,bc->...c', [2, 3], [4], None),
54+
('...b,bc->...c', [2, 3], [4], 'c'),
55+
('...ab,bc->...ac', [2, 3], [2, 4], None),
56+
('...ab,bc->...ac', [2, 4], [2, 4], 'c'),
57+
('...abc,cd->...abd', [2, 3, 4], [2, 3, 5], None),
58+
('...abc,cd->...abd', [2, 3, 4], [2, 3, 5], 'b'),
59+
('...abc,cd->...abd', [2, 3, 4], [2, 3, 5], 'd'),
60+
('...abc,cd->...abd', [2, 3, 4], [2, 3, 5], 'bd'),
61+
('...abc,cef->...abef', [2, 3, 4], [2, 3, 5, 6], None),
62+
('...abc,cef->...abef', [2, 3, 4], [2, 3, 5, 6], 'bf'),
63+
# Case C3.
64+
('ab...,bc->ac...', [2, 3], [4, 3], None),
65+
('ab...,bc->ac...', [2, 3], [4, 3], 'c'),
66+
('abc...,cd->abd...', [2, 3, 4], [2, 5, 4], None),
67+
('abc...,cd->abd...', [2, 3, 4], [2, 5, 4], 'b'),
68+
('abc...,cd->abd...', [2, 3, 4], [2, 5, 4], 'd'),
69+
('abc...,cd->abd...', [2, 3, 4], [2, 5, 4], 'bd'),
70+
('abc...,cef->abef...', [2, 3, 4], [2, 5, 6, 4], None),
71+
('abc...,cef->abef...', [2, 3, 4], [2, 5, 6, 4], 'bf'),
72+
]
73+
74+
75+
def get_einsum_layer_registry():
76+
einsum_registry = layer_registry.LayerRegistry()
77+
einsum_registry.insert(
78+
tfm.nlp.layers.EinsumDense,
79+
einsum_dense.einsum_layer_computation,
80+
)
81+
return einsum_registry
82+
83+
84+
class GradNormTest(tf.test.TestCase, parameterized.TestCase):
85+
86+
def setUp(self):
87+
super().setUp()
88+
self.strategy = tf.distribute.get_strategy()
89+
self.using_tpu = False
90+
91+
@parameterized.product(
92+
layer_name=list(get_einsum_layer_generators()),
93+
param_tuple=get_einsum_parameter_tuples(),
94+
num_microbatches=[None, 2],
95+
is_eager=[True, False],
96+
)
97+
def test_gradient_norms_on_various_models(
98+
self,
99+
layer_name,
100+
param_tuple,
101+
num_microbatches,
102+
is_eager,
103+
):
104+
# Parse inputs to generate test data. Note that each batched input is a
105+
# reshape of a `tf.range()` call.
106+
equation, input_dims, output_dims, bias_axes = param_tuple
107+
batch_size = 4
108+
example_size = tf.reduce_prod(input_dims)
109+
example_values = tf.range(batch_size * example_size, dtype=tf.float32)
110+
x_batch = tf.reshape(example_values, [batch_size] + input_dims)
111+
112+
# Make the layer generator via currying.
113+
einsum_generator = get_einsum_layer_generators()[layer_name]
114+
115+
def curried_generator(a, b):
116+
del a, b
117+
return einsum_generator(equation, output_dims, bias_axes)
118+
119+
# Load shared assets to all devices.
120+
with self.strategy.scope():
121+
model = common_test_utils.get_model_from_generator(
122+
model_generator=common_test_utils.make_one_layer_functional_model,
123+
layer_generator=curried_generator,
124+
input_dims=input_dims,
125+
output_dims=output_dims,
126+
is_eager=is_eager,
127+
)
128+
129+
# Define the main testing ops. These may be later compiled to a Graph op.
130+
def test_op(x):
131+
return common_test_utils.get_computed_and_true_norms_from_model(
132+
model=model,
133+
per_example_loss_fn=None,
134+
num_microbatches=num_microbatches,
135+
x_batch=x,
136+
registry=get_einsum_layer_registry(),
137+
)
138+
139+
# TPUs can only run `tf.function`-decorated functions.
140+
if self.using_tpu:
141+
test_op = tf.function(test_op, autograph=False)
142+
143+
# TPUs use lower precision than CPUs, so we relax our criterion.
144+
# E.g., one of the TPU runs generated the following results:
145+
#
146+
# computed_norm = 93.48296
147+
# true_norm = 93.31176
148+
# abs_diff = 0.17120361
149+
# rel_diff = 0.00183475
150+
#
151+
# which is a reasonable level of error for computing gradient norms.
152+
# Other trials also give an absolute (resp. relative) error of around
153+
# 0.05 (resp. 0.0015).
154+
rtol = 1e-2 if self.using_tpu else 1e-3
155+
atol = 5e-1 if self.using_tpu else 1e-2
156+
157+
# Set up the device ops and run the test.
158+
computed_norms, true_norms = self.strategy.run(test_op, args=(x_batch,))
159+
# TPUs return replica contexts, which must be unwrapped.
160+
if self.using_tpu:
161+
common_test_utils.assert_replica_values_are_close(self, computed_norms)
162+
common_test_utils.assert_replica_values_are_close(self, true_norms)
163+
computed_norms = computed_norms.values[0]
164+
true_norms = true_norms.values[0]
165+
expected_size = num_microbatches or batch_size
166+
self.assertEqual(tf.shape(computed_norms)[0], expected_size)
167+
self.assertAllClose(computed_norms, true_norms, rtol=rtol, atol=atol)
168+
169+
170+
if __name__ == '__main__':
171+
tf.test.main()

‎tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_dense_tpu_test.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Copyright 2023, The TensorFlow Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import tensorflow as tf
16+
from tensorflow_privacy.privacy.fast_gradient_clipping import common_test_utils
17+
from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import einsum_dense_test
18+
19+
20+
class GradNormTpuTest(einsum_dense_test.GradNormTest):
21+
22+
def setUp(self):
23+
super(einsum_dense_test.GradNormTest, self).setUp()
24+
self.strategy = common_test_utils.create_tpu_strategy()
25+
self.assertIn('TPU', self.strategy.extended.worker_devices[0])
26+
self.using_tpu = True
27+
28+
29+
if __name__ == '__main__':
30+
tf.test.main()

0 commit comments

Comments
 (0)