{
  "id": 135443,
  "title": "TPU: extreme optimizations",
  "url": "/competitions/flower-classification-with-tpus/discussion/135443",
  "author_name": "Martin Görner",
  "post_date": "2020-03-13T21:56:30.801000",
  "votes": 64,
  "comment_count": 39,
  "views": 0,
  "content": "<p><strong>Edit: the learning rate issue highlighed by <a href=\"/yihdarshieh\">@yihdarshieh</a> in the comments has been fixed</strong></p>\n\n<p>All three models for this comparison are in this notebook: <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu\">Custom Training Loop with 100+ flowers on TPU</a></p>\n\n<h3>1) Keras model.fit</h3>\n\n<p>Fine-tuning a Keras flowers model on TPU is quite fast but if you open the TPU profiler (click on the gauges in the top-right corner of your running notebook), you see that it could be a bit faster:</p>\n\n<p>\n\n</p>\n\n<p>25% idle time means that the TPU is waiting. It could be waiting for data. Or there could be some overhead in the way Keras is running its training loop.</p>\n\n<h3>2) Custom training loop</h3>\n\n<p>Let's try to write a custom training loop. Even in a distributed environment like a TPU, this can be written quite easily with two additional functions offered by the TPU distributions strategy: distribute_dataset and run.</p>\n\n<p>```</p>\n\n<h1>training step function</h1>\n\n<p>@tf.function\ndef train_step(images, labels):\n    with tf.GradientTape() as tape:\n        probabilities = model(images, training=True)\n        loss = loss_fn(labels, probabilities)\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))</p>\n\n<h1>distributed dataset</h1>\n\n<p>train_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())</p>\n\n<h1>custom training loop</h1>\n\n<p>for (images, labels) in train_dist_ds:\n    strategy.experimental_run_v2(train_step, args=(images, labels))\n```\nFor full code, including validation, refer to <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Custom-training-loop\">Custom Training Loop with 100+ flowers on TPU, section custom training loop</a></p>\n\n<p>Let's look at TPU idle time in the custom training loop:</p>\n\n<p>\n\n</p>\n\n<p>Not really an improvement...</p>\n\n<h3>3) TPU-optimized custom training loop</h3>\n\n<p>There is still room for optimization, the TPU is a remote accelerator and calling it has a cost. What if we could call it less often. If one training step does more work directly on the TPU, things could go a bit faster. This code uses a loop in the step function. It uses <code>tf.range()</code> instead of <code>range()</code> so that it can be compiled to Tensorflow operation and sent to the TPU. The compilation is forced by <code>@tf.function</code>.</p>\n\n<p>```</p>\n\n<h1>training step function</h1>\n\n<p>@tf.function\ndef train_step(data_iter):\n    def train_step_fn(images, labels):\n        with tf.GradientTape() as tape:\n            probabilities = model(images, training=True)\n            loss = loss_fn(labels, probabilities)\n        grads = tape.gradient(loss, model.trainable_variables)\n        optimizer.apply_gradients(zip(grads, model.trainable_variables))</p>\n\n<pre><code># loop executing on TPU\nfor _ in tf.range(STEPS_PER_TPU_CALL):\n    strategy.experimental_run_v2(train_step_fn, next(data_iter))\n</code></pre>\n\n<h1>distributed dataset</h1>\n\n<p>train_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())</p>\n\n<h1>dataset iterator</h1>\n\n<p>train_data_iter = iter(train_dist_ds)</p>\n\n<h1>custom training loop</h1>\n\n<p>while True:\n    train_step(train_data_iter)\n```</p>\n\n<p>For full code, including validation, refer to <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Optimized-custom-training-loop\">Custom Training Loop with 100+ flowers on TPU, section optimized custom training loop</a></p>\n\n<p>Let's look at the idle time with this TPU-optimized custom training loop:</p>\n\n<p>\n\n</p>\n\n<p>No more idle time ! This confirms that the data pipeline was not the problem. </p>\n\n<h3>Benchmark</h3>\n\n<p>From the notebook <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu\">Custom Training Loop with 100+ flowers on TPU</a>:</p>\n\n<p>| Training loop type | TPU idle time | Total training time | Typical epoch time |\n|--------------------|--------------|-------------------|-------------------|\n| Keras model.fit() | 25% | 417 s | 30 s |\n| simple custom training loop | 23% | 436 s | 32 s |\n| TPU-optimized custom training loop | 00% | 350 s | 22 s |</p>",
  "messages": [
    {
      "id": 771219,
      "postDate": "2020-03-13T21:56:30.800Z",
      "content": "<p><strong>Edit: the learning rate issue highlighed by <a href=\"/yihdarshieh\">@yihdarshieh</a> in the comments has been fixed</strong></p>\n\n<p>All three models for this comparison are in this notebook: <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu\">Custom Training Loop with 100+ flowers on TPU</a></p>\n\n<h3>1) Keras model.fit</h3>\n\n<p>Fine-tuning a Keras flowers model on TPU is quite fast but if you open the TPU profiler (click on the gauges in the top-right corner of your running notebook), you see that it could be a bit faster:</p>\n\n<p>\n\n</p>\n\n<p>25% idle time means that the TPU is waiting. It could be waiting for data. Or there could be some overhead in the way Keras is running its training loop.</p>\n\n<h3>2) Custom training loop</h3>\n\n<p>Let's try to write a custom training loop. Even in a distributed environment like a TPU, this can be written quite easily with two additional functions offered by the TPU distributions strategy: distribute_dataset and run.</p>\n\n<p>```</p>\n\n<h1>training step function</h1>\n\n<p>@tf.function\ndef train_step(images, labels):\n    with tf.GradientTape() as tape:\n        probabilities = model(images, training=True)\n        loss = loss_fn(labels, probabilities)\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))</p>\n\n<h1>distributed dataset</h1>\n\n<p>train_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())</p>\n\n<h1>custom training loop</h1>\n\n<p>for (images, labels) in train_dist_ds:\n    strategy.experimental_run_v2(train_step, args=(images, labels))\n```\nFor full code, including validation, refer to <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Custom-training-loop\">Custom Training Loop with 100+ flowers on TPU, section custom training loop</a></p>\n\n<p>Let's look at TPU idle time in the custom training loop:</p>\n\n<p>\n\n</p>\n\n<p>Not really an improvement...</p>\n\n<h3>3) TPU-optimized custom training loop</h3>\n\n<p>There is still room for optimization, the TPU is a remote accelerator and calling it has a cost. What if we could call it less often. If one training step does more work directly on the TPU, things could go a bit faster. This code uses a loop in the step function. It uses <code>tf.range()</code> instead of <code>range()</code> so that it can be compiled to Tensorflow operation and sent to the TPU. The compilation is forced by <code>@tf.function</code>.</p>\n\n<p>```</p>\n\n<h1>training step function</h1>\n\n<p>@tf.function\ndef train_step(data_iter):\n    def train_step_fn(images, labels):\n        with tf.GradientTape() as tape:\n            probabilities = model(images, training=True)\n            loss = loss_fn(labels, probabilities)\n        grads = tape.gradient(loss, model.trainable_variables)\n        optimizer.apply_gradients(zip(grads, model.trainable_variables))</p>\n\n<pre><code># loop executing on TPU\nfor _ in tf.range(STEPS_PER_TPU_CALL):\n    strategy.experimental_run_v2(train_step_fn, next(data_iter))\n</code></pre>\n\n<h1>distributed dataset</h1>\n\n<p>train_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())</p>\n\n<h1>dataset iterator</h1>\n\n<p>train_data_iter = iter(train_dist_ds)</p>\n\n<h1>custom training loop</h1>\n\n<p>while True:\n    train_step(train_data_iter)\n```</p>\n\n<p>For full code, including validation, refer to <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Optimized-custom-training-loop\">Custom Training Loop with 100+ flowers on TPU, section optimized custom training loop</a></p>\n\n<p>Let's look at the idle time with this TPU-optimized custom training loop:</p>\n\n<p>\n\n</p>\n\n<p>No more idle time ! This confirms that the data pipeline was not the problem. </p>\n\n<h3>Benchmark</h3>\n\n<p>From the notebook <a href=\"https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu\">Custom Training Loop with 100+ flowers on TPU</a>:</p>\n\n<p>| Training loop type | TPU idle time | Total training time | Typical epoch time |\n|--------------------|--------------|-------------------|-------------------|\n| Keras model.fit() | 25% | 417 s | 30 s |\n| simple custom training loop | 23% | 436 s | 32 s |\n| TPU-optimized custom training loop | 00% | 350 s | 22 s |</p>",
      "rawMarkdown": "**Edit: the learning rate issue highlighed by @yihdarshieh in the comments has been fixed**\n\nAll three models for this comparison are in this notebook: [Custom Training Loop with 100+ flowers on TPU](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu)\n\n### 1) Keras model.fit\n\nFine-tuning a Keras flowers model on TPU is quite fast but if you open the TPU profiler (click on the gauges in the top-right corner of your running notebook), you see that it could be a bit faster:\n<p>\n<img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F4003597%2F4959d4a9c096730989ec288f450cac26%2Ftrain%20fit%20TPU%20idle%20time.png?generation=1584134538223267&amp;alt=media\" width=\"300ps\">\n</p>\n\n25% idle time means that the TPU is waiting. It could be waiting for data. Or there could be some overhead in the way Keras is running its training loop.\n\n### 2) Custom training loop\n\nLet's try to write a custom training loop. Even in a distributed environment like a TPU, this can be written quite easily with two additional functions offered by the TPU distributions strategy: distribute_dataset and run.\n\n```\n# training step function\n@tf.function\ndef train_step(images, labels):\n    with tf.GradientTape() as tape:\n        probabilities = model(images, training=True)\n        loss = loss_fn(labels, probabilities)\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n\n# distributed dataset\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\n\n# custom training loop\nfor (images, labels) in train_dist_ds:\n    strategy.experimental_run_v2(train_step, args=(images, labels))\n```\nFor full code, including validation, refer to [Custom Training Loop with 100+ flowers on TPU, section custom training loop](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Custom-training-loop)\n\nLet's look at TPU idle time in the custom training loop:\n<p>\n<img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F4003597%2Ff3f20c836c4d3cecf4f0d2f1a0f5d0c2%2Fsimple%20CTL%20TPU%20idle%20time.png?generation=1584135998683394&amp;alt=media\" width=\"300ps\">\n</p>\n\nNot really an improvement...\n\n### 3) TPU-optimized custom training loop\n\nThere is still room for optimization, the TPU is a remote accelerator and calling it has a cost. What if we could call it less often. If one training step does more work directly on the TPU, things could go a bit faster. This code uses a loop in the step function. It uses `tf.range()` instead of `range()` so that it can be compiled to Tensorflow operation and sent to the TPU. The compilation is forced by `@tf.function`.\n\n```\n# training step function\n@tf.function\ndef train_step(data_iter):\n    def train_step_fn(images, labels):\n        with tf.GradientTape() as tape:\n            probabilities = model(images, training=True)\n            loss = loss_fn(labels, probabilities)\n        grads = tape.gradient(loss, model.trainable_variables)\n        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n\n    # loop executing on TPU\n    for _ in tf.range(STEPS_PER_TPU_CALL):\n        strategy.experimental_run_v2(train_step_fn, next(data_iter))\n\n# distributed dataset\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\n# dataset iterator\ntrain_data_iter = iter(train_dist_ds)\n\n# custom training loop\nwhile True:\n    train_step(train_data_iter)\n```\n\nFor full code, including validation, refer to [Custom Training Loop with 100+ flowers on TPU, section optimized custom training loop](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Optimized-custom-training-loop)\n\nLet's look at the idle time with this TPU-optimized custom training loop:\n<p>\n<img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F4003597%2F817f6ab062a59abe1556f030b0a75491%2Foptimized%20CTL%20TPU%20idle%20time.png?generation=1584136547845529&amp;alt=media\" width=\"300ps\">\n</p>\n\nNo more idle time ! This confirms that the data pipeline was not the problem. \n\n### Benchmark\nFrom the notebook [Custom Training Loop with 100+ flowers on TPU](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu):\n\n| Training loop type | TPU idle time | Total training time | Typical epoch time |\n|--------------------|--------------|-------------------|-------------------|\n| Keras model.fit() | 25% | 417 s | 30 s |\n| simple custom training loop | 23% | 436 s | 32 s |\n| TPU-optimized custom training loop | 00% | 350 s | 22 s |",
      "votes": 63
    },
    {
      "id": 775647,
      "postDate": "2020-03-16T22:55:01.327Z",
      "content": "<p>This optimization shaves 25% from the TPU training time on a single TPU (8 cores). Given the complexity of the code compared to a simple model.fit(), it might not always be worth it. On TPU pods however (32 cores and more), the training time can go down by up to 50% (twice faster).</p>",
      "rawMarkdown": "This optimization shaves 25% from the TPU training time on a single TPU (8 cores). Given the complexity of the code compared to a simple model.fit(), it might not always be worth it. On TPU pods however (32 cores and more), the training time can go down by up to 50% (twice faster).",
      "votes": 8
    },
    {
      "id": 779064,
      "postDate": "2020-03-19T02:08:06.300Z",
      "content": "<p>Thanks for the great sharing. Tensorflow GradientTape() method is really useful!</p>\n\n<p>Did you calculate the idle time when model is finally trained? Or as shown in figure won't there be a small difference in Idle time as in your screenshot between method1(2nd epoch) and method2(3rd epoch)?</p>",
      "rawMarkdown": "Thanks for the great sharing. Tensorflow GradientTape() method is really useful!\n\nDid you calculate the idle time when model is finally trained? Or as shown in figure won't there be a small difference in Idle time as in your screenshot between method1(2nd epoch) and method2(3rd epoch)?",
      "votes": 1
    },
    {
      "id": 771256,
      "postDate": "2020-03-14T00:12:29.040Z",
      "content": "<p>This is awesome <a href=\"/mgornergoogle\">@mgornergoogle</a> , I'll try for sure!</p>",
      "rawMarkdown": "This is awesome @mgornergoogle , I'll try for sure!",
      "votes": 1
    },
    {
      "id": 775339,
      "postDate": "2020-03-16T15:30:00.450Z",
      "content": "<p>I have created a notebook using theses optimizations and applying to K-fold, in case anyone wanna checkout <a href=\"https://www.kaggle.com/dimitreoliveira/flower-with-tpus-k-fold-optimized-training-loops\">here it is</a>, during my experimentations I have faced some issues related to GCS availability, anyone got similar issues?</p>",
      "rawMarkdown": "I have created a notebook using theses optimizations and applying to K-fold, in case anyone wanna checkout [here it is](https://www.kaggle.com/dimitreoliveira/flower-with-tpus-k-fold-optimized-training-loops), during my experimentations I have faced some issues related to GCS availability, anyone got similar issues?",
      "votes": 2,
      "replies": [
        {
          "id": 775670,
          "postDate": "2020-03-16T23:45:54.863Z",
          "content": "<p>Nice notebook. Can you fix the LRSchedule issue in it so that my errors from the first version do not continue to live there ? There are a couple of comments from my original notebook that can now be deleted too. Seach for \"scaled up by 8\" and \"multiply the learning rate schedule by 8\"</p>",
          "rawMarkdown": "Nice notebook. Can you fix the LRSchedule issue in it so that my errors from the first version do not continue to live there ? There are a couple of comments from my original notebook that can now be deleted too. Seach for \"scaled up by 8\" and \"multiply the learning rate schedule by 8\"",
          "votes": 1
        },
        {
          "id": 775835,
          "postDate": "2020-03-17T01:33:31.747Z",
          "content": "<p>Just did it, thanks for the head-up <a href=\"/mgornergoogle\">@mgornergoogle</a> , still, I have no idea why sometimes I get \"UnavailableError: Socket closed\" </p>",
          "rawMarkdown": "Just did it, thanks for the head-up @mgornergoogle , still, I have no idea why sometimes I get \"UnavailableError: Socket closed\" ",
          "votes": 2
        },
        {
          "id": 776735,
          "postDate": "2020-03-17T15:43:52.493Z",
          "content": "<p>That might happen when the XLA compiler crashes, for example when it runs out of memory on a large model. Lowering the batch size in that case sometimes helps.</p>",
          "rawMarkdown": "That might happen when the XLA compiler crashes, for example when it runs out of memory on a large model. Lowering the batch size in that case sometimes helps.",
          "votes": 2
        },
        {
          "id": 777543,
          "postDate": "2020-03-17T18:57:45.600Z",
          "content": "<p>This seems very strange because it can crash with any model architecture, but usually when I use <code>mode.predict()</code>, maybe the right way is to execute predictions outside of the training loop to save memory, I will do some more experimentations and report back.</p>",
          "rawMarkdown": "This seems very strange because it can crash with any model architecture, but usually when I use `mode.predict()`, maybe the right way is to execute predictions outside of the training loop to save memory, I will do some more experimentations and report back.",
          "votes": 1
        },
        {
          "id": 777577,
          "postDate": "2020-03-17T19:33:45.950Z",
          "content": "<p>I recently filed a bug where I had a repeatable crash caused by the drop_remainder=True setting. Maybe you are hitting this bug too ? The setting is not required on repeated datasets but is supposed to make things a little faster. Try without it.</p>",
          "rawMarkdown": "I recently filed a bug where I had a repeatable crash caused by the drop_remainder=True setting. Maybe you are hitting this bug too ? The setting is not required on repeated datasets but is supposed to make things a little faster. Try without it.",
          "votes": 2
        },
        {
          "id": 777871,
          "postDate": "2020-03-18T02:00:36.657Z",
          "content": "<p>I tried reducing the batch size and making predictions outside the training loop, it seems the problem is a combination of model size, memory and number of training epochs.</p>",
          "rawMarkdown": "I tried reducing the batch size and making predictions outside the training loop, it seems the problem is a combination of model size, memory and number of training epochs.",
          "votes": 2
        },
        {
          "id": 778630,
          "postDate": "2020-03-18T15:49:28.670Z",
          "content": "<p>did you try without drop_remainder=True?</p>",
          "rawMarkdown": "did you try without drop_remainder=True?"
        },
        {
          "id": 779481,
          "postDate": "2020-03-19T11:54:09.537Z",
          "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> yes, If I use small models (like EffNetB0) and small number of epochs (10) everything works fine.</p>",
          "rawMarkdown": "@mgornergoogle yes, If I use small models (like EffNetB0) and small number of epochs (10) everything works fine.",
          "votes": 1
        },
        {
          "id": 779781,
          "postDate": "2020-03-19T17:20:45.553Z",
          "content": "<p>If you can put together a small, simple reproducible example, I can file a bug. Of you can file a bug on your end and send me the link.</p>",
          "rawMarkdown": "If you can put together a small, simple reproducible example, I can file a bug. Of you can file a bug on your end and send me the link.",
          "votes": 1
        },
        {
          "id": 780097,
          "postDate": "2020-03-20T01:17:38.733Z",
          "content": "<p>Hey <a href=\"/mgornergoogle\">@mgornergoogle</a> here is a cleaner and smaller <a href=\"https://www.kaggle.com/dimitreoliveira/bug-report-unavailableerror-socket-closed/\">version of the notebook</a> I will file a bug soon and link here.</p>",
          "rawMarkdown": "Hey @mgornergoogle here is a cleaner and smaller [version of the notebook](https://www.kaggle.com/dimitreoliveira/bug-report-unavailableerror-socket-closed/) I will file a bug soon and link here.",
          "votes": 1
        },
        {
          "id": 781580,
          "postDate": "2020-03-21T13:09:40.077Z",
          "content": "<p>Hi <a href=\"/mgornergoogle\">@mgornergoogle</a> I've just filed the <a href=\"https://github.com/tensorflow/tensorflow/issues/37779\">bug here</a>, let me know if I can be of further help.</p>",
          "rawMarkdown": "Hi @mgornergoogle I've just filed the [bug here](https://github.com/tensorflow/tensorflow/issues/37779), let me know if I can be of further help.",
          "votes": 1
        },
        {
          "id": 782598,
          "postDate": "2020-03-22T13:22:19.420Z",
          "content": "<p><a href=\"/dimitreoliveira\">@dimitreoliveira</a> , would you mind to comment out the line <code>dataset = dataset.cache()</code> inside <code>get_validation_dataset()</code> and see if it fixes the issue or reduce the frequency?</p>",
          "rawMarkdown": "@dimitreoliveira , would you mind to comment out the line `dataset = dataset.cache()` inside `get_validation_dataset()` and see if it fixes the issue or reduce the frequency?",
          "votes": 2
        },
        {
          "id": 783021,
          "postDate": "2020-03-22T22:37:42.103Z",
          "content": "<p>Hey <a href=\"/yihdarshieh\">@yihdarshieh</a>  and <a href=\"/mgornergoogle\">@mgornergoogle</a>  I've got some experiments that seems to fix this issue, one thing that works is to do what <a href=\"/yihdarshieh\">@yihdarshieh</a> said and remove the <code>dataset = dataset.cache()</code> from the <code>get_validation_dataset()</code> function, And the other thing is to remove the validation loop inside the training phase, instead of calling it multiple times, if you call the validation step only once, the kernel does not crashes, it would be like the following:</p>\n\n<p><img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F1182060%2Fd4eab2c2bdebd86e94a5c6577121a60d%2FScreenshot%20from%202020-03-22%2019-36-13.png?generation=1584916603830184&amp;alt=media\" alt=\"\"></p>\n\n<p>So, for me it seems that the issue is related to memory allocation.</p>",
          "rawMarkdown": "Hey @yihdarshieh  and @mgornergoogle  I've got some experiments that seems to fix this issue, one thing that works is to do what @yihdarshieh said and remove the `dataset = dataset.cache()` from the `get_validation_dataset()` function, And the other thing is to remove the validation loop inside the training phase, instead of calling it multiple times, if you call the validation step only once, the kernel does not crashes, it would be like the following:\n\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F1182060%2Fd4eab2c2bdebd86e94a5c6577121a60d%2FScreenshot%20from%202020-03-22%2019-36-13.png?generation=1584916603830184&amp;alt=media)\n\nSo, for me it seems that the issue is related to memory allocation.",
          "votes": 2
        },
        {
          "id": 783997,
          "postDate": "2020-03-23T21:27:18.137Z",
          "content": "<p>Thank you for the bug report and the investigation.</p>",
          "rawMarkdown": "Thank you for the bug report and the investigation.",
          "votes": 1
        },
        {
          "id": 784002,
          "postDate": "2020-03-23T21:37:49.880Z",
          "content": "<p>You're welcome <a href=\"/mgornergoogle\">@mgornergoogle</a> , I think that if my supposition is correct, your example of custom training loop would also get an error if you run it for something like 30 epochs.</p>",
          "rawMarkdown": "You're welcome @mgornergoogle , I think that if my supposition is correct, your example of custom training loop would also get an error if you run it for something like 30 epochs.",
          "votes": 2
        }
      ]
    },
    {
      "id": 771480,
      "postDate": "2020-03-14T08:11:34.883Z",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> , this is amazing! Thanks.</p>\n\n<p>The reason that you have to scale the learning rate by 8 is because the callable <code>lr_schedule.lr_scaled</code> or <code>lr_schedule.lr</code> always return a fixed value, which is the value for <code>epoch=0</code>.</p>\n\n<p>I know that you use <code>lr_schedule.set_epoch(epoch)</code>, but it still doesn't work for some reason.</p>\n\n<p>I made a kernel (running on GPU, because when I use <code>tf.print</code> inside <code>TPU</code>, I got errors), there you can see the learning rate is always <code>1e-05</code> even in epoch 2 or epoch 3. Be aware that I change <code>lr_scaled()</code> to return <code>self.lrfn(self.epoch) * 1</code>, I just made a quick try, but it's not the point of the issue. </p>\n\n<p>I think that when you use <code>optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr_scaled)</code> and that optimizer is used inside a <code>tf.function</code> (so in graph mode), you can't really change the internal status of <code>lr_schedule</code> in <code>optimizer</code>.</p>\n\n<p>Here is the kernel</p>\n\n<p><a href=\"https://www.kaggle.com/yihdarshieh/problematic-lr/comments?scriptVersionId=30152935\">problematic lr</a></p>\n\n<p>I also published a quick fix by subclassing <code>tf.keras.optimizers.schedules.LearningRateSchedule</code> like</p>\n\n<pre><code>class LRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n\n    def __init__(self):\n\n        super(LRSchedule, self).__init__()\n\n    def __call__(self, step):\n\n        epoch = step // STEPS_PER_EPOCH\n\n        c1 = epoch &amp;lt; LR_RAMPUP_EPOCHS        \n        c2 = tf.math.logical_and(epoch &amp;lt; LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS, epoch &amp;gt;= LR_RAMPUP_EPOCHS)\n        c3 = epoch &amp;gt;= LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS\n\n        lr1 = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n        lr2 = LR_MAX\n        lr3 = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n\n        lr = tf.cast(c1, dtype=tf.float32) * lr1 + tf.cast(c2, dtype=tf.float32) * lr2 + tf.cast(c3, dtype=tf.float32) * lr3\n\n        return lr  \n</code></pre>\n\n<p>You can find in the following kernel that the convergence is OK now.</p>\n\n<p><a href=\"https://www.kaggle.com/yihdarshieh/quick-fix-lr?scriptVersionId=30156432\">quick fix lr</a></p>",
      "rawMarkdown": "@mgornergoogle , this is amazing! Thanks.\n\nThe reason that you have to scale the learning rate by 8 is because the callable `lr_schedule.lr_scaled` or `lr_schedule.lr` always return a fixed value, which is the value for `epoch=0`.\n\nI know that you use `lr_schedule.set_epoch(epoch)`, but it still doesn't work for some reason.\n\nI made a kernel (running on GPU, because when I use `tf.print` inside `TPU`, I got errors), there you can see the learning rate is always `1e-05` even in epoch 2 or epoch 3. Be aware that I change `lr_scaled()` to return `self.lrfn(self.epoch) * 1`, I just made a quick try, but it's not the point of the issue. \n\nI think that when you use `optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr_scaled)` and that optimizer is used inside a `tf.function` (so in graph mode), you can't really change the internal status of `lr_schedule` in `optimizer`.\n\nHere is the kernel\n\n[problematic lr](https://www.kaggle.com/yihdarshieh/problematic-lr/comments?scriptVersionId=30152935)\n\nI also published a quick fix by subclassing `tf.keras.optimizers.schedules.LearningRateSchedule` like\n\n    class LRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n        \n        def __init__(self):\n            \n            super(LRSchedule, self).__init__()\n                        \n        def __call__(self, step):\n\n            epoch = step // STEPS_PER_EPOCH\n        \n            c1 = epoch &lt; LR_RAMPUP_EPOCHS        \n            c2 = tf.math.logical_and(epoch &lt; LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS, epoch &gt;= LR_RAMPUP_EPOCHS)\n            c3 = epoch &gt;= LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS\n            \n            lr1 = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n            lr2 = LR_MAX\n            lr3 = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n        \n            lr = tf.cast(c1, dtype=tf.float32) * lr1 + tf.cast(c2, dtype=tf.float32) * lr2 + tf.cast(c3, dtype=tf.float32) * lr3\n        \n            return lr  \n\nYou can find in the following kernel that the convergence is OK now.\n\n[quick fix lr](https://www.kaggle.com/yihdarshieh/quick-fix-lr?scriptVersionId=30156432)",
      "votes": 2,
      "replies": [
        {
          "id": 775643,
          "postDate": "2020-03-16T22:49:16.133Z",
          "content": "<p>Thanks for the fix. Your solution is correct and I have updated the sample accordingly.</p>\n\n<p>I had a long exchange with the Tensorflow team about this. What happened before was:\n- I was using a Python variable for the epoch, stored in an lr_schdule object.\n- the lr_schedule.lr function got traced (compiled) for TPU when first used in:<br>\n<code>tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr)</code>\n- the only value this captured was epoch=0 so I was in effect using a constant learning rate</p>",
          "rawMarkdown": "Thanks for the fix. Your solution is correct and I have updated the sample accordingly.\n\nI had a long exchange with the Tensorflow team about this. What happened before was:\n- I was using a Python variable for the epoch, stored in an lr_schdule object.\n- the lr_schedule.lr function got traced (compiled) for TPU when first used in:<br>\n`tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr)`\n- the only value this captured was epoch=0 so I was in effect using a constant learning rate",
          "votes": 3
        }
      ]
    },
    {
      "id": 2235867,
      "postDate": "2023-04-26T11:54:18.793Z",
      "content": "<p>Thanks for sharing all these info. It would be useful</p>",
      "rawMarkdown": "Thanks for sharing all these info. It would be useful"
    },
    {
      "id": 883704,
      "postDate": "2020-06-12T20:56:00.863Z",
      "content": "<p>Superbly helpful <a href=\"/mgornergoogle\">@mgornergoogle</a>! Just a quick ask (maybe foolish): what prevents the TPU from achieving higher MXU values if idle time is zero (i.e., it is not waiting for data)? \nI am using the TPU optimized custom loop with XLMRoberta-large for jigsaw and MXU tops at 15% with 0% idle time. Does this relate to memory constraints (the batch size) at all?   </p>",
      "rawMarkdown": "Superbly helpful @mgornergoogle! Just a quick ask (maybe foolish): what prevents the TPU from achieving higher MXU values if idle time is zero (i.e., it is not waiting for data)? \nI am using the TPU optimized custom loop with XLMRoberta-large for jigsaw and MXU tops at 15% with 0% idle time. Does this relate to memory constraints (the batch size) at all?   "
    },
    {
      "id": 845909,
      "postDate": "2020-05-13T13:40:12.580Z",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> from my understanding about <a href=\"https://www.tensorflow.org/api_docs/python/tf/function\">@tf.function</a>, a tf graph is executed when the function written under it is called.</p>\n\n<p>So, in optimized custom training loop, <code>train_step(train_data_iter)</code> should be the line where the model gets trained.</p>\n\n<p>But when I checked, in an epoch,\nthe time taken by lines - \n<code>\ntrain_step(train_data_iter)\nvalid_step(valid_data_iter) \n</code>\nis 50 milliseconds</p>\n\n<p>and \nthe time taken by lines - \n<code>\nhistory.history['sparse_categorical_accuracy'].append(train_accuracy.result().numpy())\nhistory.history['val_sparse_categorical_accuracy'].append(valid_accuracy.result().numpy())\nhistory.history['loss'].append(train_loss.result().numpy() / (BATCH_SIZE*epoch_steps))\nhistory.history['val_loss'].append(valid_loss.result().numpy() / (BATCH_SIZE*valid_epoch_steps))\n</code>\nis 26 seconds</p>\n\n<p>If this second part is not the main part, Is it possible to reduce its time?</p>",
      "rawMarkdown": "@mgornergoogle from my understanding about [@tf.function](https://www.tensorflow.org/api_docs/python/tf/function), a tf graph is executed when the function written under it is called.\n\nSo, in optimized custom training loop, `train_step(train_data_iter)` should be the line where the model gets trained.\n\nBut when I checked, in an epoch,\nthe time taken by lines - \n```\ntrain_step(train_data_iter)\nvalid_step(valid_data_iter) \n```\nis 50 milliseconds\n\nand \nthe time taken by lines - \n```\nhistory.history['sparse_categorical_accuracy'].append(train_accuracy.result().numpy())\nhistory.history['val_sparse_categorical_accuracy'].append(valid_accuracy.result().numpy())\nhistory.history['loss'].append(train_loss.result().numpy() / (BATCH_SIZE*epoch_steps))\nhistory.history['val_loss'].append(valid_loss.result().numpy() / (BATCH_SIZE*valid_epoch_steps))\n```\nis 26 seconds\n\nIf this second part is not the main part, Is it possible to reduce its time?",
      "replies": [
        {
          "id": 846172,
          "postDate": "2020-05-13T15:54:03.590Z",
          "content": "<p>I noticed that too but I suspect there is some kind of lazy execution going on, where results are not computed until actually needed. The correct time is spent, but it is indeed spent in an unexpected place. Speculation only. I didn't get to the bottom of this.</p>",
          "rawMarkdown": "I noticed that too but I suspect there is some kind of lazy execution going on, where results are not computed until actually needed. The correct time is spent, but it is indeed spent in an unexpected place. Speculation only. I didn't get to the bottom of this.",
          "votes": 1
        }
      ]
    },
    {
      "id": 844358,
      "postDate": "2020-05-12T15:14:08.767Z",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> do you know if I can train LightGBM model or any other kind of GBM boosters using a TPU (does it have to TF, Kera or Pytorch models only)? Could you please provide an example or two on how to go about it?</p>",
      "rawMarkdown": "@mgornergoogle do you know if I can train LightGBM model or any other kind of GBM boosters using a TPU (does it have to TF, Kera or Pytorch models only)? Could you please provide an example or two on how to go about it?",
      "replies": [
        {
          "id": 855099,
          "postDate": "2020-05-20T15:36:58.550Z",
          "content": "<p>it looks like you would have to write your own implementation on GBM boosters using a TPU. some examples from other Kagglers (using other models) that might help with how to go about this are:</p>\n\n<ul>\n<li><a href=\"https://www.kaggle.com/c/flower-classification-with-tpus/discussion/132935\">CutMix and MixUp on GPU//TPU</a></li>\n<li><a href=\"https://www.kaggle.com/c/flower-classification-with-tpus/discussion/132191\">Rotation Augmentation on GPU//TPU</a></li>\n</ul>",
          "rawMarkdown": "it looks like you would have to write your own implementation on GBM boosters using a TPU. some examples from other Kagglers (using other models) that might help with how to go about this are:\n\n- [CutMix and MixUp on GPU//TPU](https://www.kaggle.com/c/flower-classification-with-tpus/discussion/132935)\n- [Rotation Augmentation on GPU//TPU](https://www.kaggle.com/c/flower-classification-with-tpus/discussion/132191)",
          "votes": 1
        }
      ]
    },
    {
      "id": 775673,
      "postDate": "2020-03-16T23:52:01.613Z",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> , I looked the 2 versions of custom training code several times, and I still have a question that what actually makes this big difference. The following is my guess, I am not sure at all, but it would be great if you can explain in more detail, so we can know more about the secret of TPU.</p>\n\n<p>It seems to me that, in the extreme optimized version, the whole graph <code>train\\_step</code> is run on the remote machine where the TPU actually locates on. This include <code>next(data_iter)</code>, so the data processing is done remotely, and when the data is ready, it can be used immediately by TPU. In the other version, the data processing in on local machine (on Kaggle, you said it's done on other VM which are more powoerful??), so it takes more time to be used by TPU?</p>",
      "rawMarkdown": "@mgornergoogle , I looked the 2 versions of custom training code several times, and I still have a question that what actually makes this big difference. The following is my guess, I am not sure at all, but it would be great if you can explain in more detail, so we can know more about the secret of TPU.\n\nIt seems to me that, in the extreme optimized version, the whole graph `train\\_step` is run on the remote machine where the TPU actually locates on. This include `next(data_iter)`, so the data processing is done remotely, and when the data is ready, it can be used immediately by TPU. In the other version, the data processing in on local machine (on Kaggle, you said it's done on other VM which are more powoerful??), so it takes more time to be used by TPU?",
      "replies": [
        {
          "id": 775704,
          "postDate": "2020-03-17T00:15:08.350Z",
          "content": "<p>A couple more keys for understanding:\n- The TPU only executes Tensorflow graphs of operations (not Python code)\n- the @tf.function annotation is what transforms a piece of Python into a graph.</p>\n\n<p>So the tf.range() loop is running on the TPU itself instead of going back and forth between the TPU and the Kaggle VM.</p>\n\n<p>I don't believe that the place where tf.data.Dataset operations are executed changes between the two implementations. </p>",
          "rawMarkdown": "A couple more keys for understanding:\n- The TPU only executes Tensorflow graphs of operations (not Python code)\n- the @tf.function annotation is what transforms a piece of Python into a graph.\n\nSo the tf.range() loop is running on the TPU itself instead of going back and forth between the TPU and the Kaggle VM.\n\nI don't believe that the place where tf.data.Dataset operations are executed changes between the two implementations. ",
          "votes": 4
        },
        {
          "id": 775754,
          "postDate": "2020-03-17T00:39:55.470Z",
          "content": "<p>But inside tf.range, there is a call to next(data_iter), which gets the next batch in the tf.data.dataset. So if this loop is converted to a graph and running entirely on TPU, shouldn't the necessary computation to get the next batch  also run on TPU?? Otherwise, in every step in tf.range, there is still data processing on local VM, then it has no difference to the other version.</p>\n\n<p>I hope you can get why I feel confused.</p>",
          "rawMarkdown": "But inside tf.range, there is a call to next(data_iter), which gets the next batch in the tf.data.dataset. So if this loop is converted to a graph and running entirely on TPU, shouldn't the necessary computation to get the next batch  also run on TPU?? Otherwise, in every step in tf.range, there is still data processing on local VM, then it has no difference to the other version.\n\nI hope you can get why I feel confused."
        },
        {
          "id": 775768,
          "postDate": "2020-03-17T00:47:20.990Z",
          "content": "<p>I should have said \"on the TPU side\" instead of saying \"on the TPU\". The TPU side is what lies on the other side of the XLA compiler and it is comprised of a traditional VM with a PCI-attached TPU. The data pipeline runs on this VM while forward and backward passes through the model happen on the TPU chip itself.</p>",
          "rawMarkdown": "I should have said \"on the TPU side\" instead of saying \"on the TPU\". The TPU side is what lies on the other side of the XLA compiler and it is comprised of a traditional VM with a PCI-attached TPU. The data pipeline runs on this VM while forward and backward passes through the model happen on the TPU chip itself.",
          "votes": 4
        },
        {
          "id": 775782,
          "postDate": "2020-03-17T00:56:30.937Z",
          "rawMarkdown": "",
          "isDeleted": true
        },
        {
          "id": 775791,
          "postDate": "2020-03-17T01:02:43.350Z",
          "content": "<p>Looks clear now! Thanks! To be sure, in the normal custom training version, the data processing is not on the TPU side, but on local VM that execute the python code. Right?</p>",
          "rawMarkdown": "Looks clear now! Thanks! To be sure, in the normal custom training version, the data processing is not on the TPU side, but on local VM that execute the python code. Right?"
        },
        {
          "id": 779783,
          "postDate": "2020-03-19T17:22:14.860Z",
          "content": "<p>No, any code running inside of tf.data.Dataset is running on the TPU side.</p>",
          "rawMarkdown": "No, any code running inside of tf.data.Dataset is running on the TPU side.",
          "votes": 1
        }
      ]
    },
    {
      "id": 771375,
      "postDate": "2020-03-14T05:13:38.880Z",
      "content": "<p>Wow, nice work Martin. </p>",
      "rawMarkdown": "Wow, nice work Martin. "
    },
    {
      "id": 775378,
      "postDate": "2020-03-16T16:13:05.340Z",
      "rawMarkdown": "",
      "isDeleted": true
    },
    {
      "id": 771428,
      "postDate": "2020-03-14T06:44:45.600Z",
      "rawMarkdown": "",
      "isDeleted": true
    },
    {
      "id": 775263,
      "postDate": "2020-03-16T13:13:41.010Z",
      "content": "<p>Thanks for the great sharing!</p>",
      "rawMarkdown": "Thanks for the great sharing!"
    },
    {
      "id": 771273,
      "postDate": "2020-03-14T00:40:29.863Z",
      "content": "<p>Super great. Thanks.</p>",
      "rawMarkdown": "Super great. Thanks."
    }
  ],
  "comments": [
    {
      "id": 775647,
      "author_name": "Martin Görner",
      "author_url": "",
      "post_date": "2020-03-16T22:55:01.327000",
      "content": "<p>This optimization shaves 25% from the TPU training time on a single TPU (8 cores). Given the complexity of the code compared to a simple model.fit(), it might not always be worth it. On TPU pods however (32 cores and more), the training time can go down by up to 50% (twice faster).</p>",
      "votes": 8,
      "replies": []
    },
    {
      "id": 779064,
      "author_name": "Kurian Benoy",
      "author_url": "",
      "post_date": "2020-03-19T02:08:06.300000",
      "content": "<p>Thanks for the great sharing. Tensorflow GradientTape() method is really useful!</p>\n\n<p>Did you calculate the idle time when model is finally trained? Or as shown in figure won't there be a small difference in Idle time as in your screenshot between method1(2nd epoch) and method2(3rd epoch)?</p>",
      "votes": 1,
      "replies": []
    },
    {
      "id": 771256,
      "author_name": "DimitreOliveira",
      "author_url": "",
      "post_date": "2020-03-14T00:12:29.040000",
      "content": "<p>This is awesome <a href=\"/mgornergoogle\">@mgornergoogle</a> , I'll try for sure!</p>",
      "votes": 1,
      "replies": []
    },
    {
      "id": 775339,
      "author_name": "DimitreOliveira",
      "author_url": "",
      "post_date": "2020-03-16T15:30:00.450000",
      "content": "<p>I have created a notebook using theses optimizations and applying to K-fold, in case anyone wanna checkout <a href=\"https://www.kaggle.com/dimitreoliveira/flower-with-tpus-k-fold-optimized-training-loops\">here it is</a>, during my experimentations I have faced some issues related to GCS availability, anyone got similar issues?</p>",
      "votes": 2,
      "replies": [
        {
          "id": 775670,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-16T23:45:54.863000",
          "content": "<p>Nice notebook. Can you fix the LRSchedule issue in it so that my errors from the first version do not continue to live there ? There are a couple of comments from my original notebook that can now be deleted too. Seach for \"scaled up by 8\" and \"multiply the learning rate schedule by 8\"</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 775835,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-17T01:33:31.747000",
          "content": "<p>Just did it, thanks for the head-up <a href=\"/mgornergoogle\">@mgornergoogle</a> , still, I have no idea why sometimes I get \"UnavailableError: Socket closed\" </p>",
          "votes": 2,
          "replies": []
        },
        {
          "id": 776735,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-17T15:43:52.493000",
          "content": "<p>That might happen when the XLA compiler crashes, for example when it runs out of memory on a large model. Lowering the batch size in that case sometimes helps.</p>",
          "votes": 2,
          "replies": []
        },
        {
          "id": 777543,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-17T18:57:45.600000",
          "content": "<p>This seems very strange because it can crash with any model architecture, but usually when I use <code>mode.predict()</code>, maybe the right way is to execute predictions outside of the training loop to save memory, I will do some more experimentations and report back.</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 777577,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-17T19:33:45.950000",
          "content": "<p>I recently filed a bug where I had a repeatable crash caused by the drop_remainder=True setting. Maybe you are hitting this bug too ? The setting is not required on repeated datasets but is supposed to make things a little faster. Try without it.</p>",
          "votes": 2,
          "replies": []
        },
        {
          "id": 777871,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-18T02:00:36.657000",
          "content": "<p>I tried reducing the batch size and making predictions outside the training loop, it seems the problem is a combination of model size, memory and number of training epochs.</p>",
          "votes": 2,
          "replies": []
        },
        {
          "id": 778630,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-18T15:49:28.670000",
          "content": "<p>did you try without drop_remainder=True?</p>",
          "votes": 0,
          "replies": []
        },
        {
          "id": 779481,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-19T11:54:09.537000",
          "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> yes, If I use small models (like EffNetB0) and small number of epochs (10) everything works fine.</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 779781,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-19T17:20:45.553000",
          "content": "<p>If you can put together a small, simple reproducible example, I can file a bug. Of you can file a bug on your end and send me the link.</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 780097,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-20T01:17:38.733000",
          "content": "<p>Hey <a href=\"/mgornergoogle\">@mgornergoogle</a> here is a cleaner and smaller <a href=\"https://www.kaggle.com/dimitreoliveira/bug-report-unavailableerror-socket-closed/\">version of the notebook</a> I will file a bug soon and link here.</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 781580,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-21T13:09:40.077000",
          "content": "<p>Hi <a href=\"/mgornergoogle\">@mgornergoogle</a> I've just filed the <a href=\"https://github.com/tensorflow/tensorflow/issues/37779\">bug here</a>, let me know if I can be of further help.</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 782598,
          "author_name": "Yih-Dar SHIEH",
          "author_url": "",
          "post_date": "2020-03-22T13:22:19.420000",
          "content": "<p><a href=\"/dimitreoliveira\">@dimitreoliveira</a> , would you mind to comment out the line <code>dataset = dataset.cache()</code> inside <code>get_validation_dataset()</code> and see if it fixes the issue or reduce the frequency?</p>",
          "votes": 2,
          "replies": []
        },
        {
          "id": 783021,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-22T22:37:42.103000",
          "content": "<p>Hey <a href=\"/yihdarshieh\">@yihdarshieh</a>  and <a href=\"/mgornergoogle\">@mgornergoogle</a>  I've got some experiments that seems to fix this issue, one thing that works is to do what <a href=\"/yihdarshieh\">@yihdarshieh</a> said and remove the <code>dataset = dataset.cache()</code> from the <code>get_validation_dataset()</code> function, And the other thing is to remove the validation loop inside the training phase, instead of calling it multiple times, if you call the validation step only once, the kernel does not crashes, it would be like the following:</p>\n\n<p><img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F1182060%2Fd4eab2c2bdebd86e94a5c6577121a60d%2FScreenshot%20from%202020-03-22%2019-36-13.png?generation=1584916603830184&amp;alt=media\" alt=\"\"></p>\n\n<p>So, for me it seems that the issue is related to memory allocation.</p>",
          "votes": 2,
          "replies": []
        },
        {
          "id": 783997,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-23T21:27:18.137000",
          "content": "<p>Thank you for the bug report and the investigation.</p>",
          "votes": 1,
          "replies": []
        },
        {
          "id": 784002,
          "author_name": "DimitreOliveira",
          "author_url": "",
          "post_date": "2020-03-23T21:37:49.880000",
          "content": "<p>You're welcome <a href=\"/mgornergoogle\">@mgornergoogle</a> , I think that if my supposition is correct, your example of custom training loop would also get an error if you run it for something like 30 epochs.</p>",
          "votes": 2,
          "replies": []
        }
      ]
    },
    {
      "id": 771480,
      "author_name": "Yih-Dar SHIEH",
      "author_url": "",
      "post_date": "2020-03-14T08:11:34.883000",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> , this is amazing! Thanks.</p>\n\n<p>The reason that you have to scale the learning rate by 8 is because the callable <code>lr_schedule.lr_scaled</code> or <code>lr_schedule.lr</code> always return a fixed value, which is the value for <code>epoch=0</code>.</p>\n\n<p>I know that you use <code>lr_schedule.set_epoch(epoch)</code>, but it still doesn't work for some reason.</p>\n\n<p>I made a kernel (running on GPU, because when I use <code>tf.print</code> inside <code>TPU</code>, I got errors), there you can see the learning rate is always <code>1e-05</code> even in epoch 2 or epoch 3. Be aware that I change <code>lr_scaled()</code> to return <code>self.lrfn(self.epoch) * 1</code>, I just made a quick try, but it's not the point of the issue. </p>\n\n<p>I think that when you use <code>optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr_scaled)</code> and that optimizer is used inside a <code>tf.function</code> (so in graph mode), you can't really change the internal status of <code>lr_schedule</code> in <code>optimizer</code>.</p>\n\n<p>Here is the kernel</p>\n\n<p><a href=\"https://www.kaggle.com/yihdarshieh/problematic-lr/comments?scriptVersionId=30152935\">problematic lr</a></p>\n\n<p>I also published a quick fix by subclassing <code>tf.keras.optimizers.schedules.LearningRateSchedule</code> like</p>\n\n<pre><code>class LRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n\n    def __init__(self):\n\n        super(LRSchedule, self).__init__()\n\n    def __call__(self, step):\n\n        epoch = step // STEPS_PER_EPOCH\n\n        c1 = epoch &amp;lt; LR_RAMPUP_EPOCHS        \n        c2 = tf.math.logical_and(epoch &amp;lt; LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS, epoch &amp;gt;= LR_RAMPUP_EPOCHS)\n        c3 = epoch &amp;gt;= LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS\n\n        lr1 = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n        lr2 = LR_MAX\n        lr3 = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n\n        lr = tf.cast(c1, dtype=tf.float32) * lr1 + tf.cast(c2, dtype=tf.float32) * lr2 + tf.cast(c3, dtype=tf.float32) * lr3\n\n        return lr  \n</code></pre>\n\n<p>You can find in the following kernel that the convergence is OK now.</p>\n\n<p><a href=\"https://www.kaggle.com/yihdarshieh/quick-fix-lr?scriptVersionId=30156432\">quick fix lr</a></p>",
      "votes": 2,
      "replies": [
        {
          "id": 775643,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-16T22:49:16.133000",
          "content": "<p>Thanks for the fix. Your solution is correct and I have updated the sample accordingly.</p>\n\n<p>I had a long exchange with the Tensorflow team about this. What happened before was:\n- I was using a Python variable for the epoch, stored in an lr_schdule object.\n- the lr_schedule.lr function got traced (compiled) for TPU when first used in:<br>\n<code>tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr)</code>\n- the only value this captured was epoch=0 so I was in effect using a constant learning rate</p>",
          "votes": 3,
          "replies": []
        }
      ]
    },
    {
      "id": 2235867,
      "author_name": "Степан Муленок",
      "author_url": "",
      "post_date": "2023-04-26T11:54:18.793000",
      "content": "<p>Thanks for sharing all these info. It would be useful</p>",
      "votes": 0,
      "replies": []
    },
    {
      "id": 883704,
      "author_name": "Rishabh Jha",
      "author_url": "",
      "post_date": "2020-06-12T20:56:00.863000",
      "content": "<p>Superbly helpful <a href=\"/mgornergoogle\">@mgornergoogle</a>! Just a quick ask (maybe foolish): what prevents the TPU from achieving higher MXU values if idle time is zero (i.e., it is not waiting for data)? \nI am using the TPU optimized custom loop with XLMRoberta-large for jigsaw and MXU tops at 15% with 0% idle time. Does this relate to memory constraints (the batch size) at all?   </p>",
      "votes": 0,
      "replies": []
    },
    {
      "id": 845909,
      "author_name": "Prateek",
      "author_url": "",
      "post_date": "2020-05-13T13:40:12.580000",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> from my understanding about <a href=\"https://www.tensorflow.org/api_docs/python/tf/function\">@tf.function</a>, a tf graph is executed when the function written under it is called.</p>\n\n<p>So, in optimized custom training loop, <code>train_step(train_data_iter)</code> should be the line where the model gets trained.</p>\n\n<p>But when I checked, in an epoch,\nthe time taken by lines - \n<code>\ntrain_step(train_data_iter)\nvalid_step(valid_data_iter) \n</code>\nis 50 milliseconds</p>\n\n<p>and \nthe time taken by lines - \n<code>\nhistory.history['sparse_categorical_accuracy'].append(train_accuracy.result().numpy())\nhistory.history['val_sparse_categorical_accuracy'].append(valid_accuracy.result().numpy())\nhistory.history['loss'].append(train_loss.result().numpy() / (BATCH_SIZE*epoch_steps))\nhistory.history['val_loss'].append(valid_loss.result().numpy() / (BATCH_SIZE*valid_epoch_steps))\n</code>\nis 26 seconds</p>\n\n<p>If this second part is not the main part, Is it possible to reduce its time?</p>",
      "votes": 0,
      "replies": [
        {
          "id": 846172,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-05-13T15:54:03.590000",
          "content": "<p>I noticed that too but I suspect there is some kind of lazy execution going on, where results are not computed until actually needed. The correct time is spent, but it is indeed spent in an unexpected place. Speculation only. I didn't get to the bottom of this.</p>",
          "votes": 1,
          "replies": []
        }
      ]
    },
    {
      "id": 844358,
      "author_name": "Mani Sarkar",
      "author_url": "",
      "post_date": "2020-05-12T15:14:08.767000",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> do you know if I can train LightGBM model or any other kind of GBM boosters using a TPU (does it have to TF, Kera or Pytorch models only)? Could you please provide an example or two on how to go about it?</p>",
      "votes": 0,
      "replies": [
        {
          "id": 855099,
          "author_name": "Jesse Mostipak",
          "author_url": "",
          "post_date": "2020-05-20T15:36:58.550000",
          "content": "<p>it looks like you would have to write your own implementation on GBM boosters using a TPU. some examples from other Kagglers (using other models) that might help with how to go about this are:</p>\n\n<ul>\n<li><a href=\"https://www.kaggle.com/c/flower-classification-with-tpus/discussion/132935\">CutMix and MixUp on GPU//TPU</a></li>\n<li><a href=\"https://www.kaggle.com/c/flower-classification-with-tpus/discussion/132191\">Rotation Augmentation on GPU//TPU</a></li>\n</ul>",
          "votes": 1,
          "replies": []
        }
      ]
    },
    {
      "id": 775673,
      "author_name": "Yih-Dar SHIEH",
      "author_url": "",
      "post_date": "2020-03-16T23:52:01.613000",
      "content": "<p><a href=\"/mgornergoogle\">@mgornergoogle</a> , I looked the 2 versions of custom training code several times, and I still have a question that what actually makes this big difference. The following is my guess, I am not sure at all, but it would be great if you can explain in more detail, so we can know more about the secret of TPU.</p>\n\n<p>It seems to me that, in the extreme optimized version, the whole graph <code>train\\_step</code> is run on the remote machine where the TPU actually locates on. This include <code>next(data_iter)</code>, so the data processing is done remotely, and when the data is ready, it can be used immediately by TPU. In the other version, the data processing in on local machine (on Kaggle, you said it's done on other VM which are more powoerful??), so it takes more time to be used by TPU?</p>",
      "votes": 0,
      "replies": [
        {
          "id": 775704,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-17T00:15:08.350000",
          "content": "<p>A couple more keys for understanding:\n- The TPU only executes Tensorflow graphs of operations (not Python code)\n- the @tf.function annotation is what transforms a piece of Python into a graph.</p>\n\n<p>So the tf.range() loop is running on the TPU itself instead of going back and forth between the TPU and the Kaggle VM.</p>\n\n<p>I don't believe that the place where tf.data.Dataset operations are executed changes between the two implementations. </p>",
          "votes": 4,
          "replies": []
        },
        {
          "id": 775754,
          "author_name": "Yih-Dar SHIEH",
          "author_url": "",
          "post_date": "2020-03-17T00:39:55.470000",
          "content": "<p>But inside tf.range, there is a call to next(data_iter), which gets the next batch in the tf.data.dataset. So if this loop is converted to a graph and running entirely on TPU, shouldn't the necessary computation to get the next batch  also run on TPU?? Otherwise, in every step in tf.range, there is still data processing on local VM, then it has no difference to the other version.</p>\n\n<p>I hope you can get why I feel confused.</p>",
          "votes": 0,
          "replies": []
        },
        {
          "id": 775768,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-17T00:47:20.990000",
          "content": "<p>I should have said \"on the TPU side\" instead of saying \"on the TPU\". The TPU side is what lies on the other side of the XLA compiler and it is comprised of a traditional VM with a PCI-attached TPU. The data pipeline runs on this VM while forward and backward passes through the model happen on the TPU chip itself.</p>",
          "votes": 4,
          "replies": []
        },
        {
          "id": 775782,
          "author_name": "",
          "author_url": "",
          "post_date": "2020-03-17T00:56:30.937000",
          "content": "",
          "votes": 0,
          "replies": []
        },
        {
          "id": 775791,
          "author_name": "Yih-Dar SHIEH",
          "author_url": "",
          "post_date": "2020-03-17T01:02:43.350000",
          "content": "<p>Looks clear now! Thanks! To be sure, in the normal custom training version, the data processing is not on the TPU side, but on local VM that execute the python code. Right?</p>",
          "votes": 0,
          "replies": []
        },
        {
          "id": 779783,
          "author_name": "Martin Görner",
          "author_url": "",
          "post_date": "2020-03-19T17:22:14.860000",
          "content": "<p>No, any code running inside of tf.data.Dataset is running on the TPU side.</p>",
          "votes": 1,
          "replies": []
        }
      ]
    },
    {
      "id": 771375,
      "author_name": "Chris Deotte",
      "author_url": "",
      "post_date": "2020-03-14T05:13:38.880000",
      "content": "<p>Wow, nice work Martin. </p>",
      "votes": 0,
      "replies": []
    },
    {
      "id": 775378,
      "author_name": "",
      "author_url": "",
      "post_date": "2020-03-16T16:13:05.340000",
      "content": "",
      "votes": 0,
      "replies": []
    },
    {
      "id": 771428,
      "author_name": "",
      "author_url": "",
      "post_date": "2020-03-14T06:44:45.600000",
      "content": "",
      "votes": 0,
      "replies": []
    },
    {
      "id": 775263,
      "author_name": "Tsai29",
      "author_url": "",
      "post_date": "2020-03-16T13:13:41.010000",
      "content": "<p>Thanks for the great sharing!</p>",
      "votes": 0,
      "replies": []
    },
    {
      "id": 771273,
      "author_name": "Marília Prata",
      "author_url": "",
      "post_date": "2020-03-14T00:40:29.863000",
      "content": "<p>Super great. Thanks.</p>",
      "votes": 0,
      "replies": []
    }
  ],
  "raw_markdown_by_id": {
    "771219": "**Edit: the learning rate issue highlighed by @yihdarshieh in the comments has been fixed**\n\nAll three models for this comparison are in this notebook: [Custom Training Loop with 100+ flowers on TPU](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu)\n\n### 1) Keras model.fit\n\nFine-tuning a Keras flowers model on TPU is quite fast but if you open the TPU profiler (click on the gauges in the top-right corner of your running notebook), you see that it could be a bit faster:\n<p>\n<img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F4003597%2F4959d4a9c096730989ec288f450cac26%2Ftrain%20fit%20TPU%20idle%20time.png?generation=1584134538223267&amp;alt=media\" width=\"300ps\">\n</p>\n\n25% idle time means that the TPU is waiting. It could be waiting for data. Or there could be some overhead in the way Keras is running its training loop.\n\n### 2) Custom training loop\n\nLet's try to write a custom training loop. Even in a distributed environment like a TPU, this can be written quite easily with two additional functions offered by the TPU distributions strategy: distribute_dataset and run.\n\n```\n# training step function\n@tf.function\ndef train_step(images, labels):\n    with tf.GradientTape() as tape:\n        probabilities = model(images, training=True)\n        loss = loss_fn(labels, probabilities)\n    grads = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n\n# distributed dataset\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\n\n# custom training loop\nfor (images, labels) in train_dist_ds:\n    strategy.experimental_run_v2(train_step, args=(images, labels))\n```\nFor full code, including validation, refer to [Custom Training Loop with 100+ flowers on TPU, section custom training loop](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Custom-training-loop)\n\nLet's look at TPU idle time in the custom training loop:\n<p>\n<img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F4003597%2Ff3f20c836c4d3cecf4f0d2f1a0f5d0c2%2Fsimple%20CTL%20TPU%20idle%20time.png?generation=1584135998683394&amp;alt=media\" width=\"300ps\">\n</p>\n\nNot really an improvement...\n\n### 3) TPU-optimized custom training loop\n\nThere is still room for optimization, the TPU is a remote accelerator and calling it has a cost. What if we could call it less often. If one training step does more work directly on the TPU, things could go a bit faster. This code uses a loop in the step function. It uses `tf.range()` instead of `range()` so that it can be compiled to Tensorflow operation and sent to the TPU. The compilation is forced by `@tf.function`.\n\n```\n# training step function\n@tf.function\ndef train_step(data_iter):\n    def train_step_fn(images, labels):\n        with tf.GradientTape() as tape:\n            probabilities = model(images, training=True)\n            loss = loss_fn(labels, probabilities)\n        grads = tape.gradient(loss, model.trainable_variables)\n        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n\n    # loop executing on TPU\n    for _ in tf.range(STEPS_PER_TPU_CALL):\n        strategy.experimental_run_v2(train_step_fn, next(data_iter))\n\n# distributed dataset\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\n# dataset iterator\ntrain_data_iter = iter(train_dist_ds)\n\n# custom training loop\nwhile True:\n    train_step(train_data_iter)\n```\n\nFor full code, including validation, refer to [Custom Training Loop with 100+ flowers on TPU, section optimized custom training loop](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Optimized-custom-training-loop)\n\nLet's look at the idle time with this TPU-optimized custom training loop:\n<p>\n<img src=\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F4003597%2F817f6ab062a59abe1556f030b0a75491%2Foptimized%20CTL%20TPU%20idle%20time.png?generation=1584136547845529&amp;alt=media\" width=\"300ps\">\n</p>\n\nNo more idle time ! This confirms that the data pipeline was not the problem. \n\n### Benchmark\nFrom the notebook [Custom Training Loop with 100+ flowers on TPU](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu):\n\n| Training loop type | TPU idle time | Total training time | Typical epoch time |\n|--------------------|--------------|-------------------|-------------------|\n| Keras model.fit() | 25% | 417 s | 30 s |\n| simple custom training loop | 23% | 436 s | 32 s |\n| TPU-optimized custom training loop | 00% | 350 s | 22 s |",
    "775647": "This optimization shaves 25% from the TPU training time on a single TPU (8 cores). Given the complexity of the code compared to a simple model.fit(), it might not always be worth it. On TPU pods however (32 cores and more), the training time can go down by up to 50% (twice faster).",
    "779064": "Thanks for the great sharing. Tensorflow GradientTape() method is really useful!\n\nDid you calculate the idle time when model is finally trained? Or as shown in figure won't there be a small difference in Idle time as in your screenshot between method1(2nd epoch) and method2(3rd epoch)?",
    "771256": "This is awesome @mgornergoogle , I'll try for sure!",
    "775339": "I have created a notebook using theses optimizations and applying to K-fold, in case anyone wanna checkout [here it is](https://www.kaggle.com/dimitreoliveira/flower-with-tpus-k-fold-optimized-training-loops), during my experimentations I have faced some issues related to GCS availability, anyone got similar issues?",
    "771480": "@mgornergoogle , this is amazing! Thanks.\n\nThe reason that you have to scale the learning rate by 8 is because the callable `lr_schedule.lr_scaled` or `lr_schedule.lr` always return a fixed value, which is the value for `epoch=0`.\n\nI know that you use `lr_schedule.set_epoch(epoch)`, but it still doesn't work for some reason.\n\nI made a kernel (running on GPU, because when I use `tf.print` inside `TPU`, I got errors), there you can see the learning rate is always `1e-05` even in epoch 2 or epoch 3. Be aware that I change `lr_scaled()` to return `self.lrfn(self.epoch) * 1`, I just made a quick try, but it's not the point of the issue. \n\nI think that when you use `optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule.lr_scaled)` and that optimizer is used inside a `tf.function` (so in graph mode), you can't really change the internal status of `lr_schedule` in `optimizer`.\n\nHere is the kernel\n\n[problematic lr](https://www.kaggle.com/yihdarshieh/problematic-lr/comments?scriptVersionId=30152935)\n\nI also published a quick fix by subclassing `tf.keras.optimizers.schedules.LearningRateSchedule` like\n\n    class LRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n        \n        def __init__(self):\n            \n            super(LRSchedule, self).__init__()\n                        \n        def __call__(self, step):\n\n            epoch = step // STEPS_PER_EPOCH\n        \n            c1 = epoch &lt; LR_RAMPUP_EPOCHS        \n            c2 = tf.math.logical_and(epoch &lt; LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS, epoch &gt;= LR_RAMPUP_EPOCHS)\n            c3 = epoch &gt;= LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS\n            \n            lr1 = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n            lr2 = LR_MAX\n            lr3 = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n        \n            lr = tf.cast(c1, dtype=tf.float32) * lr1 + tf.cast(c2, dtype=tf.float32) * lr2 + tf.cast(c3, dtype=tf.float32) * lr3\n        \n            return lr  \n\nYou can find in the following kernel that the convergence is OK now.\n\n[quick fix lr](https://www.kaggle.com/yihdarshieh/quick-fix-lr?scriptVersionId=30156432)",
    "2235867": "Thanks for sharing all these info. It would be useful",
    "883704": "Superbly helpful @mgornergoogle! Just a quick ask (maybe foolish): what prevents the TPU from achieving higher MXU values if idle time is zero (i.e., it is not waiting for data)? \nI am using the TPU optimized custom loop with XLMRoberta-large for jigsaw and MXU tops at 15% with 0% idle time. Does this relate to memory constraints (the batch size) at all?   ",
    "845909": "@mgornergoogle from my understanding about [@tf.function](https://www.tensorflow.org/api_docs/python/tf/function), a tf graph is executed when the function written under it is called.\n\nSo, in optimized custom training loop, `train_step(train_data_iter)` should be the line where the model gets trained.\n\nBut when I checked, in an epoch,\nthe time taken by lines - \n```\ntrain_step(train_data_iter)\nvalid_step(valid_data_iter) \n```\nis 50 milliseconds\n\nand \nthe time taken by lines - \n```\nhistory.history['sparse_categorical_accuracy'].append(train_accuracy.result().numpy())\nhistory.history['val_sparse_categorical_accuracy'].append(valid_accuracy.result().numpy())\nhistory.history['loss'].append(train_loss.result().numpy() / (BATCH_SIZE*epoch_steps))\nhistory.history['val_loss'].append(valid_loss.result().numpy() / (BATCH_SIZE*valid_epoch_steps))\n```\nis 26 seconds\n\nIf this second part is not the main part, Is it possible to reduce its time?",
    "844358": "@mgornergoogle do you know if I can train LightGBM model or any other kind of GBM boosters using a TPU (does it have to TF, Kera or Pytorch models only)? Could you please provide an example or two on how to go about it?",
    "775673": "@mgornergoogle , I looked the 2 versions of custom training code several times, and I still have a question that what actually makes this big difference. The following is my guess, I am not sure at all, but it would be great if you can explain in more detail, so we can know more about the secret of TPU.\n\nIt seems to me that, in the extreme optimized version, the whole graph `train\\_step` is run on the remote machine where the TPU actually locates on. This include `next(data_iter)`, so the data processing is done remotely, and when the data is ready, it can be used immediately by TPU. In the other version, the data processing in on local machine (on Kaggle, you said it's done on other VM which are more powoerful??), so it takes more time to be used by TPU?",
    "771375": "Wow, nice work Martin. ",
    "775378": "",
    "771428": "",
    "775263": "Thanks for the great sharing!",
    "771273": "Super great. Thanks."
  }
}