|
|
import os
|
|
|
import tensorflow as tf
|
|
|
from tensorflow.keras.models import Sequential
|
|
|
from tensorflow.keras.layers import Dense, Dropout, Conv2D, MaxPooling2D, Flatten
|
|
|
from tensorflow.keras.optimizers import Adam
|
|
|
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
|
|
|
|
|
|
import tensorflow as tf
|
|
|
print(tf.__version__)
|
|
|
|
|
|
# 导入数据
|
|
|
mnist = tf.keras.datasets.mnist
|
|
|
(train_data, train_target), (test_data, test_target) = mnist.load_data()
|
|
|
|
|
|
# 改变数据维度
|
|
|
# 改变数据维度
|
|
|
train_data = train_data.reshape(-1, 28, 28, 1)
|
|
|
test_data = test_data.reshape(-1, 28, 28, 1)
|
|
|
# 注：在TensorFlow中，在做卷积的时候需要把数据变成4维的格式
|
|
|
# 这4个维度分别是：数据数量，图片高度，图片宽度，图片通道数
|
|
|
|
|
|
# 归一化（有助于提升训练速度）
|
|
|
train_data = train_data / 255.0
|
|
|
test_data = test_data / 255.0
|
|
|
|
|
|
# 独热编码
|
|
|
train_target = tf.keras.utils.to_categorical(train_target, num_classes=10)
|
|
|
test_target = tf.keras.utils.to_categorical(test_target, num_classes=10) # 10种结果
|
|
|
|
|
|
# 配置早停
|
|
|
early_stopping = EarlyStopping(monitor='val_loss', patience=3, verbose=1, restore_best_weights=True)
|
|
|
|
|
|
# 配置GPU加速
|
|
|
strategy = tf.distribute.MirroredStrategy()
|
|
|
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
|
|
|
|
|
|
with strategy.scope():
|
|
|
# 构建更复杂的模型
|
|
|
model = tf.keras.Sequential([
|
|
|
Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1), padding = 'same'),
|
|
|
MaxPooling2D((2, 2), padding = 'same'),
|
|
|
Conv2D(64, (3, 3), activation='relu', padding = 'same'),
|
|
|
MaxPooling2D((2, 2), padding = 'same'),
|
|
|
Flatten(),
|
|
|
Dense(1024, activation='relu'),
|
|
|
Dropout(0.5),
|
|
|
Dense(10, activation='softmax')
|
|
|
])
|
|
|
|
|
|
# 编译模型
|
|
|
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
|
|
|
model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy'])
|
|
|
|
|
|
# # 配置模型检查点，保存最优模型
|
|
|
# checkpoint_path = "./model/number_model.h5"
|
|
|
# model_checkpoint = ModelCheckpoint(checkpoint_path, monitor='val_loss',
|
|
|
# save_best_only=True, save_weights_only=False, verbose=1)
|
|
|
|
|
|
# 训练模型
|
|
|
model.fit(train_data, train_target, epochs=5, validation_data=(test_data, test_target), callbacks=[early_stopping])
|
|
|
|
|
|
# 保存模型为 .h5 文件
|
|
|
model.save("./model/number_model.h5", save_format='h5') |
