How to use keras layers to augment both image and label data in image segmentation tasks？
Hi @wangdada-love ,

For utilizing distribution training on the augmented images we need to create a `strategy` like   `strategy = tf.distribute.MirroredStrategy()` and under that `strategy.scope` context we can create the required model with all augmented layers along with other layers inside it. You can find different augmentation layers available in tensorflow [here](https://www.tensorflow.org/api_docs/python/tf/keras/layers).

Once you created a model under the strategy you can train the model simply with model.fit and keras will take care of the distribution training.Here only we need to ensure batch_size should be multiple of GPUs like an example below.

```
BATCH_SIZE_PER_REPLICA = 64
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
```

You can find a demo tutorial for distribution training [here](https://www.tensorflow.org/tutorials/distribute/keras). I hope this helps for you.

Thanks! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61197">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61197">No</a>

### Issue type

Others

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

tf 2.6

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

python 3.7

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

How to use tool a to enhance both image and label data in image segmentation tasks.
I saw an example in the official document of using Keras layer for data augmentation, which is very useful in the training process of classification models because the data augmentation does not require synchronized operations on labels. However, in the segmentation task, if I embed the data enhancement layer into the model structure, I cannot do operations like rotate 、zoom .etc on the label and image at the same time, because the fit method only feed the original image into the model for inference, which results in the label being isolated from the inference process, so it is impossible to do synchronize Affine transformations with the original image.

Thanks.


### Standalone code to reproduce the issue

```shell
I cannot copy my code from the company computer. I just want to know how to do efficient data augmentation that can apply GPU acceleration and distribute strategy.Thanks.
```


### Relevant log output

_No response_