前言

Tensorflow 2.x (简称TF2.x) 发布也有一段时间了,TF1.x 时代的API对用户以及初学者非常不友好,笔者也是因为TF1.x API晦涩难懂一直没有深入学习TF而选择了PyTorch. 非常看好PyTorch的发展前景,但是由于工作需要着手学习TF2.x. 本文选择深度学习领域的HelloWord---MNIST手写体识别数据集做为开始. 主要内容:

• 了解TF中数据输入的Pipeline
• 学习tf.keras API创建CNN网络
• 学习tf.keras 中Loss, Optimizer, 以及网络训练的流程

---

Requirement

• OS: Ubuntu 18.04
• Tensorflow>=2.0.0

---

TF2.x MNIST 训练例子

数据集准备

Tensorflow keras中默认自带MNIST数据集

关于MNIST 数据集:

• x_train:(60000, 28, 28) 训练集(28x28x1的图像),60000张
• y_train:(60000,) 训练集对应的label
• x_test:(10000, 28, 28) 测试集(28x28x1的图像),10000张
• y_test:(10000,) 测试集对应的label

流程: 导入数据(numpy.ndarray格式) ----> 数据归一化 ----> 使用TF Dataset API读取,以及划分Batch

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model
from net import MyModel  # 自定义的脚本

# ----------------------Prepare Dataset------------------------------
mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
print('='*20)
print('x_train:{}'.format(x_train.shape))
print('y_train:{}'.format(y_train.shape))
print('x_test:{}'.format(x_test.shape))
print('y_test:{}'.format(y_test.shape))

# 数据归一化
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dim
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)

train_ds = tf.data.Dataset.from_tensor_slices(
    (x_test, y_test)
).batch(32)

创建CNN网络/模型

TF2.x 推荐采用tf.keras API创建模型, tf.keras API与PyTorch的torch.nn API用法和类似. 注意:TF中tensor默认存储方式为 NxHxWxC, 即channels_last

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


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


if __name__ == '__main__':
    model = MyModel()
    # dump_input = tf.keras.Input(shape=[28, 28, 3], batch_size=1)
    model.build(input_shape=(1, 28, 28, 3))
    model.summary()

调用model.summary() API可以打印出模型的结构:

Model: "my_model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              multiple                  896       
_________________________________________________________________
flatten (Flatten)            multiple                  0         
_________________________________________________________________
dense (Dense)                multiple                  2769024   
_________________________________________________________________
dense_1 (Dense)              multiple                  1290      
=================================================================
Total params: 2,771,210
Trainable params: 2,771,210
Non-trainable params: 0
_________________________________________________________________

定义Loss函数以及Loss优化器

# -----------------------Loss and Optimizer----------------------
# 创建Loss函数,以及Loss优化函数
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

定义测量指标,方便观察模型的准确度以及Loss变化

• Loss
• Accuracy

# ----------------------metrics----------------------------------
# 设置测量指标
# For train
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
# For test
test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')# ----------------------metrics----------------------------------
# 设置测量指标
# For train
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
# For test
test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

在整个数据集上进行训练

# 训练
# @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 eopch in  range(EPOCHS):
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()

    for images, labels in train_ds:
        train_step(images, labels)
    
    for images, labels in train_ds:
        test_step(images, labels)
    
    template = 'Epoch {}, Loss: {}, Accuarcy: {}, Test Loss: {}, Test Accuracy: {}'
    print(template.format(
        eopch + 1,
        train_loss.result(),
        train_accuracy.result() * 100,
        test_loss.result(),
        test_accuracy.result()*100 
    ))

训练结果

图片.png

完整代码

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model
from net import MyModel


# ----------------------Prepare Dataset------------------------------
mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
print('='*20)
print('x_train:{}'.format(x_train.shape))
print('y_train:{}'.format(y_train.shape))
print('x_test:{}'.format(x_test.shape))
print('y_test:{}'.format(y_test.shape))

# 数据归一化
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dim
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)

train_ds = tf.data.Dataset.from_tensor_slices(
    (x_test, y_test)
).batch(32)

# -----------------------Create a CNN model----------------------
model = MyModel()

# -----------------------Loss and Optimizer----------------------
# 创建Loss函数,以及Loss优化函数
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

# ----------------------metrics----------------------------------
# 设置测量指标
# For train
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
# For test
test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

# 训练
# @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 eopch in  range(EPOCHS):
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()

    for images, labels in train_ds:
        train_step(images, labels)
    
    for images, labels in train_ds:
        test_step(images, labels)
    
    template = 'Epoch {}, Loss: {}, Accuarcy: {}, Test Loss: {}, Test Accuracy: {}'
    print(template.format(
        eopch + 1,
        train_loss.result(),
        train_accuracy.result() * 100,
        test_loss.result(),
        test_accuracy.result()*100 
    ))