如何用TensorFlow加载自己的二分类图像数据集 并搭建CNN模型以及训练模型代码

1. 加载数据集

首先，我们需要将数据集准备好并加载到 TensorFlow 中。在这里，我们假设数据集已经被分成了训练集和测试集，并分别存储在两个文件夹下。

import tensorflow as tf

train_dir = '/path/to/train_set/'
test_dir = '/path/to/test_set/'

train_data = tf.keras.preprocessing.image_dataset_from_directory(
    train_dir,
    batch_size=32,
    image_size=(224, 224),
    seed=42,
    validation_split=0.2,
    subset='training'
)

val_data = tf.keras.preprocessing.image_dataset_from_directory(
    train_dir,
    batch_size=32,
    image_size=(224, 224),
    seed=42,
    validation_split=0.2,
    subset='validation'
)

test_data = tf.keras.preprocessing.image_dataset_from_directory(
    test_dir,
    batch_size=32,
    image_size=(224, 224)
)

2. 搭建模型

在搭建模型之前，我们需要定义一些超参数，例如卷积核的数量、大小以及池化操作的大小等。在这里，我们使用了 3 个卷积层和 2 个全连接层。

num_classes = 2

model = tf.keras.Sequential([
    tf.keras.layers.experimental.preprocessing.Rescaling(1./255, input_shape=(224, 224, 3)),
    tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Conv2D(128, kernel_size=(3, 3), activation='relu'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])

3. 训练模型

在模型搭建完成之后，我们需要对其进行编译，并使用训练集进行训练。

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

history = model.fit(
    train_data,
    validation_data=val_data,
    epochs=10
)

4. 评估模型

训练完成后，我们可以使用测试集对模型进行评估。

test_loss, test_acc = model.evaluate(test_data)

print('Test Loss:', test_loss)
print('Test Accuracy:', test_acc)