import tensorflow as tf import os import numpy as np import random

定义训练、验证、测试集的路径

train_dir = 'data/train/' val_dir = 'data/val/' test_dir = 'data/test/'

定义每个类别的标签

label_dict = {'bulbasaur': 0, 'charmander': 1, 'squirtle': 2, 'pikachu': 3, 'mewtwo': 4}

定义超参数

batch_size = 32 epochs = 50 learning_rate = 0.001

批量读取图像路径，生成每张图像的路径和标签并打乱顺序

def read_images_path(dir_path): images_path = [] labels = [] for root, dirs, files in os.walk(dir_path): for file in files: if file.endswith('.jpg'): images_path.append(os.path.join(root, file)) labels.append(label_dict[root.split('/')[-1]]) images_labels = list(zip(images_path, labels)) random.shuffle(images_labels) images_path, labels = zip(*images_labels) return np.array(images_path), np.array(labels)

构建卷积神经网络

def build_model(): model = tf.keras.Sequential([ tf.keras.layers.Conv2D(16, (3, 3), padding='same', activation='relu', input_shape=(224, 224, 3)), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu'), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu'), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Flatten(), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dropout(0.2), tf.keras.layers.Dense(5, activation='softmax') ]) return model

定义准确率和损失的tensorboard可视化

def tensorboard_callback(): log_dir = 'logs/' tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1) return tensorboard_callback

加载数据集，并将其划分为训练集、验证集、测试集

train_images_path, train_labels = read_images_path(train_dir) val_images_path, val_labels = read_images_path(val_dir) test_images_path, test_labels = read_images_path(test_dir)

train_images = tf.data.Dataset.from_tensor_slices((train_images_path, train_labels)) val_images = tf.data.Dataset.from_tensor_slices((val_images_path, val_labels)) test_images = tf.data.Dataset.from_tensor_slices((test_images_path, test_labels))

对数据集进行预处理

def preprocess_images(images_path, label): image = tf.io.read_file(images_path) image = tf.image.decode_jpeg(image, channels=3) image = tf.image.resize(image, (224, 224)) image = tf.cast(image, tf.float32) image = image / 255.0 return image, label

train_images = train_images.map(preprocess_images).batch(batch_size).shuffle(buffer_size=1000) val_images = val_images.map(preprocess_images).batch(batch_size).shuffle(buffer_size=1000) test_images = test_images.map(preprocess_images).batch(batch_size).shuffle(buffer_size=1000)

创建模型

model = build_model()

编译模型

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

训练模型

history = model.fit(train_images, epochs=epochs, validation_data=val_images, callbacks=[tensorboard_callback()])

评估模型

test_loss, test_acc = model.evaluate(test_images, verbose=2) print('\nTest accuracy:', test_acc)

保存模型

model.save('pokemon_classification_model'