Keras: The Gold Standard for TensorFlow

Keras：Tensorflow的黄金标准 苏剑林 · 2019-11-06 · Machine Learning Keras

Over the past two weeks I've put quite a lot of effort into developing bert4keras. Besides some API standardization work, most of the effort went into building the pretraining code. As of yesterday, the pretraining code is basically complete, and I've tested it successfully on both TPU and multi-GPU setups. So for anyone who has the ambition (and the compute) to improve pretrained models, there's now one more option—this may be the clearest and most understandable implementation of BERT and its pretraining code currently available.

Pretraining code link: https://github.com/bojone/bert4keras/tree/master/pretraining

After these two weeks of development (bug-fixing), my biggest takeaway is this: Keras has become the gold standard for TensorFlow. As long as you write your code following Keras conventions, you can migrate it to tf.keras with ease, and from there, training on TPU or multi-GPU setups becomes remarkably painless—truly a "do it once, benefit forever" situation. On the other hand, if your code is too clever, including many of the "creative hack" style Keras tricks I've introduced before, you may run into quite a few problems. It's even possible that code which already runs fine on multi-GPU simply refuses to work on TPU no matter what you try.

Keras and Tensorflowmore

Relentless Support

Everyone says TensorFlow 2.0 is pushing tf.keras as its main API, but in fact, tf.keras has been the gold standard since TensorFlow 1.14. So if you want to experience just how relentlessly Google supports Keras, all you need is TensorFlow 1.14+—you don't necessarily need to upgrade to 2.0. Currently, bert4keras supports both the original Keras and tf.keras. For typical single-GPU fine-tuning tasks, you can use either; but if you want multi-GPU or TPU training, tf.keras is definitely the better choice.

To get started with tf.keras, I'd first recommend a really good website: https://tf.wiki/

In tf.keras, turning a model from single-GPU to multi-GPU training is remarkably simple:

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = create_a_model()
    model.compile(loss='mse', optimizer='adam')

model.fit(train_x, train_y, epochs=10)

In other words, all you need to do is define a strategy, and then build your model under the scope of that strategy, and it automatically becomes a multi-GPU model. Multi-GPU training has never been this simple.

Incidentally, Keras itself comes with a built-in multi_gpu_model function for multi-GPU training, but from my own testing, multi_gpu_model isn't very reliable—sometimes it just doesn't work. In short, I'd still recommend tf.keras. Also, the example above is for single-machine multi-GPU; multi-machine multi-GPU is similar, but I don't have the corresponding setup to test it, so I won't give an example here. If you want to try it out, please refer to the introduction at https://tf.wiki/.

What about TPU? Just as simple—just swap out strategy (with a small change in TensorFlow 2.0):

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_address)
tf.config.experimental_connect_to_host(resolver.master())
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

Doesn't it feel almost unbelievably simple? Now you should understand what I meant earlier by "relentless" support for Keras. Since TensorFlow 1.14, as long as you write standard Keras code, things pretty much just work, everywhere.

What Counts as "Standard"?

I've been emphasizing "standard" Keras code all along—so what actually counts as standard? Here I'll summarize some of the lessons I've learned.

1. As far as possible, implement whatever functionality you need using Keras's built-in layers, loss functions, and optimizers. If a model is built entirely from Keras's built-in layers, loss functions, and optimizers, you can be reasonably confident it will run on multi-GPU or TPU without issues.
2. If you need to write custom layers, follow the conventions strictly, and in particular make sure to properly implement the get_config method. One way to test whether your implementation follows conventions correctly is: build a model using your custom layer, and see whether it can be successfully cloned by the clone_model function. If it can, your layer definition follows the conventions.
3. If you're training on TPU, don't use add_loss at the end of the model to define a custom loss, and don't use add_metric to add a metric. If you need to define a complex custom loss or metric, define it as the output of a layer instead—see this approach for reference.
4. If you're training on TPU, absolutely avoid dynamic (variable-length) code during training. For example, when using tf.where, the x and y arguments can't be None, since otherwise the length of tf.where's output becomes indeterminate. Also, almost any TensorFlow function with "dynamic" in its name is off-limits.

As you can see, "standard" here essentially means imitating Keras's existing conventions as closely as possible, and inventing as little as possible yourself. If you follow points 1 and 2, you can easily train with multi-GPU under tf.keras; points 3 and 4 are there specifically to avoid TPU pitfalls. In short, everything needs to be static.

Although TensorFlow 2.0 defaults to eager execution now, I personally don't recommend relying on it. I believe we should get used to the workflow of building models with static graphs. While eager execution is convenient for debugging, it makes us overly dependent on immediate output results, which weakens our debugging skills when facing more complex problems. Similarly, I don't recommend relying too heavily on tools like code completion or code suggestions—these tools can make us overly dependent on them, preventing us from truly understanding the functions we're using. (This is just my personal opinion, no need to get upset if you disagree.)

A Victory of Human-Centered Design

If I remember correctly, I first encountered Keras in early 2015. Back then there weren't many deep learning frameworks around; I just wanted a handy tool to implement a few simple models, and that's how I found Keras—and I've used it ever since. Not just me, but perhaps even Keras's own authors never imagined that Keras would one day become the gold standard for TensorFlow.

I don't think this is a coincidence. TensorFlow has had no shortage of higher-level API frameworks over the years—TensorLayer, tf.slim, TFLearn, to name a few—so why did Keras end up being chosen in the end? Beyond Keras's "long history," it's because Keras truly deserves recognition as an elegant piece of engineering. Over this past year, I've occasionally gone back to read Keras's source code, and every single time I'm struck by its rigor and elegance. It is, without question, a work of art—a genuinely human-centered creation.

So, this is a victory of human-centered design.