<p dir="auto"><strong>System information</strong>.</p>
<p dir="auto">TensorFlow version (you are using): 2.5.0<br>
Are you willing to contribute it (Yes/No) : No</p>
<p dir="auto"><strong>Describe the feature and the current behavior/state</strong>.</p>
<p dir="auto">When I subclass <code class="notranslate">Layer</code> to implement a custom layer, sometimes I need to storage temporary values (e.g. hidden features)  in forward propogation (<code class="notranslate">call</code> method) for later use, here is an example:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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 &lt; 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)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>
<span class="pl-k">from</span> <span class="pl-s1">tensorflow</span>.<span class="pl-s1">keras</span> <span class="pl-k">import</span> <span class="pl-s1">layers</span>

<span class="pl-k">class</span> <span class="pl-v">MyLayer</span>(<span class="pl-s1">layers</span>.<span class="pl-v">Layer</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>):
        <span class="pl-en">super</span>().<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">convs</span> <span class="pl-c1">=</span> [
            <span class="pl-s1">layers</span>.<span class="pl-v">Conv2D</span>(<span class="pl-c1">64</span> <span class="pl-k">if</span> <span class="pl-s1">i</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">2</span> <span class="pl-k">else</span> <span class="pl-c1">3</span>, (<span class="pl-c1">3</span>, <span class="pl-c1">3</span>), <span class="pl-s1">strides</span><span class="pl-c1">=</span><span class="pl-c1">2</span>, <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'same'</span>)
            <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">3</span>)]

    <span class="pl-k">def</span> <span class="pl-en">call</span>(<span class="pl-s1">self</span>, <span class="pl-s1">inputs</span>):
        <span class="pl-s1">gt</span> <span class="pl-c1">=</span> <span class="pl-s1">inputs</span>[<span class="pl-s">'gt'</span>]
        <span class="pl-s1">pred</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">forward</span>(<span class="pl-s1">inputs</span>)
        <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">compute_loss</span>(<span class="pl-s1">gt</span>, <span class="pl-s1">pred</span>)
        <span class="pl-k">return</span> <span class="pl-s1">loss</span>

    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">inputs</span>):
        <span class="pl-c"># implement forward propogation and reserve intermediate feature map sizes</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">featmap_sizes</span> <span class="pl-c1">=</span> []
        <span class="pl-s1">feat</span> <span class="pl-c1">=</span> <span class="pl-s1">image</span> <span class="pl-c1">=</span> <span class="pl-s1">inputs</span>[<span class="pl-s">'image'</span>]
        <span class="pl-k">for</span> <span class="pl-s1">conv</span> <span class="pl-c1">in</span> <span class="pl-s1">self</span>.<span class="pl-s1">convs</span>:
            <span class="pl-s1">feat</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">relu</span>(<span class="pl-en">conv</span>(<span class="pl-s1">feat</span>))
            <span class="pl-s1">self</span>.<span class="pl-s1">featmap_sizes</span>.<span class="pl-en">append</span>(<span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">feat</span>)[<span class="pl-c1">1</span>:<span class="pl-c1">3</span>])
        <span class="pl-s1">pred</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">image</span>.<span class="pl-en">resize</span>(<span class="pl-s1">feat</span>, <span class="pl-s1">tf</span>.<span class="pl-en">shape</span>(<span class="pl-s1">image</span>)[<span class="pl-c1">1</span>:<span class="pl-c1">3</span>], <span class="pl-s">'bilinear'</span>) <span class="pl-c1">+</span> <span class="pl-s1">image</span>
        <span class="pl-k">return</span> <span class="pl-s1">pred</span>

    <span class="pl-k">def</span> <span class="pl-en">compute_loss</span>(<span class="pl-s1">self</span>, <span class="pl-s1">y_true</span>, <span class="pl-s1">y_pred</span>):
        <span class="pl-c"># maybe a simple mse loss</span>
        <span class="pl-k">return</span> <span class="pl-s1">tf</span>.<span class="pl-en">reduce_mean</span>((<span class="pl-s1">y_true</span> <span class="pl-c1">-</span> <span class="pl-s1">y_pred</span>)<span class="pl-c1">**</span><span class="pl-c1">2</span>)

<span class="pl-s1">my_layer</span> <span class="pl-c1">=</span> <span class="pl-v">MyLayer</span>()
<span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">my_layer</span>({
    <span class="pl-s">'image'</span>: <span class="pl-s1">tf</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>((<span class="pl-c1">1</span>, <span class="pl-c1">64</span>, <span class="pl-c1">64</span>, <span class="pl-c1">3</span>)),
    <span class="pl-s">'gt'</span>: <span class="pl-s1">tf</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>((<span class="pl-c1">1</span>, <span class="pl-c1">64</span>, <span class="pl-c1">64</span>, <span class="pl-c1">3</span>))
})
<span class="pl-en">print</span>(<span class="pl-s1">my_layer</span>.<span class="pl-s1">featmap_sizes</span>)</pre></div>
<p dir="auto">It works fine when not using distribute strategy (or default strategy <code class="notranslate">tf.distribute.get_strategy()</code>), <code class="notranslate">self.featmap_sizes</code> is a list of three tensors (in this case [[32, 32], [16, 16], [8, 8]]). But when running under MirroredStrategy, I found that <code class="notranslate">self.featmap_sizes</code> may contain more than three tensors, and the order of tensors may be wrong:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# 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)"><pre class="notranslate"><span class="pl-c"># assume we have 2 GPUs</span>
<span class="pl-s1">distribute_strategy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">distribute</span>.<span class="pl-v">MirroredStrategy</span>()
<span class="pl-k">with</span> <span class="pl-s1">distribute_strategy</span>.<span class="pl-en">scope</span>():
    <span class="pl-s1">my_layer2</span> <span class="pl-c1">=</span> <span class="pl-v">MyLayer</span>()
 
<span class="pl-s1">data</span> <span class="pl-c1">=</span> <span class="pl-s1">distribute_strategy</span>.<span class="pl-en">experimental_distribute_values_from_function</span>(
    <span class="pl-k">lambda</span> <span class="pl-s1">ctx</span>: {
        <span class="pl-s">'image'</span>: <span class="pl-s1">tf</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>((<span class="pl-c1">1</span>, <span class="pl-c1">64</span>, <span class="pl-c1">64</span>, <span class="pl-c1">3</span>)),
        <span class="pl-s">'gt'</span>: <span class="pl-s1">tf</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>((<span class="pl-c1">1</span>, <span class="pl-c1">64</span>, <span class="pl-c1">64</span>, <span class="pl-c1">3</span>))
    })
<span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s1">distribute_strategy</span>.<span class="pl-en">run</span>(<span class="pl-s1">my_layer2</span>, <span class="pl-s1">args</span><span class="pl-c1">=</span>(<span class="pl-s1">data</span>,))
<span class="pl-en">print</span>(<span class="pl-s1">my_layer2</span>.<span class="pl-s1">featmap_sizes</span>)</pre></div>
<p dir="auto">As expected, <code class="notranslate">loss</code> is a <code class="notranslate">PerReplica</code> tensor but <code class="notranslate">self.featmap_sizes</code> is not, <code class="notranslate">self.featmap_sizes</code> contains normal <code class="notranslate">tf.Tensor</code> 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).</p>
<p dir="auto">Now I have two questions:</p>
<ul dir="auto">
<li>Whether tensor aggregation (not the reduction of <code class="notranslate">tf.Variable</code>) cross replicas is only for output tensors of function run by <code class="notranslate">distribute_strategy.run</code> ? (As in the example above, the returned loss is aggregated as <code class="notranslate">PerReplica</code>, but the property <code class="notranslate">self.featmap_sizes</code> is not). Sorry I didn't find the description of this part in the <a href="https://tensorflow.google.cn/api_docs/python/tf/distribute/Strategy?hl=en#run" rel="nofollow">tensorflow document</a>. 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.</li>
<li>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 <code class="notranslate">Layer</code>, I found that it defines a <code class="notranslate">self._thread_local</code> attribute in the constructor, Is this related to the problem I stated? Do I need to save the temporary tensors I need in <code class="notranslate">self._thread_local</code> like <code class="notranslate">self._thread_local.featmap_sizes = []</code>? Sorry I didn't see anything about <code class="notranslate">self._thread_local</code> in the <a href="https://tensorflow.google.cn/api_docs/python/tf/keras/layers/Layer?hl=en" rel="nofollow">document of keras Layer</a>.</li>
</ul>
<p dir="auto">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~</p>
<p dir="auto"><strong>Will this change the current api? How?</strong><br>
No</p>
<p dir="auto"><strong>Who will benefit from this feature?</strong><br>
Anyone who want to customize complex functions in keras Layer subclassing under MirroredStrategy.</p>
<p dir="auto"><strong><a href="https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md">Contributing</a></strong></p>
<ul dir="auto">
<li>Do you want to contribute a PR? (yes/no): no</li>
<li>If yes, please read <a href="https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md">this page</a> for instructions</li>
<li>Briefly describe your candidate solution(if contributing):</li>
</ul>
<p dir="auto">Hello,</p>
<p dir="auto">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</p>
<p dir="auto">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.</p>
<p dir="auto">The class <a href="https://keras.io/utils/" rel="nofollow"><code class="notranslate">keras.utils.Sequence</code></a> is a very useful work-around, but it's not very well documented and not easy to use (it requires an undocumented <code class="notranslate">OrderedEnqueuer</code> wrapper, etc.)  I am interested in creating a <code class="notranslate">fit_sequence</code> method that takes care of this.</p>
<p dir="auto">Thoughts?</p>
<p dir="auto">Thanks,<br>
Jose</p>