Integrate with Cloud TPU about minigo HOT 13 CLOSED

@amj @Kashomon So to get started on this, I'm wondering what our TPU instance management looks like. Would we be spinning up a instance manually, git-pulling, and pip install'ing to be able to train on TPU? Or is there a kubeified solution to this?

from minigo.

Notes to self:

This issue is for training on TPUs. Selfplay on TPUs is still a difficult problem

Anticipated issues:

• rewriting DualNetworkTrainer to use Estimator api (should be easy; I architected the input_fn and network definition functions based on estimator api.)
• during estimator rewrite, figure out how to continue logging all the custom metrics we've implemented. Hopefully these metrics can continue to be exported in an efficient fashion during TPU operation.
• figuring out how many parallel readers we can use; how large the files should be for the parallel readers; whether a sloppy parallel reader can be used.
• profile the cloud TPU execution trace to verify we're getting good throughput.
• get an idea of how expensive TPU training is, overhead-wise, and related, figure out what's the optimal saving schedule.

from minigo.

TPUs for GKE are quickly approaching Alpha. I'll put more details here soon

from minigo.

@brilee for now though, spinning up an instance, attaching at TPU to it, pulling the code, etc, is the way to get moving. It should be pretty painless to setup and hopefully you shouldn't have to do it very often ;)

from minigo.

more elaboration/notes to self:

• rewriting DualNetworkTrainer to use Estimator api (hah, I thought this would be easy)
  • figure out how to initialize variable to previous model's weights, but save current run under a new model name.
    • estimator assumes that model_dir will contain all things related to the training run - checkpoints, saved models, and logs. this means that tensorboard logs may be broken up if we use a separate directory for each model generation... perhaps use a single model_dir but periodically create a model export to GCS? this would require also building in a way to reinitialize training from an exported GCS model.
  • figure out how to create a set of bootstrap weights
  • replace StatisticsCollector with eval_metric_ops({...tf.metrics.mean()}
  • create a training hook that evaluates the update ratios every so often
  • figure out save/checkpoint schedule
  • verify that tensorboard shows all the right things
  • figure out how to export model as SavedModel. (I don't think the default exporter works, because it exports a graph that takes as input a serialized tf.Example.)

from minigo.

This is definitely going to require some sort of baking-in period to ensure I haven't messed up my estimator rewrite in some way...

from minigo.

That's a pretty good list. is it worth breaking the work up to two people?

from minigo.

Don't think this is parallelizable. I'm pretty much ripping out 75% of dual_net.

from minigo.

I'm pretty much ripping out 75% of dual_net.

This comment worries me quite a bit. It sounds like we're going to sacrifice readability here to make this work on TPUs. Is that the tradeoff? Should we keep around the old implementation for clarity and because it works?

from minigo.

Estimator does a host of things for you, like managing checkpointing frequency, initializing model weights, setting up logs, and streaming eval metrics (which obsoletes the homegrown StatisticsCollector stuff). So the code will actually get a lot shorter/cleaner as a result. The main pain point is twisting Estimator to do some of the things we'd done manually, but don't really fit into how Estimator wants to do things.

from minigo.

The TPUEstimator's marketing pitch is that you just use Estimator, swap out the Estimator impl, for TPUEstimator, and it should Just Work. But we'll see :)

from minigo.

I stumbled across https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_main.py#L248 while looking at examples. Seems like getting summaries logged during training requires a workaround at the moment.

from minigo.

we have cloud tpu training up now.

from minigo.