2 ms·
You can use tensorflow's tf.data, I think.
by yablak 7y ago
You can use tensorflow's tf.data, I think.
- sillysaurusx 7y agoSort of. The way you do it is to scope operations to tf.device(None), which selects the TPU's CPU. (I think it's equivalent to using the first device in sess.list_devices(), which is the CPU.) You can scope operations in tf.data, sure, but you can also execute arbitrary training operations. I use this technique to train GPT-2 117M with a 25k context window, which requires about 300GB memory. Only a TPU's CPU has so much. That's why it was surprising to hear Jax can't do that. It's one of the best features of a TPU.