TensorFlow Lite 开发手册（2）——TensorFlow 2.0使用简介

（一）训练模型的一般流程

训练模型的基本流程如下：

• 将 Tensorflow 导入程序：

import tensorflow as tf

from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model

• 加载数据

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

• 使用 tf.data 来将数据集切分为 batch 以及shuffle数据集：

train_ds = tf.data.Dataset.from_tensor_slices(
    (x_train, y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

• 使用 Keras 模型子类化（model subclassing） API 构建 tf.keras 模型：

class MyModel(Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.conv1 = Conv2D(32, 3, activation='relu')
    self.flatten = Flatten()
    self.d1 = Dense(128, activation='relu')
    self.d2 = Dense(10, activation='softmax')

  def call(self, x):
    x = self.conv1(x)
    x = self.flatten(x)
    x = self.d1(x)
    return self.d2(x)

model = MyModel()

• 为训练选择优化器与损失函数：

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()

optimizer = tf.keras.optimizers.Adam()

• 选择衡量指标来度量模型的损失值（loss）和准确率（accuracy）。这些指标在 epoch 上累积值，然后打印出整体结果。

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

• 使用 tf.GradientTape 来训练模型：

@tf.function
def train_step(images, labels):
  with tf.GradientTape() as tape:
    predictions = model(images)
    loss = loss_object(labels, predictions)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

  train_loss(loss)
  train_accuracy(labels, predictions)

• 测试模型：

@tf.function
def test_step(images, labels):
  predictions = model(images)
  t_loss = loss_object(labels, predictions)

  test_loss(t_loss)
  test_accuracy(labels, predictions)
EPOCHS = 5

for epoch in range(EPOCHS):
  for images, labels in train_ds:
    train_step(images, labels)

  for test_images, test_labels in test_ds:
    test_step(test_images, test_labels)

  template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
  print (template.format(epoch+1,
                         train_loss.result(),
                         train_accuracy.result()*100,
                         test_loss.result(),
                         test_accuracy.result()*100))  

• 输出如下：

Epoch 1, Loss: 0.13822732865810394, Accuracy: 95.84833526611328, Test Loss: 0.07067110389471054, Test Accuracy: 97.75
Epoch 2, Loss: 0.09080979228019714, Accuracy: 97.25, Test Loss: 0.06446609646081924, Test Accuracy: 97.95999908447266
Epoch 3, Loss: 0.06777264922857285, Accuracy: 97.93944549560547, Test Loss: 0.06325332075357437, Test Accuracy: 98.04000091552734
Epoch 4, Loss: 0.054447807371616364, Accuracy: 98.33999633789062, Test Loss: 0.06611879169940948, Test Accuracy: 98.00749969482422
Epoch 5, Loss: 0.04556874558329582, Accuracy: 98.60433197021484, Test Loss: 0.06510476022958755, Test Accuracy: 98.10400390625

（二）训练示例模型

基于以下代码训练示例模型，并得到模型文件model.h5：

import tensorflow as tf
from tensorflow import keras

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = tf.expand_dims(x_train, 3)
y_train = keras.utils.to_categorical(y_train, num_classes=10)
datasets = tf.data.Dataset.from_tensor_slices((x_train, y_train))
datasets = datasets.repeat(1).batch(10)

img = keras.Input(shape=[28, 28, 1])

x = keras.layers.Conv2D(filters=64, kernel_size=4,strides=1, padding='SAME',activation='relu')(img)
x = keras.layers.AveragePooling2D(pool_size=2, strides=2, padding='SAME')(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(0.15)(x)

x = keras.layers.Conv2D(filters=64, kernel_size=4,strides=1, padding='SAME',activation='relu')(img)
x = keras.layers.MaxPooling2D(pool_size=2, strides=2, padding='SAME')(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(0.15)(x)

x = keras.layers.Flatten()(x)
x = keras.layers.Dense(512, activation='relu')(x)
y_pred = keras.layers.Dense(10, activation='softmax')(x)

model = keras.Model(inputs=img, outputs=y_pred)
model.compile(optimizer= keras.optimizers.Adam(0.01),
             loss= keras.losses.categorical_crossentropy,
             metrics = ['AUC', 'accuracy'])

model.fit(datasets, epochs=1)
model.save('/home/ai/model.h5')