导入所需的模块
import tensorflow as tf
import numpy as np
from tensorflow.keras.optimizers import SGD
加载数据
(train_images, train_labels), (valid_images, valid_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(-1, 28, 28, 1)
valid_images = valid_images.reshape(-1, 28, 28, 1)
train_images, valid_images = train_images / 255.0, valid_images / 255.0
print(train_images.shape)
print(train_labels.shape)
print(train_images.dtype)
print(train_labels.dtype)
构建模型
model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(6, (5,5), activation='tanh', input_shape=(28,28,1)),
tf.keras.layers.MaxPooling2D(2,2),
tf.keras.layers.Conv2D(16, (5,5), activation='tanh'),
tf.keras.layers.MaxPooling2D(2,2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(120, activation='tanh'),
tf.keras.layers.Dense(84, activation='tanh'),
tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(loss='sparse_categorical_crossentropy',
optimizer=SGD(lr=0.05, decay=1e-6, momentum=0.9, nesterov=True),
metrics=['acc'])
训练模型
history = model.fit(train_images, train_labels, batch_size=32, epochs=2,\
verbose=1, shuffle=True, validation_data=(valid_images, valid_labels))
评估模型
model.evaluate(valid_images, valid_labels, verbose=2)
保存模型
model.save("./test-model.h5")
new_model = tf.keras.models.load_model("./test-model.h5")
迁移学习
加载预训练权重
base_model = tf.keras.models.load_model("./model_h5/test-model.h5")
flatten = base_model.get_layer("flatten_1").output
dense = tf.keras.layers.Dense(256, activation='relu')(flatten)
dense = tf.keras.layers.Dense(128, activation='relu')(dense)
pred = tf.keras.layers.Dense(10, activation='softmax')(dense)
fine_tune_model = tf.keras.models.Model(inputs=base_model.input, outputs=pred)
fine_tune_model.summary()
Freeze
for i, layer in enumerate(fine_tune_model.layers):
print(i, layer.name)
for layer in fine_tune_model.layers[:6]:
layer.trainable = False
Fine-tune
fine_tune_model.compile(loss='sparse_categorical_crossentropy', optimizer=SGD(lr=0.05, decay=1e-6, momentum=0.9, nesterov=True), metrics=['acc'])
fine_tune_model.fit(train_images, train_labels, batch_size=32, epochs=1, verbose=1, shuffle=True, validation_data=(valid_images, valid_labels))
并行训练
strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
with strategy.scope():
model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(6, (5,5), activation='tanh', input_shape=(28,28,1)),
tf.keras.layers.MaxPooling2D(2,2),
tf.keras.layers.Conv2D(16, (5,5), activation='tanh'),
tf.keras.layers.MaxPooling2D(2,2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(120, activation='tanh'),
tf.keras.layers.Dense(84, activation='tanh'),
tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(loss='sparse_categorical_crossentropy', optimizer=SGD(lr=0.05, decay=1e-6, momentum=0.9, nesterov=True), metrics=['acc'])
model.fit(train_images, train_labels, batch_size=32, epochs=1, verbose=1, shuffle=True, validation_data=(valid_images, valid_labels))
参考
- Keras Guide
- keras训练