三分鍾快速上手TensorFlow 2.0 （中）——常用模塊和模型的部署

本文學習筆記參照來源：https://tf.wiki/zh/basic/basic.html

前文：三分鍾快速上手TensorFlow 2.0 （上）——前置基礎、模型建立與可視化

 

tf.train.Checkpoint ：變量的保存與恢復

只保存模型的參數，不保存模型的計算過程

需要導出模型（無需源代碼也能運行模型），請參考 SavedModel

可以使用其 save() 和 restore() 方法將 TensorFlow 中所有包含 Checkpointable State 的對象進行保存和恢復。

tf.keras.optimizer 、 tf.Variable 、 tf.keras.Layer 或者 tf.keras.Model 實例都可以被保存。

checkpoint = tf.train.Checkpoint(model=model)

舉例，

 

checkpoint.save(save_path_with_prefix)

save()　　　　輸入的是保存文件的目錄 + 前綴。比如'./save/model.ckpt'

model_to_be_restored = MyModel()                                        # 待恢復參數的同一模型
checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)   # 鍵名保持為“myAwesomeModel”
checkpoint.restore(save_path_with_prefix_and_index)

restore（）

save_path_with_prefix_and_index 是之前保存的文件的目錄 + 前綴 + 編號。例如，調用 checkpoint.restore('./save/model.ckpt-1') 就可以載入前綴為 model.ckpt ，序號為 1 的文件來恢復模型。

載入最近的一個。可以使用 tf.train.latest_checkpoint(save_path) 這個輔助函數返回目錄下最近一次 checkpoint 的文件名。

 

總框架

# train.py 模型訓練階段

model = MyModel()
# 實例化Checkpoint，指定保存對象為model（如果需要保存Optimizer的參數也可加入）
checkpoint = tf.train.Checkpoint(myModel=model)
# ...（模型訓練代碼）
# 模型訓練完畢后將參數保存到文件（也可以在模型訓練過程中每隔一段時間就保存一次）
checkpoint.save('./save/model.ckpt')

# test.py 模型使用階段

model = MyModel()
checkpoint = tf.train.Checkpoint(myModel=model)             # 實例化Checkpoint，指定恢復對象為model
checkpoint.restore(tf.train.latest_checkpoint('./save'))    # 從文件恢復模型參數
# 模型使用代碼

 

import tensorflow as tf
import numpy as np
import argparse
from zh.model.mnist.mlp import MLP
from zh.model.utils import MNISTLoader

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('--mode', default='train', help='train or test')
parser.add_argument('--num_epochs', default=1)
parser.add_argument('--batch_size', default=50)
parser.add_argument('--learning_rate', default=0.001)
args = parser.parse_args()
data_loader = MNISTLoader()


def train():
    model = MLP()
    optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
    num_batches = int(data_loader.num_train_data // args.batch_size * args.num_epochs)
    checkpoint = tf.train.Checkpoint(myAwesomeModel=model)      # 實例化Checkpoint，設置保存對象為model
    for batch_index in range(1, num_batches+1):                 
        X, y = data_loader.get_batch(args.batch_size)
        with tf.GradientTape() as tape:
            y_pred = model(X)
            loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
            loss = tf.reduce_mean(loss)
            print("batch %d: loss %f" % (batch_index, loss.numpy()))
        grads = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
        if batch_index % 100 == 0:                              # 每隔100個Batch保存一次
            path = checkpoint.save('./save/model.ckpt')         # 保存模型參數到文件
            print("model saved to %s" % path)


def test():
    model_to_be_restored = MLP()
    # 實例化Checkpoint，設置恢復對象為新建立的模型model_to_be_restored
    checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)      
    checkpoint.restore(tf.train.latest_checkpoint('./save'))    # 從文件恢復模型參數
    y_pred = np.argmax(model_to_be_restored.predict(data_loader.test_data), axis=-1)
    print("test accuracy: %f" % (sum(y_pred == data_loader.test_label) / data_loader.num_test_data))


if __name__ == '__main__':
    if args.mode == 'train':
        train()
    if args.mode == 'test':
        test()

 

 

• 保留最后的幾個 Checkpoint；

• Checkpoint 默認從 1 開始編號，每次累加 1，但我們可能希望使用別的編號方式（例如使用當前 Batch 的編號作為文件編號）。

checkpoint = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(checkpoint, directory='./save', checkpoint_name='model.ckpt', max_to_keep=k)

directory 參數為文件保存的路徑， checkpoint_name 為文件名前綴（不提供則默認為 ckpt ）， max_to_keep 為保留的 Checkpoint 數目。

在需要保存模型的時候，我們直接使用 manager.save() 即可。如果我們希望自行指定保存的 Checkpoint 的編號，則可以在保存時加入 checkpoint_number 參數。例如 manager.save(checkpoint_number=100) 。

import tensorflow as tf
import numpy as np
import argparse
from zh.model.mnist.mlp import MLP
from zh.model.utils import MNISTLoader

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('--mode', default='train', help='train or test')
parser.add_argument('--num_epochs', default=1)
parser.add_argument('--batch_size', default=50)
parser.add_argument('--learning_rate', default=0.001)
args = parser.parse_args()
data_loader = MNISTLoader()


def train():
    model = MLP()
    optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
    num_batches = int(data_loader.num_train_data // args.batch_size * args.num_epochs)
    checkpoint = tf.train.Checkpoint(myAwesomeModel=model)      
    # 使用tf.train.CheckpointManager管理Checkpoint
    manager = tf.train.CheckpointManager(checkpoint, directory='./save', max_to_keep=3)
    for batch_index in range(1, num_batches):
        X, y = data_loader.get_batch(args.batch_size)
        with tf.GradientTape() as tape:
            y_pred = model(X)
            loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
            loss = tf.reduce_mean(loss)
            print("batch %d: loss %f" % (batch_index, loss.numpy()))
        grads = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
        if batch_index % 100 == 0:
            # 使用CheckpointManager保存模型參數到文件並自定義編號
            path = manager.save(checkpoint_number=batch_index)         
            print("model saved to %s" % path)


def test():
    model_to_be_restored = MLP()
    checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)      
    checkpoint.restore(tf.train.latest_checkpoint('./save'))
    y_pred = np.argmax(model_to_be_restored.predict(data_loader.test_data), axis=-1)
    print("test accuracy: %f" % (sum(y_pred == data_loader.test_label) / data_loader.num_test_data))


if __name__ == '__main__':
    if args.mode == 'train':
        train()
    if args.mode == 'test':
        test()

 

 

 

 

TensorBoard：訓練過程可視化

summary_writer = tf.summary.create_file_writer('./tensorboard')     # 參數為記錄文件所保存的目錄

summary_writer = tf.summary.create_file_writer('./tensorboard')
# 開始模型訓練
for batch_index in range(num_batches):
    # ...（訓練代碼，當前batch的損失值放入變量loss中）
    with summary_writer.as_default():                               # 希望使用的記錄器
        tf.summary.scalar("loss", loss, step=batch_index)
        tf.summary.scalar("MyScalar", my_scalar, step=batch_index)  # 還可以添加其他自定義的變量

tensorboard --logdir=./tensorboard

tf.summary.trace_on(graph=True, profiler=True)  # 開啟Trace，可以記錄圖結構和profile信息
# 進行訓練
with summary_writer.as_default():
    tf.summary.trace_export(name="model_trace", step=0, profiler_outdir=log_dir)    # 保存Trace信息到文件

• 如果需要重新訓練，需要刪除掉記錄文件夾內的信息並重啟 TensorBoard（或者建立一個新的記錄文件夾並開啟 TensorBoard， --logdir 參數設置為新建立的文件夾）；

• 記錄文件夾目錄保持全英文。

import tensorflow as tf
from zh.model.mnist.mlp import MLP
from zh.model.utils import MNISTLoader

num_batches = 1000
batch_size = 50
learning_rate = 0.001
log_dir = 'tensorboard'
model = MLP()
data_loader = MNISTLoader()
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
summary_writer = tf.summary.create_file_writer(log_dir)     # 實例化記錄器
tf.summary.trace_on(profiler=True)  # 開啟Trace（可選）
for batch_index in range(num_batches):
    X, y = data_loader.get_batch(batch_size)
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        print("batch %d: loss %f" % (batch_index, loss.numpy()))
        with summary_writer.as_default():                           # 指定記錄器
            tf.summary.scalar("loss", loss, step=batch_index)       # 將當前損失函數的值寫入記錄器
    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
with summary_writer.as_default():
    tf.summary.trace_export(name="model_trace", step=0, profiler_outdir=log_dir)    # 保存Trace信息到文件（可選）

 

 

 

 

 

 

tf.data ：數據集的構建與預處理

快速、高效地構建數據輸入的流水線，尤其適用於數據量巨大的場景。

數據集對象的建立

tf.data 的核心是 tf.data.Dataset 類，提供了對數據集的高層封裝。tf.data.Dataset由一系列的可迭代訪問的元素（element）組成

import tensorflow as tf
import numpy as np

X = tf.constant([2013, 2014, 2015, 2016, 2017])
Y = tf.constant([12000, 14000, 15000, 16500, 17500])

# 也可以使用NumPy數組，效果相同
# X = np.array([2013, 2014, 2015, 2016, 2017])
# Y = np.array([12000, 14000, 15000, 16500, 17500])

dataset = tf.data.Dataset.from_tensor_slices((X, Y))

for x, y in dataset:
    print(x.numpy(), y.numpy()) 

例：

數據集對象的預處理

最常用的如：

• Dataset.map(f) ：對數據集中的每個元素應用函數 f ，得到一個新的數據集（這部分往往結合 tf.io 進行讀寫和解碼文件， tf.image 進行圖像處理）；

• Dataset.shuffle(buffer_size) ：將數據集打亂（設定一個固定大小的緩沖區（Buffer），取出前 buffer_size 個元素放入，並從緩沖區中隨機采樣，采樣后的數據用后續數據替換）；

• Dataset.batch(batch_size) ：將數據集分成批次，即對每 batch_size 個元素，使用 tf.stack() 在第 0 維合並，成為一個元素。

還有 Dataset.repeat() （重復數據集的元素）、 Dataset.reduce() （與 Map 相對的聚合操作）、 Dataset.take() （截取數據集中的前若干個元素）等

def rot90(image, label):
    image = tf.image.rot90(image)
    return image, label

mnist_dataset = mnist_dataset.map(rot90)

for image, label in mnist_dataset:
    plt.title(label.numpy())
    plt.imshow(image.numpy()[:, :, 0])
    plt.show(

mnist_dataset = mnist_dataset.batch(4)

for images, labels in mnist_dataset:    # image: [4, 28, 28, 1], labels: [4]
    fig, axs = plt.subplots(1, 4)
    for i in range(4):
        axs[i].set_title(labels.numpy()[i])
        axs[i].imshow(images.numpy()[i, :, :, 0])
    plt.show()

mnist_dataset = mnist_dataset.shuffle(buffer_size=10000).batch(4)

for images, labels in mnist_dataset:
    fig, axs = plt.subplots(1, 4)
    for i in range(4):
        axs[i].set_title(labels.numpy()[i])
        axs[i].imshow(images.numpy()[i, :, :, 0])
    plt.show()

• 設定一個固定大小為 buffer_size 的緩沖區（Buffer）；

• 初始化時，取出數據集中的前 buffer_size 個元素放入緩沖區；

• 每次需要從數據集中取元素時，即從緩沖區中隨機采樣一個元素並取出，然后從后續的元素中取出一個放回到之前被取出的位置，以維持緩沖區的大小。

因此，緩沖區的大小需要根據數據集的特性和數據排列順序特點來進行合理的設置。比如：

• 當 buffer_size 設置為 1 時，其實等價於沒有進行任何打散；

• 當數據集的標簽順序分布極為不均勻（例如二元分類時數據集前 N 個的標簽為 0，后 N 個的標簽為 1）時，較小的緩沖區大小會使得訓練時取出的 Batch 數據很可能全為同一標簽，從而影響訓練效果。一般而言，數據集的順序分布若較為隨機，則緩沖區的大小可較小，否則則需要設置較大的緩沖區。

使用 tf.data 的並行化策略提高訓練流程效率

mnist_dataset=mnist_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

mnist_dataset = mnist_dataset.map(map_func=rot90, num_parallel_calls=2)

數據集元素的獲取與使用

dataset = tf.data.Dataset.from_tensor_slices((A, B, C, ...))
for a, b, c, ... in dataset:
    # 對張量a, b, c等進行操作，例如送入模型進行訓練

dataset = tf.data.Dataset.from_tensor_slices((A, B, C, ...))
it = iter(dataset)
a_0, b_0, c_0, ... = next(it)
a_1, b_1, c_1, ... = next(it)

Keras 支持使用 tf.data.Dataset 直接作為輸入。當調用 tf.keras.Model 的 fit() 和 evaluate() 方法時，可以將參數中的輸入數據 x 指定為一個元素格式為 (輸入數據, 標簽數據) 的 Dataset ，並忽略掉參數中的標簽數據 y 。例如，對於上述的 MNIST 數據集

model.fit(x=train_data, y=train_label, epochs=num_epochs, batch_size=batch_size)

model.fit(mnist_dataset, epochs=num_epochs)

 

import tensorflow as tf
import os

num_epochs = 10
batch_size = 32
learning_rate = 0.001
data_dir = 'C:/datasets/cats_vs_dogs'
train_cats_dir = data_dir + '/train/cats/'
train_dogs_dir = data_dir + '/train/dogs/'
test_cats_dir = data_dir + '/valid/cats/'
test_dogs_dir = data_dir + '/valid/dogs/'

def _decode_and_resize(filename, label):
    image_string = tf.io.read_file(filename)            # 讀取原始文件
    image_decoded = tf.image.decode_jpeg(image_string)  # 解碼JPEG圖片
    image_resized = tf.image.resize(image_decoded, [256, 256]) / 255.0
    return image_resized, label

if __name__ == '__main__':
    # 構建訓練數據集
    train_cat_filenames = tf.constant([train_cats_dir + filename for filename in os.listdir(train_cats_dir)])
    train_dog_filenames = tf.constant([train_dogs_dir + filename for filename in os.listdir(train_dogs_dir)])
    train_filenames = tf.concat([train_cat_filenames, train_dog_filenames], axis=-1)
    train_labels = tf.concat([
        tf.zeros(train_cat_filenames.shape, dtype=tf.int32), 
        tf.ones(train_dog_filenames.shape, dtype=tf.int32)], 
        axis=-1)

    train_dataset = tf.data.Dataset.from_tensor_slices((train_filenames, train_labels))
    train_dataset = train_dataset.map(
        map_func=_decode_and_resize, 
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    # 取出前buffer_size個數據放入buffer，並從其中隨機采樣，采樣后的數據用后續數據替換
    train_dataset = train_dataset.shuffle(buffer_size=23000)    
    train_dataset = train_dataset.batch(batch_size)
    train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(256, 256, 3)),
        tf.keras.layers.MaxPooling2D(),
        tf.keras.layers.Conv2D(32, 5, activation='relu'),
        tf.keras.layers.MaxPooling2D(),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(2, activation='softmax')
    ])

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=[tf.keras.metrics.sparse_categorical_accuracy]
    )

    model.fit(train_dataset, epochs=num_epochs)

 # 構建測試數據集
    test_cat_filenames = tf.constant([test_cats_dir + filename for filename in os.listdir(test_cats_dir)])
    test_dog_filenames = tf.constant([test_dogs_dir + filename for filename in os.listdir(test_dogs_dir)])
    test_filenames = tf.concat([test_cat_filenames, test_dog_filenames], axis=-1)
    test_labels = tf.concat([
        tf.zeros(test_cat_filenames.shape, dtype=tf.int32), 
        tf.ones(test_dog_filenames.shape, dtype=tf.int32)], 
        axis=-1)

    test_dataset = tf.data.Dataset.from_tensor_slices((test_filenames, test_labels))
    test_dataset = test_dataset.map(_decode_and_resize)
    test_dataset = test_dataset.batch(batch_size)

    print(model.metrics_names)
    print(model.evaluate(test_dataset))

 

 

TFRecord ：TensorFlow 數據集存儲格式

TensorFlow 中的數據集存儲格式。當我們將數據集整理成 TFRecord 格式后，TensorFlow 就可以高效地讀取和處理這些數據集

TFRecord 可以理解為一系列序列化的 tf.train.Example 元素所組成的列表文件，而每一個 tf.train.Example 又由若干個 tf.train.Feature 的字典組成。形式如下：

# dataset.tfrecords
[ { # example 1 (tf.train.Example) 'feature_1': tf.train.Feature, ... 'feature_k': tf.train.Feature }, ... { # example N (tf.train.Example) 'feature_1': tf.train.Feature, ... 'feature_k': tf.train.Feature } ] 

為了將形式各樣的數據集整理為 TFRecord 格式，我們可以對數據集中的每個元素進行以下步驟：

• 讀取該數據元素到內存；

• 將該元素轉換為 tf.train.Example 對象（每一個 tf.train.Example 由若干個 tf.train.Feature 的字典組成，因此需要先建立 Feature 的字典）；

• 將該 tf.train.Example 對象序列化為字符串，並通過一個預先定義的 tf.io.TFRecordWriter 寫入 TFRecord 文件。

而讀取 TFRecord 數據則可按照以下步驟：

• 通過 tf.data.TFRecordDataset 讀入原始的 TFRecord 文件（此時文件中的 tf.train.Example 對象尚未被反序列化），獲得一個 tf.data.Dataset 數據集對象；

• 通過 Dataset.map 方法，對該數據集對象中的每一個序列化的 tf.train.Example 字符串執行 tf.io.parse_single_example 函數，從而實現反序列化。

 

將 cats_vs_dogs 二分類數據集的訓練集部分轉換為 TFRecord 文件，並讀取該文件 

存儲為 TFRecord 文件

import tensorflow as tf
import os

data_dir = 'C:/datasets/cats_vs_dogs'
train_cats_dir = data_dir + '/train/cats/'
train_dogs_dir = data_dir + '/train/dogs/'
tfrecord_file = data_dir + '/train/train.tfrecords'

train_cat_filenames = [train_cats_dir + filename for filename in os.listdir(train_cats_dir)]
train_dog_filenames = [train_dogs_dir + filename for filename in os.listdir(train_dogs_dir)]
train_filenames = train_cat_filenames + train_dog_filenames
train_labels = [0] * len(train_cat_filenames) + [1] * len(train_dog_filenames)  # 將 cat 類的標簽設為0，dog 類的標簽設為1

with tf.io.TFRecordWriter(tfrecord_file) as writer:
    for filename, label in zip(train_filenames, train_labels):
        image = open(filename, 'rb').read()     # 讀取數據集圖片到內存，image 為一個 Byte 類型的字符串
        feature = {                             # 建立 tf.train.Feature 字典
            'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),  # 圖片是一個 Bytes 對象
            'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))   # 標簽是一個 Int 對象
        }
        example = tf.train.Example(features=tf.train.Features(feature=feature)) # 通過字典建立 Example
        writer.write(example.SerializeToString())   # 將Example序列化並寫入 TFRecord 文件

tf.train.Feature 支持三種數據格式：

• tf.train.BytesList ：字符串或原始 Byte 文件（如圖片），通過 bytes_list 參數傳入一個由字符串數組初始化的 tf.train.BytesList 對象；

• tf.train.FloatList ：浮點數，通過 float_list 參數傳入一個由浮點數數組初始化的 tf.train.FloatList 對象；

• tf.train.Int64List ：整數，通過 int64_list 參數傳入一個由整數數組初始化的 tf.train.Int64List 對象。

如果只希望保存一個元素而非數組，傳入一個只有一個元素的數組即可

讀取 TFRecord 文件

raw_dataset = tf.data.TFRecordDataset(tfrecord_file)    # 讀取 TFRecord 文件

feature_description = { # 定義Feature結構，告訴解碼器每個Feature的類型是什么
    'image': tf.io.FixedLenFeature([], tf.string),
    'label': tf.io.FixedLenFeature([], tf.int64),
}

def _parse_example(example_string): # 將 TFRecord 文件中的每一個序列化的 tf.train.Example 解碼
    feature_dict = tf.io.parse_single_example(example_string, feature_description)
    feature_dict['image'] = tf.io.decode_jpeg(feature_dict['image'])    # 解碼JPEG圖片
    return feature_dict['image'], feature_dict['label']

dataset = raw_dataset.map(_parse_example)

運行以上代碼后，我們獲得一個數據集對象  dataset ，這已經是一個可以用於訓練的  tf.data.Dataset 對象

import matplotlib.pyplot as plt 

for image, label in dataset:
    plt.title('cat' if label == 0 else 'dog')
    plt.imshow(image.numpy())
    plt.show()

 

 

@tf.function ：圖執行模式 *

追求高性能或部署模型，將模型轉換為高效的 TensorFlow 圖模型

import tensorflow as tf
import time
from zh.model.mnist.cnn import CNN
from zh.model.utils import MNISTLoader

num_batches = 400
batch_size = 50
learning_rate = 0.001
data_loader = MNISTLoader()

@tf.function
def train_one_step(X, y):    
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        # 注意這里使用了TensorFlow內置的tf.print()。@tf.function不支持Python內置的print方法
        tf.print("loss", loss)  
    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))

if __name__ == '__main__':
    model = CNN()
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
    start_time = time.time()
    for batch_index in range(num_batches):
        X, y = data_loader.get_batch(batch_size)
        train_one_step(X, y)
    end_time = time.time()
    print(end_time - start_time)      

當被 @tf.function 修飾的函數第一次被調用的時候，進行以下操作：

在即時執行模式關閉的環境下，函數內的代碼依次運行。也就是說，每個 tf. 方法都只是定義了計算節點，而並沒有進行任何實質的計算。這與 TensorFlow 1.X 的圖執行模式是一致的；
使用 AutoGraph 將函數中的 Python 控制流語句轉換成 TensorFlow 計算圖中的對應節點（比如說 while 和 for 語句轉換為 tf.while ， if 語句轉換為 tf.cond 等等；
基於上面的兩步，建立函數內代碼的計算圖表示（為了保證圖的計算順序，圖中還會自動加入一些 tf.control_dependencies 節點）；
運行一次這個計算圖；
基於函數的名字和輸入的函數參數的類型生成一個哈希值，並將建立的計算圖緩存到一個哈希表中。
在被 @tf.function 修飾的函數之后再次被調用的時候，根據函數名和輸入的函數參數的類型計算哈希值，檢查哈希表中是否已經有了對應計算圖的緩存。如果是，則直接使用已緩存的計算圖，否則重新按上述步驟建立計算圖。

import tensorflow as tf

@tf.function
def square_if_positive(x):
    if x > 0:
        x = x * x
    else:
        x = 0
    return x

a = tf.constant(1)
b = tf.constant(-1)
print(square_if_positive(a), square_if_positive(b))
print(tf.autograph.to_code(square_if_positive.python_function))

 optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)
    num_batches = int(data_loader.num_train_data // batch_size * num_epochs)
    # 建立計算圖
    X_placeholder = tf.compat.v1.placeholder(name='X', shape=[None, 28, 28, 1], dtype=tf.float32)
    y_placeholder = tf.compat.v1.placeholder(name='y', shape=[None], dtype=tf.int32)
    y_pred = model(X_placeholder)
    loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y_placeholder, y_pred=y_pred)
    loss = tf.reduce_mean(loss)
    train_op = optimizer.minimize(loss)
    sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
    # 建立Session
    with tf.compat.v1.Session() as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        for batch_index in range(num_batches):
            X, y = data_loader.get_batch(batch_size)
            # 使用Session.run()將數據送入計算圖節點，進行訓練以及計算損失函數
            _, loss_value = sess.run([train_op, loss], feed_dict={X_placeholder: X, y_placeholder: y})
            print("batch %d: loss %f" % (batch_index, loss_value))

        num_batches = int(data_loader.num_test_data // batch_size)
        for batch_index in range(num_batches):
            start_index, end_index = batch_index * batch_size, (batch_index + 1) * batch_size
            y_pred = model.predict(data_loader.test_data[start_index: end_index])
            sess.run(sparse_categorical_accuracy.update(y_true=data_loader.test_label[start_index: end_index], y_pred=y_pred))
        print("test accuracy: %f" % sess.run(sparse_categorical_accuracy.result()))

tf.TensorArray ：TensorFlow 動態數組 *

在即時執行模式下，你可以直接使用一個 Python 列表（List）存放數組。

不過，如果你需要基於計算圖的特性（例如使用 @tf.function 加速模型運行或者使用 SavedModel 導出模型），就無法使用這種方式了。

因此，TensorFlow 提供了 tf.TensorArray ，一種支持計算圖特性的 TensorFlow 動態數組。

其聲明的方式為：

• arr = tf.TensorArray(dtype, size, dynamic_size=False) ：聲明一個大小為 size ，類型為 dtype 的 TensorArray arr 。如果將 dynamic_size 參數設置為 True ，則該數組會自動增長空間。

其讀取和寫入的方法為：

• write(index, value) ：將 value 寫入數組的第 index 個位置；

• read(index) ：讀取數組的第 index 個值；

除此以外，TensorArray 還包括 stack() 、 unstack() 等常用操作，可參考 文檔 以了解詳情。

請注意，由於需要支持計算圖， tf.TensorArray 的 write() 方法是不可以忽略左值的！也就是說，在圖執行模式下，必須按照以下的形式寫入數組：

arr = arr.write(index, value) 

這樣才可以正常生成一個計算圖操作，並將該操作返回給 arr 。而不可以寫成：

arr.write(index, value) # 生成的計算圖操作沒有左值接收，從而丟失

import tensorflow as tf

@tf.function
def array_write_and_read():
    arr = tf.TensorArray(dtype=tf.float32, size=3)
    arr = arr.write(0, tf.constant(0.0))
    arr = arr.write(1, tf.constant(1.0))
    arr = arr.write(2, tf.constant(2.0))
    arr_0 = arr.read(0)
    arr_1 = arr.read(1)
    arr_2 = arr.read(2)
    return arr_0, arr_1, arr_2

a, b, c = array_write_and_read()
print(a, b, c)

tf.config：GPU 的使用與分配 *

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_visible_devices(devices=gpus[0:2], device_type='GPU')

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(device=gpu, enable=True)

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
     tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])

 

下一篇：三分鍾快速上手TensorFlow 2.0 （下）——模型的部署 、大規模訓練、加速