PyTorch TPU support

Kale-ab_Tessera: <p>Any news on PyTorch TPU support?</p>
<p>Last I heard (<a href="https://twitter.com/soumithchintala/status/1009112034242453506" rel="nofollow noopener">https://twitter.com/soumithchintala/status/1009112034242453506</a>).</p>

plm: <p>Was wondering the same… <a href="https://twitter.com/plouismarie/status/1048196291132514305" rel="nofollow noopener">https://twitter.com/plouismarie/status/1048196291132514305</a></p>

fmassa: <p>There was an announcement in the PyTorch Developer Conference this week, where a ResNet50 was successfully trained on a TPU.<br>
For more info, check <a href="https://cloud.google.com/blog/products/ai-machine-learning/introducing-pytorch-across-google-cloud">https://cloud.google.com/blog/products/ai-machine-learning/introducing-pytorch-across-google-cloud</a></p>

plm: <p>That’s great to read, thank you very much for sharing.</p>
<blockquote>
<p>we’re pleased to announce that engineers on Google’s TPU team are actively collaborating with core PyTorch developers to connect PyTorch to Cloud TPUs</p>
</blockquote>
<p>There is already some keras code to show how to do TPU in a very concrete and practical way :</p>
<pre><code class="lang-auto">tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)
tpu_model.compile(
    optimizer=tf.train.AdamOptimizer(learning_rate=1e-3, ),
    loss=tf.keras.losses.sparse_categorical_crossentropy,
    metrics=['sparse_categorical_accuracy']
)

def train_gen(batch_size):
  while True:
    offset = np.random.randint(0, x_train.shape[0] - batch_size)
    yield x_train[offset:offset+batch_size], y_train[offset:offset + batch_size]
    

tpu_model.fit_generator(
    train_gen(1024),
    epochs=10,
    steps_per_epoch=100,
    validation_data=(x_test, y_test),
)
</code></pre>
<p>I am hoping to see some pytorch demo codes that enables TPU usage.<br>
I guess it is coming soon. Can not wait !!!</p>

DuaneNielsen: <p>Anyone heard anything new on Pytorch on TPU?</p>

RicCu: <p>I believe it’s already usable but a little rough around the edges, although I haven’t tried it myself. <a href="https://github.com/pytorch/xla" rel="nofollow noopener">https://github.com/pytorch/xla</a></p>

Zuanazzi: <p>Any developments on that in the mean while? XLA is a nice temporary solution, but it would be good to know if we can expect an official solution soon.</p>

Kale-ab_Tessera: <p>Example: <a href="https://github.com/pytorch/xla/tree/master/contrib/colab" rel="nofollow noopener">https://github.com/pytorch/xla/tree/master/contrib/colab</a></p>

williamFalcon: <p>Thanks to the awesome work by the PyTorch and XLA team, we were able to get TPU support fully working out of the box with PyTorch Lightning.</p>
<p><a href="https://medium.com/@_willfalcon/from-pytorch-to-pytorch-lightning-a-gentle-introduction-b371b7caaf09" rel="nofollow noopener">Check out this nifty guide on going from PyTorch to PyTorch Lightning</a></p>
<aside class="onebox twitterstatus">
  <header class="source">
      <a href="https://twitter.com/PyTorchLightnin/status/1232813118507692033?s=20" target="_blank" rel="nofollow noopener">twitter.com</a>
  </header>
  <article class="onebox-body">
    <img src="https://pbs.twimg.com/profile_images/1157284227001540609/TtIfU1GJ_400x400.jpg" class="thumbnail onebox-avatar" width="400" height="400">

<h4><a href="https://twitter.com/PyTorchLightnin/status/1232813118507692033?s=20" target="_blank" rel="nofollow noopener">PyTorchLightning (PyTorchLightnin)</a></h4>

<div class="tweet"> You can now run <a href="https://twitter.com/PyTorch" target="_blank" rel="nofollow noopener">@PyTorch</a> code on TPUs and GPUs with Lightning... wait for it... 

WITHOUT CODE CHANGES!!! 

Retweet if this is useful to you! <a href="https://t.co/vyfFeAl5Yq" target="_blank" rel="nofollow noopener">https://t.co/vyfFeAl5Yq</a><div class="tweet-images"><iframe class="tweet-video" src="https://twitter.com/i/videos/1232813118507692033" width="695" height="390" frameborder="0" allowfullscreen="" seamless="seamless" sandbox="allow-same-origin allow-scripts allow-forms allow-popups allow-popups-to-escape-sandbox allow-presentation"></iframe></div>
</div>

<div class="date">
  <a href="https://twitter.com/PyTorchLightnin/status/1232813118507692033?s=20" target="_blank" rel="nofollow noopener">11:42 PM - 26 Feb 2020</a>
    <span class="like">
      <svg viewbox="0 0 512 512" width="14px" height="16px" aria-hidden="true">
        <path d="M462.3 62.6C407.5 15.9 326 24.3 275.7 76.2L256 96.5l-19.7-20.3C186.1 24.3 104.5 15.9 49.7 62.6c-62.8 53.6-66.1 149.8-9.9 207.9l193.5 199.8c12.5 12.9 32.8 12.9 45.3 0l193.5-199.8c56.3-58.1 53-154.3-9.8-207.9z"></path>
      </svg> 620
    </span>
    <span class="retweet">
      <svg viewbox="0 0 640 512" width="14px" height="16px" aria-hidden="true">
        <path d="M629.657 343.598L528.971 444.284c-9.373 9.372-24.568 9.372-33.941 0L394.343 343.598c-9.373-9.373-9.373-24.569 0-33.941l10.823-10.823c9.562-9.562 25.133-9.34 34.419.492L480 342.118V160H292.451a24.005 24.005 0 0 1-16.971-7.029l-16-16C244.361 121.851 255.069 96 276.451 96H520c13.255 0 24 10.745 24 24v222.118l40.416-42.792c9.285-9.831 24.856-10.054 34.419-.492l10.823 10.823c9.372 9.372 9.372 24.569-.001 33.941zm-265.138 15.431A23.999 23.999 0 0 0 347.548 352H160V169.881l40.416 42.792c9.286 9.831 24.856 10.054 34.419.491l10.822-10.822c9.373-9.373 9.373-24.569 0-33.941L144.971 67.716c-9.373-9.373-24.569-9.373-33.941 0L10.343 168.402c-9.373 9.373-9.373 24.569 0 33.941l10.822 10.822c9.562 9.562 25.133 9.34 34.419-.491L96 169.881V392c0 13.255 10.745 24 24 24h243.549c21.382 0 32.09-25.851 16.971-40.971l-16.001-16z"></path>
      </svg> 200
    </span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


Borda: <p>Does pytorch TPU supports also float16 precision?</p>

111305: <p>Yes. Moreover, it switches to 16 bit by default.</p>