Crash happened when running this cell:

Random performance without fine-tuning
get_accuracy(params_repl)

from vision_transformer.

Hi Tyler

Do you run out of CPU RAM or GPU/TPU RAM?

Also, how much RAM do you have?
(You can check with !free -mh)

The provided Colab works fine with CIFAR datasets and the default settings (default Colab currently has 12G of RAM).

from vision_transformer.

I cannot run the code on colab and I have the same problem when I try to use colab which crashed the blocks.

from vision_transformer.

I just checked, the Colab runs fine (at least up to the "Fine-tune" section that is below the "Random performance without fine-tuning." comment that you point out above).

from vision_transformer.

I find an error at the last two steps of Fine Tune:

The world's simplest training loop.

Completes in ~20 min on the TPU runtime.

for step, batch, lr_repl in zip(
tqdm.notebook.trange(1, total_steps + 1),
ds_train.as_numpy_iterator(),
lr_iter
):

opt_repl, loss_repl, update_rngs = update_fn_repl(
opt_repl, lr_repl, batch, update_rngs)

Thank you very much! I'm looking forward to your reply.

from vision_transformer.

Hi Tyler,

Crash happened when running this cell:

Random performance without fine-tuning
get_accuracy(params_repl)

Did you run
TPU setup : Boilerplate for connecting JAX to TPU?

If you double click this cell, you will find there are multiple codes there:

#@markdown TPU setup : Boilerplate for connecting JAX to TPU.

import os
if 'google.colab' in str(get_ipython()) and 'COLAB_TPU_ADDR' in os.environ:
  # Make sure the Colab Runtime is set to Accelerator: TPU.
  import requests
  if 'TPU_DRIVER_MODE' not in globals():
    url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20191206'
    resp = requests.post(url)
    TPU_DRIVER_MODE = 1

  # The following is required to use TPU Driver as JAX's backend.
  from jax.config import config
  config.FLAGS.jax_xla_backend = "tpu_driver"
  config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
  print('Registered TPU:', config.FLAGS.jax_backend_target)
else:
  print('No TPU detected. Can be changed under "Runtime/Change runtime type".')

I think these codes will register your TPU.

from vision_transformer.

@lucasliunju

I find an error at the last two steps of Fine Tune:

The world's simplest training loop.

Completes in ~20 min on the TPU runtime.

for step, batch, lr_repl in zip(
tqdm.notebook.trange(1, total_steps + 1),
ds_train.as_numpy_iterator(),
lr_iter
):

opt_repl, loss_repl, update_rngs = update_fn_repl(
opt_repl, lr_repl, batch, update_rngs)

Thank you very much! I'm looking forward to your reply.

#81 (comment)

To be specific, import flax.optim as optim

from vision_transformer.

Hi @GuardianWang

Thanks for your reply. That does work!
Did you know why the official optimizer cannot work?
I would like to ask whether you use this code on multi-host tpu (such as v3-32, v3-64).

Thank you very much!

from vision_transformer.

I just re-ran the code in the provided Colab on TPU and it worked without problem.

@Tylersuard can you try again?
(I think it was some temporary regression in the Colab setup and/or JAX code that was independent from the code in this repo)

from vision_transformer.

@andsteing I just tried again, same issue:

The world's simplest training loop.

Completes in ~20 min on the TPU runtime.

for step, batch, lr_repl in zip(
tqdm.notebook.trange(1, total_steps + 1),
ds_train.as_numpy_iterator(),
lr_iter
):

opt_repl, loss_repl, update_rngs = update_fn_repl(
opt_repl, lr_repl, batch, update_rngs)

Response:
Your session crashed after using all available RAM.

Log:

from vision_transformer.

Just to confirm:

1. You're loading this notebook here : https://colab.sandbox.google.com/github/google-research/vision_transformer/blob/master/vit_jax.ipynb
2. You're using a Runtime type = TPU kernel
3. You start with a fresh kernel.
4. You run end-to-end without modifications.

Because I tried multiple times and was not able to reproduce your error.

from vision_transformer.

I used a slightly different link, the one shown in this repo's readme:

https://colab.research.google.com/github/google-research/vision_transformer/blob/master/vit_jax.ipynb

from vision_transformer.

@andsteing Ok, I got it to work. For some reason, it does not work with the "high-ram" instance option, but it does work with the regular option. Thank you for your help.

from vision_transformer.