Issue 1:
System information.

TensorFlow version (you are using): 2.5.0

Are you willing to contribute it (Yes/No) : No

Describe the feature and the current behavior/state.

When I subclass Layer to implement a custom layer, sometimes I need to storage temporary values (e.g. hidden features) in forward propogation (call method) for later use, here is an example:

import tensorflow as tf
from tensorflow.keras import layers

class MyLayer(layers.Layer):
def __init__(self):
super().__init__()
self.convs = [
layers.Conv2D(64 if i < 2 else 3, (3, 3), strides=2, padding='same')
for i in range(3)]

def call(self, inputs):
gt = inputs['gt']
pred = self.forward(inputs)
loss = self.compute_loss(gt, pred)
return loss

def forward(self, inputs):
# implement forward propogation and reserve intermediate feature map sizes
self.featmap_sizes = []
feat = image = inputs['image']
for conv in self.convs:
feat = tf.nn.relu(conv(feat))
self.featmap_sizes.append(tf.shape(feat)[1:3])
pred = tf.image.resize(feat, tf.shape(image)[1:3], 'bilinear') + image
return pred

def compute_loss(self, y_true, y_pred):
# maybe a simple mse loss
return tf.reduce_mean((y_true - y_pred)**2)

my_layer = MyLayer()
loss = my_layer({
'image': tf.random.normal((1, 64, 64, 3)),
'gt': tf.random.normal((1, 64, 64, 3))
})
print(my_layer.featmap_sizes)

It works fine when not using distribute strategy (or default strategy tf.distribute.get_strategy()), self.featmap_sizes is a list of three tensors (in this case [[32, 32], [16, 16], [8, 8]]). But when running under MirroredStrategy, I found that self.featmap_sizes may contain more than three tensors, and the order of tensors may be wrong:

# assume we have 2 GPUs
distribute_strategy = tf.distribute.MirroredStrategy()
with distribute_strategy.scope():
my_layer2 = MyLayer()

data = distribute_strategy.experimental_distribute_values_from_function(
lambda ctx: {
'image': tf.random.normal((1, 64, 64, 3)),
'gt': tf.random.normal((1, 64, 64, 3))
})
loss = distribute_strategy.run(my_layer2, args=(data,))
print(my_layer2.featmap_sizes)

As expected, loss is a PerReplica tensor but self.featmap_sizes is not, self.featmap_sizes contains normal tf.Tensor object. What's even weirder is that it sometimes contains more than 3 values (This is just a simple example, in such a simple case the exception won't always appear).

Now I have two questions:

Whether tensor aggregation (not the reduction of tf.Variable) cross replicas is only for output tensors of function run by distribute_strategy.run ? (As in the example above, the returned loss is aggregated as PerReplica, but the property self.featmap_sizes is not). Sorry I didn't find the description of this part in the tensorflow document. Since I'm using keras to customize the nearul network component, I'm wondering what happens when I assign or append values to a layer's properties in a distributed situation.

Why does the problem of abnormal property value occur? I'm guessing that since MirroredStrategy is single-process, the layer's property assignment or append operation may be duplicated or overwritten between different devices (maybe they are running in different threads? I have no idea). When I reviewed the source code of keras Layer, I found that it defines a self._thread_local attribute in the constructor, Is this related to the problem I stated? Do I need to save the temporary tensors I need in self._thread_local like self._thread_local.featmap_sizes = []? Sorry I didn't see anything about self._thread_local in the document of keras Layer.

My hope is that each device (replica) should only operate on its own values in a distributed case, and should not flow between replicas without artificial aggregation. If this mechanism has been implicitly implemented, I hope that the official can supplement this description in the document or tutorials. Thank you~

Will this change the current api? How?

No

Who will benefit from this feature?

Anyone who want to customize complex functions in keras Layer subclassing under MirroredStrategy.

Contributing

Do you want to contribute a PR? (yes/no): no

If yes, please read this page for instructions

Briefly describe your candidate solution(if contributing):

Issue 2:
Hello,

Never contributed to Keras before, but I am interested in creating a PR with a new feature, but before I do so, I want to get feedback and check if it would be accepted :P

The function ``fit_generator'' is extremely useful, but it is difficult to use parallelized. It expects a generator, which are not thread-safe in Python. Alternatively, it can run each generator on a different process, but then it is likely that you would end with duplicated observations across processes since generators can't talk to each other.

The class keras.utils.Sequence is a very useful work-around, but it's not very well documented and not easy to use (it requires an undocumented OrderedEnqueuer wrapper, etc.) I am interested in creating a fit_sequence method that takes care of this.

Thoughts?

Thanks,

Jose