tensorflow2.x主要的两种编程方式

tensorflow从1.x变到2.x的过程中，编程方式发生了很大的改变，且从tensorflow2.x开始官方推荐tf.keras的编程方式，这样的变化让很多初学者非常困惑，从网上随便搜索一下，可以搜索到各种版本的tensorflow编程方式，且1.x的教程偏多，让初学者不知该选择哪种版本。这里总结了我认为比较受欢迎的两种编程方式，分别是基于tf.keras的编程方式和基于tensorflow2.x原生的编程方式，可以用于验证tensorflow是否安装成功。

基于tf.keras的编程方式

基于tf.keras的编程方式集成度比较高，不适合基于自己的需求进行修改，所以该方式仅适合于初学者，如下是两个代码，第二个代码稍微增加了一些可定制性。

import tensorflow as tf
import numpy as np
import os
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# 0. 准备数据
mnist = tf.keras.datasets.mnist
((train_x, train_y), (test_x, test_y)) = mnist.load_data()
train_x, test_x = np.expand_dims(train_x/255, 3), np.expand_dims(test_x/255, 3)

# 1. 定义模型
input = tf.keras.layers.Input(shape=(28,28,1))
conv1 = tf.keras.layers.Conv2D(filters=8, kernel_size=3, strides=1, padding="valid", activation="relu", use_bias=True)(input)
conv2 = tf.keras.layers.Conv2D(filters=4, kernel_size=3, strides=2, padding="valid", activation="relu", use_bias=True)(conv1)
fc0 = tf.keras.layers.Flatten()(conv2)
fc0 = tf.keras.layers.Dropout(0.25)(fc0)
fc1 = tf.keras.layers.Dense(units=256, activation="relu", use_bias=True)(fc0)
output = tf.keras.layers.Dense(units=10, activation="softmax", use_bias=True)(fc1)
model = tf.keras.Model(input, output)
model.summary()

# 2. 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(), 
                loss='sparse_categorical_crossentropy', 
                metrics=['sparse_categorical_accuracy'])

# 3. 训练模型
model.fit(train_x, train_y, epochs=5)

# 4. 测试模型
model.evaluate(test_x, test_y)

import tensorflow as tf
from tensorflow.keras.callbacks import Callback
import numpy as np
import os
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# 0. 准备数据
mnist = tf.keras.datasets.mnist
((train_x, train_y), (test_x, test_y)) = mnist.load_data()
train_x, test_x = np.expand_dims(train_x/255, 3), np.expand_dims(test_x/255, 3)

# 1. 定义模型
input = tf.keras.layers.Input(shape=(28,28,1))
conv1 = tf.keras.layers.Conv2D(filters=8, kernel_size=3, strides=1, padding="valid", activation="relu", use_bias=True)(input)
conv2 = tf.keras.layers.Conv2D(filters=4, kernel_size=3, strides=2, padding="valid", activation="relu", use_bias=True)(conv1)
fc0 = tf.keras.layers.Flatten()(conv2)
fc0 = tf.keras.layers.Dropout(0.25)(fc0)
fc1 = tf.keras.layers.Dense(units=256, activation="relu", use_bias=True)(fc0)
output = tf.keras.layers.Dense(units=10, activation="softmax", use_bias=True)(fc1)
model = tf.keras.Model(input, output)
model.summary()

# 2. 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(), 
                loss='sparse_categorical_crossentropy', 
                metrics=['sparse_categorical_accuracy'])

# 3. 自定义回调函数
# 学习率函数
def scheduler(epoch):
    if epoch < 2:
        return 0.001
    else:
        return 0.001 * tf.math.exp(0.1 * (2 - epoch))
learningreatescheduler = tf.keras.callbacks.LearningRateScheduler(scheduler)
# 自定义函数
class MyCallback(Callback):
    def __init__(self):
        super().__init__()
    def on_epoch_end(self, num, logs: dict):
        # print("On train end (%d, %s)" % (num, logs))
        print("On train end (epuch: %d, lr: %0.5f, loss: %0.4f, acc: %0.4f, val_loss: %0.4f, val_loss: %0.4f)" % 
                (num, logs["lr"], logs["loss"], logs["sparse_categorical_accuracy"], logs["val_loss"], logs["val_sparse_categorical_accuracy"]))
        print("")
        model.save_weights('./tmp/model_file_%d.h5' % num)
        return

# 4. 训练模型
# model.fit(train_x, train_y, epochs=5)
history = model.fit(train_x, train_y, 
                    validation_data=(test_x, test_y), 
                    epochs=5, verbose=2, 
                    callbacks=[learningreatescheduler, MyCallback()])  # 回调函数的顺序有影响，需要先学习率再自定义的，这样logs里才有lr
print(history.history)

# 5. 测试模型
loss, acc = model.evaluate(test_x, test_y)

基于tensorflow的编程方式

基于tensorflow的编程方式，可以选择原生的语句进行模型定义，但是这样比较繁琐，所以推荐采用tf.keras的方式进行模型定义，然后模型训练采用tensorflow原生的方式，这样可定制性比较高。

import tensorflow as tf
import numpy as np
import time
import os
import shutil


# 0. 准备数据
mnist = tf.keras.datasets.mnist
((train_x, train_y), (test_x, test_y)) = mnist.load_data()
train_x, test_x = np.expand_dims(train_x/255, 3), np.expand_dims(test_x/255, 3)
train_y, test_y = tf.one_hot(tf.constant(train_y), depth=10).numpy(), tf.one_hot(tf.constant(test_y), depth=10).numpy()

# 1. 定义模型
input = tf.keras.layers.Input(shape=(28,28,1))
conv1 = tf.keras.layers.Conv2D(filters=8, kernel_size=3, strides=1, padding="valid", activation="relu", use_bias=True)(input)
conv2 = tf.keras.layers.Conv2D(filters=4, kernel_size=3, strides=2, padding="valid", activation="relu", use_bias=True)(conv1)
fc0 = tf.keras.layers.Flatten()(conv2)
fc0 = tf.keras.layers.Dropout(0.25)(fc0)
fc1 = tf.keras.layers.Dense(units=256, activation="relu", use_bias=True)(fc0)
output = tf.keras.layers.Dense(units=10, activation="softmax", use_bias=True)(fc1)
model = tf.keras.Model(input, output)
model.summary()

# 2. 定义优化器和日志记录器
optimizer = tf.keras.optimizers.Adam()
logdir = "./log"
global_steps = tf.Variable(0, trainable=False, dtype=tf.int64)
if os.path.exists(logdir): shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)

# 3. 训练网络no1，耗时18s
start = time.time()
for i in range(5):
    for j in range(int(train_x.shape[0]/200)):
        # 3.1 计算梯度
        with tf.GradientTape() as tape:
            y = model(train_x[j*200:(j+1)*200,:,:,:])
            loss = -tf.reduce_sum(train_y[j*200:(j+1)*200,:]*tf.math.log(y))
        gradients = tape.gradient(loss, model.trainable_variables)
        # 3.2 更新权重
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    #3.3 测试网络
    total_accuracy = []
    for k in range(int(test_x.shape[0]/200)):
        y = model(test_x[k*200:(k+1)*200,:,:,:])
        correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(test_y[k*200:(k+1)*200,:],1))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
        total_accuracy.append(accuracy.numpy())
    print("step: %d, accuracy: %f" % (i+1, np.mean(total_accuracy, 0)))
end = time.time()
print(end-start)

# 3. 训练网络no2，耗时5s
start = time.time()
@tf.function
def train(gt_x, gt_y):  # 定义单次训练函数
    # 3.1 计算梯度
    with tf.GradientTape() as tape:
        y = model(gt_x, training=True)  # model()函数需要指定training属性，该属性影响BN和dropout层，指定当前是训练模式还是测试模式，与trainable属性是两码事
        loss = -tf.reduce_sum(gt_y*tf.math.log(y))
    gradients = tape.gradient(loss, model.trainable_variables)
    # 3.2 设置学习率和更新权重
    optimizer.lr.assign(0.001)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    # 3.3 写一次日志
    global_steps.assign_add(1)
    with writer.as_default():
        tf.summary.scalar("lr", optimizer.lr, step=global_steps)  # global_steps必须这样写，如果改成整形数字n，不仅有警告，还运行贼慢！
        tf.summary.scalar("total_loss", loss, step=global_steps)
    writer.flush()

@tf.function
def test(gt_x, gt_y):  # 定义单次测试函数
    #3.4 测试网络
    y = model(gt_x, training=False)
    correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(gt_y,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy

for i in range(5):     # 训练5个epoch
    # 训练网络
    for j in range(int(train_x.shape[0]/200)):
        gt_x, gt_y = train_x[j*200:(j+1)*200,:,:,:], train_y[j*200:(j+1)*200,:]
        train(gt_x, gt_y)
    # 测试网络
    total_accuracy = []
    for k in range(int(test_x.shape[0]/200)):
        gt_x, gt_y = test_x[k*200:(k+1)*200,:,:,:], test_y[k*200:(k+1)*200,:]
        accuracy = test(gt_x, gt_y)
        total_accuracy.append(accuracy.numpy())
    # 打印结果
    print("step: %d, accuracy: %f" % (i+1, np.mean(total_accuracy, 0)))
end = time.time()
print(end-start)


# model.save_weights('./model_file.h5')
# model.save('./model_file')