Tensorflow2.0分批訓練模型

首先引入相關的庫

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

加載數據集

(mnist_images, mnist_labels), _ = tf.keras.datasets.mnist.load_data()
print(type(mnist_images))
print(mnist_images.dtype)
print(mnist_images.shape)

此時可以查看數據集的相關信息：

<class 'numpy.ndarray'>
uint8
(60000, 28, 28)

切分產生dataset

dataset = tf.data.Dataset.from_tensor_slices(
  (tf.cast(mnist_images[...,tf.newaxis]/255, tf.float32),
   tf.cast(mnist_labels,tf.int64)))
print(dataset)

其中：
mnist_images[…,tf.newaxis] 是指在mnist_images（維度爲(60000, 28, 28)）後面再加一個維度，即變成（60000, 28, 28, 1）。省略號表示的是很多個冒號，也就是說[…,tf.newaxis]即爲[:, :, :, tf.newaxis]。
tf.cast() 的作用是將mnist_images原來的數據類型（uint8）轉化爲float32。
tf.data.Dataset.from_tensor_slices() 用來將輸入按照第一維度進行切分，所以最終得到的dataset爲：

<TensorSliceDataset shapes: ((28, 28, 1), ()), types: (tf.float32, tf.int64)>

將dataset打亂、分批

dataset = dataset.shuffle(1000).batch(32)
print(dataset)

得到的dataset爲：

<BatchDataset shapes: ((None, 28, 28, 1), (None,)), types: (tf.float32, tf.int64)>

這個操作的目的是，每次調用dataset.take(1)就會得到32個樣本的組合，如果調用的時dataset.take(x)，則得到的是x個這樣的組合。

構建模型

mnist_model = tf.keras.Sequential([
  tf.keras.layers.Conv2D(16,[3,3], activation='relu', 
                         input_shape=(None, None, 1)),
  tf.keras.layers.Conv2D(16,[3,3], activation='relu'),
  tf.keras.layers.GlobalAveragePooling2D(),
  tf.keras.layers.Dense(10)
])

先展示一波dataset.take()的作用

for images,labels in dataset.take(2):
    print(images.shape)
    print("Logits: ", mnist_model(images[0:1]).numpy())

得到：

(32, 28, 28, 1)
Logits:  [[-0.01555313  0.02404156  0.01355665 -0.03020764  0.0092129   0.02039536
   0.01312329 -0.01198696  0.00587582  0.00114011]]
(32, 28, 28, 1)
Logits:  [[-0.03150982  0.04380588  0.02759154 -0.04930525  0.0152742   0.03806124
   0.02004753 -0.01281266  0.01789122  0.00325774]]

由此可見，每次得到的要輸入模型的樣本數都是32個。Logits是取這32個樣本中的第一個輸入模型時得到的結果。

優化器+損失函數

optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

訓練過程

loss_history = []
for (batch, (images, labels)) in enumerate(dataset.take(400)):
    if batch % 10 == 0:
        print('.', end='')
    with tf.GradientTape() as tape:
        logits = mnist_model(images, training=True)
        loss_value = loss_object(labels, logits)
        
    loss_history.append(loss_value.numpy().mean())
    grads = tape.gradient(loss_value, mnist_model.trainable_variables)
    optimizer.apply_gradients(zip(grads, mnist_model.trainable_variables))

以上涉及梯度帶的使用，請參考：Tensorflow中的梯度帶（GradientTape）以及梯度更新。

畫出損失圖

plt.plot(loss_history)
plt.xlabel('Batch #')
plt.ylabel('Loss [entropy]')
plt.show()