AIDog改造手记：使用TensorFlow 2.0
之前开发过一款人工智能的微信小程序，其功能是拍一张狗狗的照片，识别出狗狗的类别。程序虽小，功能虽然单一，但五脏俱全，涉及到机器学习的各个方面，以及微信小程序的开发，非常适合作为机器学习的上手项目。这个项目是一边学习，一边写出来的，在这个过程中，进一步掌握了机器学习的知识，具体情况请参考我之前写的文章：
当微信小程序遇上TensorFlow：Server端实现补充
当微信小程序遇上TensorFlow：接收base64编码图像数据
项目地址：
https://github.com/mogoweb/AIDog
最近一段时间一直在研究微信小程序中的tensorflow.js，取得了一些成果，有兴趣可以阅读阅读一下：
当微信小程序遇上TensorFlow - tensorflow.js篇
经过这些研究，确定在微信小程序中使用TensorFlow是可行的，接下来，我准备将AIDog小程序改造一番，主要改造两点：
- 将训练模型的python脚本改造为使用TensorFlow 2.0 API编写
- 使用tensorflow.js取代原来的微信小程序 + TensorFlow serving模式
经过紧锣密鼓的开发，TensorFlow已经进化到2.0 beta版本，虽然不是最终正式版，但到了beta版本，API不会再有大的变化，其实是可以尝试一下的。关于TensorFlow 2.0 版本，可以参看我之前写的文章：
[译]高效的TensorFlow 2.0：应用最佳实践以及有什么变化
[译]标准化Keras：TensorFlow 2.0中的高级API指南
构建模型
AIDog是机器学习中图像分类的一种应用，有现成的分类模型可以借用，采用迁移学习，在新的数据集上重新训练模型。回过头去看以前的retrain.py脚本，写得相当复杂，当时我也是根据TensorFlow文档，在现有脚本上修改。这次使用TensorFlow 2.0进行改写，当然采用推荐的keras接口进行实现。一尝试，发现采用keras接口实现，太简洁了：
def build_model(num_classes):
# Create the base model from the pre-trained model Inception V3
base_model = keras.applications.InceptionV3(input_shape=IMG_SHAPE,
# We cannot use the top classification layer of the pre-trained model as it contains 1000 classes.
# It also restricts our input dimensions to that which this model is trained on (default: 299x299)
include_top=False,
weights='imagenet')
base_model.trainable = False
# Using Sequential API to stack up the layers
model = keras.Sequential([
base_model,
keras.layers.GlobalAveragePooling2D(),
keras.layers.Dense(num_classes,
activation='softmax')
])
# Compile the model to configure training parameters
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
return model代码相当简短，第一步加载预训练好的Inception V3模型，需要注意一点，Inception V3是在ImageNet数据集上进行训练，包含1000个分类，而在我们的应用中，并没有这么多分类，所以需要传入 include_top=False ，其含义是不包含最后一个softmax层。
接下来一行语句
base_model.trainable = False表示基础模型中的参数不参与训练，其实这也容易理解，模型已经在ImageNet数据集上进行训练，学会了从图像提取特征，这是在超大规模数据集上训练出的参数，我们只需训练后面自行增加层的参数。如果你希望Inception模型某些层参与训练，以更好的匹配新的数据集，你也可以这样写：
# Unfreeze all layers of InceptionV3
base_model.trainable = True
# Refreeze layers until the layers we want to fine-tune
for layer in base_model.layers[:100]:
layer.trainable = False可以自行决定多少个层不参与参数更新。要记住一点，参与训练的参数越多，训练速度就越慢。
接下来一行代码，在基础模型Inception V3的基础上加入一个平均池化层和全连接层，为什么这样定义？是参考网上的资料，最简单的方法可以仅仅加入一个使用softmax激活函数的全连接层。
数据处理
对于图像预处理，在原来的retrain.py脚本中，处理得非常复杂，在tensorflow 2.0中，可以采用tf提供的解码和缩放函数：
def preprocess_image(image):
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.resize(image, [299, 299])
image /= 255.0 # normalize to [0,1] range
# normalized to the[-1, 1] range
image = 2 * image - 1
return imageDataset API是TensorFlow 1.3版本中引入的一个新的模块，主要服务于数据读取，构建输入数据的pipeline。如果想要用到TensorFlow新出的Eager模式，就必须要使用Dataset API来读取数据。
图像数据文件按照9:1的比例划分为训练数据集和验证数据集。虽然dataset提供了shuffle方法随机打乱输入数据，但实际测试下来，特别耗费内存，会因为内存耗尽无法进行后续的模型训练，一种解决方法就是在读取图片文件列表时，打乱顺序：
random.shuffle(all_image_paths)虽然在后面的fit调用中会警告Dataset未进行shuffle，但由于我们在传入文件列表时，就已经随机打乱了次序，可以忽略这个警告：
# shuffle already when loading images, so no need shuffle dataset
ds = train_image_label_ds
ds = ds.repeat()
ds = ds.batch(BATCH_SIZE)
# `prefetch` lets the dataset fetch batches, in the background while the model is training.
ds = ds.prefetch(buffer_size=AUTOTUNE)
model.fit(ds,
epochs=FLAGS.epochs,
steps_per_epoch=steps_per_epoch,
validation_data=validate_image_label_ds.repeat().batch(BATCH_SIZE),
validation_steps=validation_steps,
callbacks=[tensorboard_callback, model_checkpoint_callback])训练模型
涉及到这种图像分类的深度学习模型，通常训练起来非常耗费时间，特别是没有强劲的GPU的条件下，几乎无法进行像样的模型训练。对于本项目采用的Inception V3模型，属于那种大型深度学习模型，虽然不是训练全部的参数，但仅仅前向传递计算，就非常耗时，如果没有GTX 1080这种级别以上的显卡，不要轻易尝试。
天无绝人之路，这个时候我们可以薅一薅Google的羊毛，之前我写过一篇文章：
详细介绍过如何使用谷歌GPU云计算平台。有一点需要注意，Google Colab目前默认使用的是TensorFlow r1.14的版本，如果要使用TensorFlow 2.0 beta版本，需要在开始位置执行：
!pip install tensorflow-gpu==2.0.0-beta1训练结束，可以在google drive上看到TensorFlow saved model格式的模型。
至此，狗狗的分类模型训练完毕。接下来，需要将saved model格式的模型，转换为tensorflow.js可用模型，且听下回分解。
以上完整源代码，可以访问我google云端硬盘：
https://colab.research.google.com/drive/1KSEky1xfBP5-R5WwUoYdpmXy2K5JzL5t
腾讯云开发者
