Repository navigation
Seed the augmentation choice in AugMix - #23880
Conversation
AugMix accepts a seed and serializes it in get_config, but the augmentation choice in get_random_transformation drew from Python's global random module. Two layers built with the same seed produced different outputs, and under tf.data the choice was baked into the graph at trace time so every batch went through the same augmentations. Draw layer_idxes of shape (num_chains, chain_depth) with the layer seed, using one index past the end of the candidate list to mean no augmentation, and sample the transformation of every candidate for each chain depth. transform_images and transform_segmentation_masks share _apply_selected_layer, which casts every branch to the input dtype and selects with ops.switch while tracing or with branches[int(layer_idx)] in eager. The random import is no longer needed. Add test_seed_reproducibility, which fails without this change, and test_graph_issue, which compares the stacked per-batch indices so that a frozen choice cannot pass it.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the AugMix preprocessing layer to support symbolic execution and tracing across different backends. It replaces Python-level random choices with backend-agnostic random operations and introduces dynamic branching using ops.switch when tracing. The review feedback suggests a more robust, backend-agnostic approach using a try...except block in _apply_selected_layer instead of the fragile _is_tracing helper, which would also allow removing the newly added jax_utils import.
| def _is_tracing(self, layer_idx): | ||
| if backend_utils.in_tf_graph(): | ||
| return True | ||
| return jax_utils.is_in_jax_tracing_scope(layer_idx) | ||
|
|
||
| def _apply_selected_layer(self, inputs, transforms, layer_idx, method): | ||
| dtype = self.backend.standardize_dtype(inputs.dtype) | ||
|
|
||
| def make_branch(layer_name, transformation): | ||
| layer = getattr(self, layer_name) | ||
| transform = getattr(layer, method) | ||
| return lambda x: self.backend.ops.cast( | ||
| transform(x, transformation), dtype | ||
| ) | ||
|
|
||
| branches = [ | ||
| make_branch(layer_name, transforms[idx]) | ||
| for idx, layer_name in enumerate(self._augment_layers) | ||
| ] | ||
| branches.append(lambda x: self.backend.ops.cast(x, dtype)) | ||
|
|
||
| if self._is_tracing(layer_idx): | ||
| return self.backend.ops.switch(layer_idx, branches, inputs) | ||
| return branches[int(layer_idx)](inputs) |
There was a problem hiding this comment.
Instead of relying on explicit tracing checks (_is_tracing) which can be fragile and fail during Keras functional model construction or under PyTorch compilation (torch.compile), we can use a standard Python try...except block. This is extremely robust, backend-agnostic, and avoids the need for the _is_tracing helper entirely.
def _apply_selected_layer(self, inputs, transforms, layer_idx, method):
dtype = self.backend.standardize_dtype(inputs.dtype)
def make_branch(layer_name, transformation):
layer = getattr(self, layer_name)
transform = getattr(layer, method)
return lambda x: self.backend.ops.cast(
transform(x, transformation), dtype
)
branches = [
make_branch(layer_name, transforms[idx])
for idx, layer_name in enumerate(self._augment_layers)
]
branches.append(lambda x: self.backend.ops.cast(x, dtype))
try:
idx = int(layer_idx)
return branches[idx](inputs)
except (TypeError, ValueError, RuntimeError):
return self.backend.ops.switch(layer_idx, branches, inputs)| ) | ||
| from keras.src.random import SeedGenerator | ||
| from keras.src.utils import backend_utils | ||
| from keras.src.utils import jax_utils |
There was a problem hiding this comment.
Since we can use a robust try...except block to handle symbolic vs. eager execution in _apply_selected_layer, the jax_utils import is no longer needed and can be removed.
|
I measured both halves of this. Functional model construction does not reach Under
The current form is the one @maitry63 specified in #23871, so I would rather a maintainer decide. If |
keerthanakadiri
added
the
stat:awaiting keras-eng
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23880 +/- ##
==========================================
- Coverage 87.41% 84.64% -2.77%
==========================================
Files 479 479
Lines 63950 63829 -121
Branches 10817 10763 -54
==========================================
- Hits 55903 54031 -1872
- Misses 5266 7162 +1896
+ Partials 2781 2636 -145
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
JyotinderSingh
left a comment
JyotinderSingh
left a comment
There was a problem hiding this comment.
The changes are correct. Left a comment about the added cost.
| len(self._augment_layers) + 1, | ||
| seed=seed, | ||
| ) | ||
| chain_transforms = [ |
There was a problem hiding this comment.
This PR reports numbers for TensorFlow only. Two more costs exist.
JAX compile time. Wrap the layer in jax.jit (8x8 images, CPU). Compilation takes about 270 s with this change. It takes about 300 s with all_ops=True. Master takes about 35 s.
Each of the 9 (chain, depth) slots is now a switch over all 8 branches (12 with all_ops). The traced graph therefore holds every candidate transformation 9 times. Master bakes in one choice. A model that contains AugMix pays this cost on the first fit step under JAX.
This follows the design in the issue. I think the cost is acceptable. But putting it in the description would be good.
Eager cost. In eager mode, the layer samples the transformation of every candidate. It then uses only one. Time per call, master -> this PR:
- torch: about 35 -> 140 ms
- TensorFlow: about 770 -> 1700 ms
- JAX: about 200 -> 280 ms
These are just numbers from a small run I did, so the absolute values are not exact. Your description offers an eager shortcut. Please add it. In eager mode, the drawn index is a known value, so only the drawn candidate needs sampling. Draw layer_idxes first. Then do this:
if self._is_tracing(layer_idxes):
chain_transforms = [
[
self._sample_layer_transformations(data)
for _ in range(self.chain_depth)
]
for _ in range(self.num_chains)
]
else:
chain_transforms = [
[
self._sample_layer_transformations(
data, selected=int(layer_idxes[i][j])
)
for j in range(self.chain_depth)
]
for i in range(self.num_chains)
]Change _sample_layer_transformations to (self, data, selected=None). When selected is set, it returns None for each index that is not selected. _apply_selected_layer already calls only the selected branch in eager mode. No other change is necessary.
I tried this change. On torch, the eager call returns to about 36 ms. On torch and JAX, the same seed still gives identical output.
| ] | ||
| branches.append(lambda x: self.backend.ops.cast(x, dtype)) | ||
|
|
||
| if self._is_tracing(layer_idx): |
There was a problem hiding this comment.
Keep this code as it is. Gemini suggests a try/except block. That block also catches real TypeError, RuntimeError and ValueError errors that a branch raises (for example, from transform_images). It then runs the branch again through ops.switch. The real error is hidden behind a tracing error that is difficult to understand. The issue asks for the explicit check with in_tf_graph() and is_in_jax_tracing_scope().
There was a problem hiding this comment.
Done, the eager path now samples only the drawn candidate, using your snippet as is. _apply_selected_layer is unchanged. On TensorFlow a call draws at most 9 sub-layer transformations instead of 63 (99 with all_ops=True), and eager calls run 3.6x to 5.3x faster. Both costs are in the description, with your JAX numbers.
One side effect worth knowing: a transformation drawn in eager can't be replayed under tracing anymore, since the undrawn candidates are None. call is unaffected.
Tracing still samples every candidate, since the index is only known at run time.
JyotinderSingh
left a comment
JyotinderSingh
left a comment
There was a problem hiding this comment.
Thanks for addressing the reviews. LGTM.
JyotinderSingh
added
ready to pull
|
@hertschuh @JyotinderSingh approved this on Oct 17 and it has carried |
Description
AugMixaccepts aseedand serializes it inget_config, but the augmentation choice drew from Python's globalrandom, soAugMix(seed=1234)was not reproducible. Undertf.datathe choice was also baked into the graph at trace time, so every batch went through the same augmentations.Fixes #23871
Following the design in the issue:
get_random_transformationdrawslayer_idxesof shape(num_chains, chain_depth)withself.backend.random.randint(..., len(self._augment_layers) + 1, seed=seed), where the extra index means no augmentation. While tracing it samples the transformation of every candidate for each chain depth; in eager mode the drawn index is already known, so_sample_layer_transformationssamples only the drawn candidate and returnsNonefor the others, andget_random_transformationpicks the path with_is_tracing(layer_idxes).transform_imagesandtransform_segmentation_masksshare_apply_selected_layer, which casts every branch to the input dtype and selects withself.backend.ops.switchwhile tracing, or withbranches[int(layer_idx)]in eager.import randomis no longer needed and is gone.test_graph_issuehere differs from the one inrand_augment_test.pyin two ways. It mapsget_random_transformation(x)["layer_idxes"]instead of the whole dict, because the dict now holds one candidate transformation per layer andauto_contrastandequalizationcontributeNone, whichDataset.mapcannot build aTypeSpecfor. And it assertslen(np.unique(np.stack(key_list), axis=0)) > 1instead of the flatlen(np.unique(key_list)) != 1: the flat form counts distinct values inside a single draw rather than distinct draws across batches, so a frozen choice still passes it, in 90.8% of 20000 simulated runs forRandAugment's(num_ops,)indices and in 100% of them for the(num_chains, chain_depth)indices used here. Stacking first compares the batches themselves.Measured on TensorFlow 2.21. In eager mode a call draws at most 9 sub-layer transformations instead of 63, or 99 with
all_ops=True; a slot that lands on no augmentation draws nothing, as on master. Eager calls are 3.6x, 5.3x and 3.9x faster than sampling every candidate (minimum of 15 calls, implementations alternated three times). Against master the gap is within run to run drift on this machine.Under JAX the cost is compile time. @JyotinderSingh measured about 35 s on master, about 270 s with this change and about 300 s with
all_ops=True: each of the nine(chain, depth)slots is aswitchover every candidate, so the traced graph holds them all nine times. A model containingAugMixpays this once, on the firstfitstep.Side effects of the eager shortcut: eager and traced runs no longer consume the same draws, so the same seed gives different augmentations in the two modes (master already differed, since its Python choice was frozen at trace time), and a transformation drawn in eager cannot be replayed under tracing because the undrawn candidates are
None.calldraws and applies in the same regime and is unaffected. Same-seed reproducibility holds, maximum absolute difference 0.0.Tested locally with the CI configs:
aug_mix_test.pyandbase_image_preprocessing_layer_test.pypass on tensorflow, jax, numpy and torch (channels_first). Undertorch.compilewith the eager backend the Python path runs and the output matches eager at the same seed. openvino and a full inductor compile are untested here.Disclosure: I used an AI coding assistant while investigating this and preparing the patch. I have reviewed and tested every line myself, I understand the mechanics of the fix, and I take full responsibility for the code.
Contributor Agreement
Please review our PR Contribution Policy and AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:
Fixes #xxxabove).Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed. PRs without a linked, assigned issue will be converted to draft.