English (unofficial) translations of posts at kexue.fm
Source

Keras: The Gold Standard of TensorFlow

Translated by DeepSeek V4 Pro. Translations can be inaccurate, please refer to the original post for important stuff.

Jianlin Su November 6, 2019 #7055

Over the past two weeks, I have devoted considerable effort to the development of bert4keras. In addition to some API standardization work, the main workload has been constructing the pre-training portion of the code. As of yesterday, the pre-training code is basically complete and has been tested in both TPU and multi-GPU environments. This provides another option for students who are ambitious (and have the computing power) to improve pre-trained models. This is likely the most clear and easy-to-understand implementation of BERT and its pre-training code currently available.

Pre-training code link: https://github.com/bojone/bert4keras/tree/master/pretraining

After these two weeks of development (and bug-fixing), my biggest takeaway is that Keras has become the gold standard for TensorFlow. As long as your code is written according to Keras standards and specifications, it can be easily migrated to tf.keras, and subsequently trained in TPU or multi-GPU environments with great ease. It truly is a "once and for all" solution. Conversely, if your coding style is too flexible—including many of the "grafting" style Keras tricks I have introduced before—you may encounter many problems. It is even possible that even if you have successfully run your code on multiple GPUs, you might struggle to get it working on a TPU no matter how much you debug.

Keras and TensorFlow

Unsparing Support

Everyone says TensorFlow 2.0 promotes tf.keras as the primary focus, but in fact, starting from TensorFlow 1.14, tf.keras has already become its gold standard. Therefore, if you want to experience how Google is sparing no effort to support Keras, you only need TensorFlow 1.14+; you don’t necessarily need to upgrade to 2.0. Currently, the bert4keras code supports both the original Keras and tf.keras. For typical single-card fine-tuning tasks, you can use either, but if you want to use multi-GPU training or even TPU training, the best choice is tf.keras.

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

In tf.keras, converting a model from single-card to multi-card training is extremely 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, as long as you define a strategy and build your model within that strategy.scope(), it becomes a multi-card model. Multi-card training has never been this simple!

By the way, Keras itself comes with a multi_gpu_model function to implement multi-GPU training, but my personal tests found that multi_gpu_model is not very reliable and sometimes fails to take effect. In short, I still recommend using tf.keras. Additionally, the above is the syntax for single-machine multi-card training; multi-machine multi-card is similar, but I do not have the environment to test it, so I haven’t provided an example. For those who want to test it, please refer to the introduction in https://tf.wiki/.

What about TPUs? It’s just as simple—just replace the strategy (with minor changes 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)

Does it feel incredibly simple? By now, everyone should understand what I meant by "sparing no effort" to support Keras. Starting from TensorFlow 1.14, as long as you use standard Keras coding practices, you can almost always succeed.

What Counts as "Standard"?

I have repeatedly emphasized using standard Keras coding practices, but what exactly counts as standard? Here is a summary of some experiences:

  1. Use Keras’s built-in layers, loss functions, and optimizers as much as possible to implement the required functions. If a model consists entirely of Keras’s built-in components, it is basically guaranteed to work on multi-GPU or TPU.

  2. If you need to customize layers, follow the specifications strictly, especially the get_config method. One way to test if your implementation is standard is to build a model with your custom layer and see if it can be successfully cloned using the clone_model function. If it can, your layer definition is standardized.

  3. If you want to train on a TPU, do not use add_loss at the end of the model to define a custom loss, and do not use add_metric to add metrics. If you need custom complex losses or metrics, define them as layer outputs, referring to this approach.

  4. If you want to train on a TPU, avoid using dynamic (variable-length) logic during the training process. For example, when using tf.where, the parameters x and y cannot be None, otherwise the length of the result of tf.where is uncertain. Almost all functions in TensorFlow with the word "dynamic" in them cannot be used.

As you can see, the so-called "standard" is actually imitating existing Keras practices to the greatest extent possible and minimizing self-invention. As long as you follow points 1 and 2, you can easily use multi-GPU training under tf.keras; points 3 and 4 are specific pitfalls for TPUs. Generally speaking, everything must be static.

Although TensorFlow 2.0 has defaulted to eager execution (dynamic graphs), I do not recommend using it. Personally, I believe we should get used to the model construction process of static graphs. While dynamic graphs are convenient for debugging, they can make us heavily dependent on immediate output results, reducing our ability to debug complex problems. Similarly, I do not recommend using tools like code completion or code suggestions; these tools create too much dependency and prevent us from truly understanding the functions themselves. (This is a personal opinion; please don’t flame me.)

The Victory of Humanization

If I remember correctly, I first encountered Keras in early 2015. At that time, there weren’t many deep learning frameworks, and I just wanted a handy tool to implement a few simple models, so I found Keras and have used it ever since. Not just myself, but perhaps even the authors of Keras did not expect back then that Keras would become the gold standard of TensorFlow today.

I feel this is no accident. TensorFlow has had many high-level API frameworks, such as TensorLayer, tf.slim, and TFLearn, but why was Keras ultimately chosen? Besides Keras’s "long history," it is because Keras truly deserves to be called an elegant encapsulation. Over the past year, I have occasionally read the Keras source code, and every time I do, I am struck by its rigor and elegance. It is a masterpiece worthy of the name—a humanized creation.

Therefore, this is a victory for humanization.

When reposting, please include the original address of this article: https://kexue.fm/archives/7055

For more detailed reposting matters, please refer to: "Scientific Space FAQ"