Closed james77777778 closed 11 months ago
Patch coverage: 98.94%
and project coverage change: +0.06%
:tada:
Comparison is base (
bb21710
) 76.49% compared to head (8f06862
) 76.56%.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
Thanks for the PR, this is great! What performance changes did you observe in e.g. BatchNormalization when using moments
in different backends compared to the existing manual implementation?
Thanks for the PR, this is great! What performance changes did you observe in e.g. BatchNormalization when using
moments
in different backends compared to the existing manual implementation?
EDITED: Please see the newest comment and let me know which implementation should be taken.
ORIGINAL: @fchollet
I have reordered the operations in ops.nn.moments
to get a significant speed-up. The key is to downsize the tensor before performing element-wise operations (such as jnp.subtract
)
# actually, tf.nn.moments could be faster...
# original
variance = jnp.mean(jnp.square(x - jax.lax.stop_gradient(mean)), axis=axes, keepdims=True)
# faster version
variance = jnp.mean(jnp.square(x), axis=axes, keepdims=True) - jnp.square(
jax.lax.stop_gradient(mean)
)
I have also updated the Normalization layers. I observed better performance in TF & JAX but torch showed similar performance.
Benchmark script:
from keras_core import layers
from keras_core import models
from keras_core import ops
x_train = ops.random.uniform(shape=(1024, 224, 224, 3))
y_train = ops.random.uniform(shape=(1024, 224, 224, 3))
# layers.BatchNormalization
# layers.GroupNormalization
# layers.LayerNormalization
normalization_cls = layers.BatchNormalization
normalization_args = {}
if normalization_cls is layers.GroupNormalization:
normalization_args = {"groups": 3}
model = models.Sequential(
[
layers.InputLayer(shape=(224, 224, 3)),
normalization_cls(**normalization_args),
normalization_cls(**normalization_args),
normalization_cls(**normalization_args),
]
)
model.compile(loss="mse", optimizer="adam")
model.fit(x_train, y_train, batch_size=128, epochs=3)
Results (with 1080 8gb card):
backend | layer | current implementation | using ops.nn.moments |
notes |
---|---|---|---|---|
tensorflow | BatchNormalization | 58ms | 58ms | fair |
jax | BatchNormalization | 63ms | 63ms | fair |
torch | BatchNormalization | 73ms | 72ms | fair |
tensorflow | GroupNormalization | 89ms | 61ms | ⬇️ |
jax | GroupNormalization | 96ms | 72ms | ⬇️ |
torch | GroupNormalization | 72ms | 74ms | fair |
tensorflow | LayerNormalization | 52ms | 47ms | ⬇️ |
jax | LayerNormalization | 68ms | 59ms | ⬇️ |
torch | LayerNormalization | 88ms | 90ms | fair |
Hi @fchollet I have a question about the choice of implementation:
tf.nn.moments
uses a slower but more numerically stable version to compute variance
variance = math_ops.reduce_mean(
math_ops.squared_difference(y, array_ops.stop_gradient(mean)),
axes,
keepdims=True,
name="variance")
Keras Core and Flax (the NN library for JAX) use a faster but less stable version for computing variance
# use_fast_variance=True by default in Flax
if use_fast_variance:
mu, mu2 = maybe_distributed_mean(x, _abs_sq(x))
# mean2 - _abs_sq(mean) is not guaranteed to be non-negative due
# to floating point round-off errors.
var = jnp.maximum(0.0, mu2 - _abs_sq(mu))
else:
mu = maybe_distributed_mean(x)
var = maybe_distributed_mean(_abs_sq(x - jnp.expand_dims(mu, axes)))
I just checked the cpu version and torch should use the same approach as tensorflow
backend | layer | current implementation | fast variance computation | tf.nn.moments |
---|---|---|---|---|
tensorflow | BatchNormalization | 58ms | 58ms | 69ms |
tensorflow | GroupNormalization | 89ms | 61ms | 73ms |
tensorflow | LayerNormalization | 52ms | 47ms | 48ms |
Which one should we take?
Thanks for the analysis -- what does "current implementation" refer to? The implementation in this PR?
Thanks for the analysis -- what does "current implementation" refer to? The implementation in this PR?
Sorry for the confusion. I have updated the table as follows:
backend | layer | manual implementation (before this PR) | ops.nn.moments (fast but unstable variance computation, this PR's implementation) |
ops.nn.moments (stable but slower variance computation) |
---|---|---|---|---|
tensorflow | BatchNormalization | 58ms | 58ms | 69ms |
jax | BatchNormalization | 63ms | 63ms | 75ms |
torch | BatchNormalization | 73ms | 72ms | 74ms |
tensorflow | GroupNormalization | 89ms | 61ms | 73ms |
jax | GroupNormalization | 96ms | 72ms | 83ms |
torch | GroupNormalization | 72ms | 74ms | 76ms |
tensorflow | LayerNormalization | 52ms | 47ms | 48ms |
jax | LayerNormalization | 68ms | 59ms | 60ms |
torch | LayerNormalization | 88ms | 90ms | 91ms |
References:
tf.nn.moments
using stable but slower version (link)use_fast_variance
(link)The question remains: Should we adopt the fast but unstable variance computation or the stable but slower version?
My take is that until we see reports of users running into stability issues, then the fast implementation should be fine. The fact that Flax defaults to it is evidence that there's little issue. I did a quick search within the google codebase and found only a couple of usages of use_fast_variance=False
, among thousands of usages of Flax normalization layers (they aren't commented, so unclear why they went with False
). So it seems that in practice the problem doesn't really surface.
The code looks good! Do you want to include the normalization layer changes in this PR, or merge this PR first and then create another one?
My take is that until we see reports of users running into stability issues, then the fast implementation should be fine. The fact that Flax defaults to it is evidence that there's little issue. I did a quick search within the google codebase and found only a couple of usages of
use_fast_variance=False
, among thousands of usages of Flax normalization layers (they aren't commented, so unclear why they went withFalse
). So it seems that in practice the problem doesn't really surface.
Thanks for the valuable insights about the usage of Flax. After a lot of searching, I can only find this comment defending the numerically stable computation of variance (without evidence?) https://github.com/tensorflow/tensorflow/pull/4198#issuecomment-244824858
The code looks good! Do you want to include the normalization layer changes in this PR, or merge this PR first and then create another one?
I think the changes are already in this PR. Please let me know if I missed anything.
Hi @fchollet
This PR should be ready. We now have fast mean & variance computation using ops.nn.moments
, and it is applied to BatchNormalization, GroupNormalization and LayerNormalization to achieve some speed-ups.
Addtionally, this PR addresses the overflow and underflow issue that occur when the input is float16. (I encountered this before when using GroupNormalization with mixed_float16) Credits to @fsx950223 (https://github.com/tensorflow/tensorflow/pull/52217)
EDITED: This PR adds
ops.nn.moments
and improves some normalization layers by fast mean and variance computation. This PR also addresses the overflow & underflow issue when the input tensor is float16.There are two approaches to compute variance:
ops.nn.moments
(fast but unstable, this PR)ops.nn.moments
(stable but slower)References:
tf.nn.moments
using stable but slower version (link)use_fast_variance
(link)