<p dir="auto"><em>Please make sure that this is a bug. As per our <a href="https://github.com/tensorflow/tensorflow/blob/master/ISSUES.md">GitHub Policy</a>, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template</em></p>
<p dir="auto"><strong>System information</strong></p>
<ul dir="auto">
<li>Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes</li>
<li>OS Platform and Distribution (e.g., Linux Ubuntu 16.04): MacOS 10.14.5</li>
<li>Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:</li>
<li>TensorFlow installed from (source or binary): binary</li>
<li>TensorFlow version (use command below): 2.0.0-beta1</li>
<li>Python version: 3.6</li>
<li>Bazel version (if compiling from source):</li>
<li>GCC/Compiler version (if compiling from source):</li>
<li>CUDA/cuDNN version:</li>
<li>GPU model and memory:</li>
</ul>
<p dir="auto"><strong>Describe the current behavior</strong><br>
Using <code class="notranslate">tf.function</code> when enumerating a dataset will cause an infinite loop.</p>
<p dir="auto"><strong>Describe the expected behavior</strong><br>
Using <code class="notranslate">tf.function</code> when enumerating a dataset should not change the looping behavior.</p>
<p dir="auto"><strong>Code to reproduce the issue</strong><br>
The code snippet below will hang after the last function call. I'm not printing anything because calling <code class="notranslate">tf.print</code> results in a syntax error on colab and I know that these snippets are being run on colab by you. When printing the variable <code class="notranslate">i</code>, it's clear that the loop just never stops, i.e. <code class="notranslate">i</code> increase indefinitely.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow as tf

ds = tf.data.Dataset.from_tensor_slices([1,2,3,4])

def test_loop_without_enumerate_without_decorator(ds):
    for val in ds:
        pass

@tf.function
def test_loop_without_enumerate_with_decorator(ds):
    for val in ds:
        pass

def test_loop_with_enumerate_without_decorator(ds):
    for i, val in enumerate(ds):
        pass

@tf.function
def test_loop_with_enumerate_with_decorator(ds):
    for i, val in enumerate(ds):
        pass

print(&quot;Without tf.function and without enumerate&quot;)
test_loop_without_enumerate_without_decorator(ds)

print(&quot;Without tf.function and with enumerate&quot;)
test_loop_with_enumerate_without_decorator(ds)

print(&quot;With tf.function and without enumerate&quot;)
test_loop_without_enumerate_with_decorator(ds)

print(&quot;With tf.function and with enumerate&quot;)
test_loop_with_enumerate_with_decorator(ds)"><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-s1">ds</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">data</span>.<span class="pl-v">Dataset</span>.<span class="pl-en">from_tensor_slices</span>([<span class="pl-c1">1</span>,<span class="pl-c1">2</span>,<span class="pl-c1">3</span>,<span class="pl-c1">4</span>])

<span class="pl-k">def</span> <span class="pl-en">test_loop_without_enumerate_without_decorator</span>(<span class="pl-s1">ds</span>):
    <span class="pl-k">for</span> <span class="pl-s1">val</span> <span class="pl-c1">in</span> <span class="pl-s1">ds</span>:
        <span class="pl-k">pass</span>

<span class="pl-en">@<span class="pl-s1">tf</span>.<span class="pl-s1">function</span></span>
<span class="pl-k">def</span> <span class="pl-en">test_loop_without_enumerate_with_decorator</span>(<span class="pl-s1">ds</span>):
    <span class="pl-k">for</span> <span class="pl-s1">val</span> <span class="pl-c1">in</span> <span class="pl-s1">ds</span>:
        <span class="pl-k">pass</span>

<span class="pl-k">def</span> <span class="pl-en">test_loop_with_enumerate_without_decorator</span>(<span class="pl-s1">ds</span>):
    <span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">val</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">ds</span>):
        <span class="pl-k">pass</span>

<span class="pl-en">@<span class="pl-s1">tf</span>.<span class="pl-s1">function</span></span>
<span class="pl-k">def</span> <span class="pl-en">test_loop_with_enumerate_with_decorator</span>(<span class="pl-s1">ds</span>):
    <span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">val</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">ds</span>):
        <span class="pl-k">pass</span>

<span class="pl-en">print</span>(<span class="pl-s">"Without tf.function and without enumerate"</span>)
<span class="pl-en">test_loop_without_enumerate_without_decorator</span>(<span class="pl-s1">ds</span>)

<span class="pl-en">print</span>(<span class="pl-s">"Without tf.function and with enumerate"</span>)
<span class="pl-en">test_loop_with_enumerate_without_decorator</span>(<span class="pl-s1">ds</span>)

<span class="pl-en">print</span>(<span class="pl-s">"With tf.function and without enumerate"</span>)
<span class="pl-en">test_loop_without_enumerate_with_decorator</span>(<span class="pl-s1">ds</span>)

<span class="pl-en">print</span>(<span class="pl-s">"With tf.function and with enumerate"</span>)
<span class="pl-en">test_loop_with_enumerate_with_decorator</span>(<span class="pl-s1">ds</span>)</pre></div>
<p dir="auto"><strong>Other info / logs</strong><br>
Output of the above snippet:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Without tf.function and without enumerate
Without tf.function and with enumerate
With tf.function and without enumerate
With tf.function and with enumerate
**HANGS HERE**"><pre class="notranslate"><code class="notranslate">Without tf.function and without enumerate
Without tf.function and with enumerate
With tf.function and without enumerate
With tf.function and with enumerate
**HANGS HERE**
</code></pre></div>
<h3 dir="auto">Description</h3>
<p dir="auto">I'm trying to understand how TF generates batches from <code class="notranslate">TFRecord</code> file format, how to implement the basic idea of evaluate the whole validation dataset after a full epoch so I've done a small experiment.</p>
<p dir="auto">Basically I've added an <code class="notranslate">image/id</code> key in the <code class="notranslate">TFRecord</code> file that looks like this:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="'image/id': tf.FixedLenFeature(shape=[], dtype=tf.int64)
'id': slim.tfexample_decoder.Tensor('image/id')"><pre class="notranslate"><code class="notranslate">'image/id': tf.FixedLenFeature(shape=[], dtype=tf.int64)
'id': slim.tfexample_decoder.Tensor('image/id')
</code></pre></div>
<p dir="auto">The values of these ids are just numbers increasing from 0 to <code class="notranslate">num_samples - 1</code>. Then I'm reading the data like this:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="...
data_provider = slim.dataset_data_provider.DatasetDataProvider(dataset=dataset, shuffle=False)
raw_image, instance_id, label = data_provider.get(items=['image', 'id', 'label'])
...
num_threads = 1
images, instance_ids, labels = tf.train.batch(tensors=[image, instance_id, label],
                                                      batch_size=batch_size,
                                                      num_threads=num_threads,
                                                      capacity = batch_size,
                                                      allow_smaller_final_batch=False)        
..."><pre class="notranslate"><code class="notranslate">...
data_provider = slim.dataset_data_provider.DatasetDataProvider(dataset=dataset, shuffle=False)
raw_image, instance_id, label = data_provider.get(items=['image', 'id', 'label'])
...
num_threads = 1
images, instance_ids, labels = tf.train.batch(tensors=[image, instance_id, label],
                                                      batch_size=batch_size,
                                                      num_threads=num_threads,
                                                      capacity = batch_size,
                                                      allow_smaller_final_batch=False)        
...
</code></pre></div>
<p dir="auto">I've used <code class="notranslate">shuffle=False</code>, <code class="notranslate">num_threads=1</code> and <code class="notranslate">capacity=batch_size</code> to ensure we are reading the instances in the validation set in order.</p>
<p dir="auto">I've then defined a <code class="notranslate">train_step_fn</code> that evaluates a mini-batch from the validation set after each step. I'm doing this to check and test we are not returning instances with duplicate ids in the same batch. The code looks like this:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="    def train_step_fn(session, *args, **kwargs):
        total_loss, should_stop = train_step(session, *args, **kwargs)
        curr_global_step = tf.train.global_step(session, global_step)
        curr_epoch = curr_global_step / validation_every_n_steps
        image, val_ids, val_loss, accuracy = session.run([validation_images, validation_ids, total_validation_loss, validation_accuracy])
        print(val_ids)
        def float_formatter(x): return &quot;%.4f&quot; % x
        tf.logging.info('after global step {} (epoch {}): validation loss = {}, validation accuracy = {}'
                            .format(curr_global_step, curr_epoch, float_formatter(val_loss), float_formatter(accuracy)))
        return [total_loss, should_stop]"><pre class="notranslate"><code class="notranslate">    def train_step_fn(session, *args, **kwargs):
        total_loss, should_stop = train_step(session, *args, **kwargs)
        curr_global_step = tf.train.global_step(session, global_step)
        curr_epoch = curr_global_step / validation_every_n_steps
        image, val_ids, val_loss, accuracy = session.run([validation_images, validation_ids, total_validation_loss, validation_accuracy])
        print(val_ids)
        def float_formatter(x): return "%.4f" % x
        tf.logging.info('after global step {} (epoch {}): validation loss = {}, validation accuracy = {}'
                            .format(curr_global_step, curr_epoch, float_formatter(val_loss), float_formatter(accuracy)))
        return [total_loss, should_stop]
</code></pre></div>
<p dir="auto">Most of the output I'm seeing makes sense but the output from the first batch looks weird.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="INFO:tensorflow:Starting Session.
INFO:tensorflow:Starting Queues.
INFO:tensorflow:global_step/sec: 0
INFO:tensorflow:global step 1: loss = 4.1099 (3.70 sec/step)
[17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32]
INFO:tensorflow:after global step 1 (epoch 0): validation loss = 4.6337, validation accuracy = 0.0000
INFO:tensorflow:global step 2: loss = 4.2118 (0.52 sec/step)
[33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48]
INFO:tensorflow:after global step 2 (epoch 0): validation loss = 3.9538, validation accuracy = 0.0625
INFO:tensorflow:global step 3: loss = 4.0593 (0.75 sec/step)
[49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64]
INFO:tensorflow:after global step 3 (epoch 0): validation loss = 4.0773, validation accuracy = 0.0000
INFO:tensorflow:global step 4: loss = 3.7749 (0.64 sec/step)
[65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80]
INFO:tensorflow:after global step 4 (epoch 0): validation loss = 4.0847, validation accuracy = 0.0625
INFO:tensorflow:global step 5: loss = 3.7973 (0.68 sec/step)
[81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96]
INFO:tensorflow:after global step 5 (epoch 0): validation loss = 3.5591, validation accuracy = 0.1875
INFO:tensorflow:global step 6: loss = 3.5584 (0.72 sec/step)
[ 97  98  99 100 101 102 103 104 105 106 107 108 109 110 111 112]
INFO:tensorflow:after global step 6 (epoch 0): validation loss = 3.6849, validation accuracy = 0.1250
INFO:tensorflow:global step 7: loss = 3.5303 (0.70 sec/step)
[113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128]
INFO:tensorflow:after global step 7 (epoch 0): validation loss = 3.7812, validation accuracy = 0.1875
INFO:tensorflow:global step 8: loss = 3.2698 (0.85 sec/step)
[129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144]
INFO:tensorflow:after global step 8 (epoch 0): validation loss = 3.3004, validation accuracy = 0.1875
INFO:tensorflow:global step 9: loss = 3.4842 (0.94 sec/step)
INFO:tensorflow:global_step/sec: 0.915008
[145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160]
INFO:tensorflow:after global step 9 (epoch 0): validation loss = 4.0377, validation accuracy = 0.1875
INFO:tensorflow:global step 10: loss = 3.5635 (1.48 sec/step)
[178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193]
INFO:tensorflow:after global step 10 (epoch 0): validation loss = 3.7854, validation accuracy = 0.1875
INFO:tensorflow:global step 11: loss = 3.1773 (0.65 sec/step)
[194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209]"><pre class="notranslate"><code class="notranslate">INFO:tensorflow:Starting Session.
INFO:tensorflow:Starting Queues.
INFO:tensorflow:global_step/sec: 0
INFO:tensorflow:global step 1: loss = 4.1099 (3.70 sec/step)
[17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32]
INFO:tensorflow:after global step 1 (epoch 0): validation loss = 4.6337, validation accuracy = 0.0000
INFO:tensorflow:global step 2: loss = 4.2118 (0.52 sec/step)
[33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48]
INFO:tensorflow:after global step 2 (epoch 0): validation loss = 3.9538, validation accuracy = 0.0625
INFO:tensorflow:global step 3: loss = 4.0593 (0.75 sec/step)
[49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64]
INFO:tensorflow:after global step 3 (epoch 0): validation loss = 4.0773, validation accuracy = 0.0000
INFO:tensorflow:global step 4: loss = 3.7749 (0.64 sec/step)
[65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80]
INFO:tensorflow:after global step 4 (epoch 0): validation loss = 4.0847, validation accuracy = 0.0625
INFO:tensorflow:global step 5: loss = 3.7973 (0.68 sec/step)
[81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96]
INFO:tensorflow:after global step 5 (epoch 0): validation loss = 3.5591, validation accuracy = 0.1875
INFO:tensorflow:global step 6: loss = 3.5584 (0.72 sec/step)
[ 97  98  99 100 101 102 103 104 105 106 107 108 109 110 111 112]
INFO:tensorflow:after global step 6 (epoch 0): validation loss = 3.6849, validation accuracy = 0.1250
INFO:tensorflow:global step 7: loss = 3.5303 (0.70 sec/step)
[113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128]
INFO:tensorflow:after global step 7 (epoch 0): validation loss = 3.7812, validation accuracy = 0.1875
INFO:tensorflow:global step 8: loss = 3.2698 (0.85 sec/step)
[129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144]
INFO:tensorflow:after global step 8 (epoch 0): validation loss = 3.3004, validation accuracy = 0.1875
INFO:tensorflow:global step 9: loss = 3.4842 (0.94 sec/step)
INFO:tensorflow:global_step/sec: 0.915008
[145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160]
INFO:tensorflow:after global step 9 (epoch 0): validation loss = 4.0377, validation accuracy = 0.1875
INFO:tensorflow:global step 10: loss = 3.5635 (1.48 sec/step)
[178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193]
INFO:tensorflow:after global step 10 (epoch 0): validation loss = 3.7854, validation accuracy = 0.1875
INFO:tensorflow:global step 11: loss = 3.1773 (0.65 sec/step)
[194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209]
</code></pre></div>
<p dir="auto">I'm assuming the first mini-batch returned should be <code class="notranslate">[0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15]</code> instead of <code class="notranslate">[17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32]</code>. Plus the fact that we are starting from 17 instead of 16 also bothers me a bit. It seems for some unknown reason TF is dropping a batch.</p>
<p dir="auto">I understand this probably won't affect the model training / validation process at all but I'm still pointing this out in case this is hiding a more serious root cause.</p>