TensorFlow 训练自己的图像分类器

使用 TensorFlow 训练图像分类器，可以快速搭建属于自己的模型来识别不同类别的图片。下面是一个简明的流程和示例。

一、准备数据集

  • 收集并整理训练数据，通常每个类别放在一个独立文件夹内。
  • 例如：
1
2
3
4
5
6
7
dataset/
├── cats/
│   ├── cat1.jpg
│   ├── cat2.jpg
├── dogs/
│   ├── dog1.jpg
│   ├── dog2.jpg

二、环境准备

安装 TensorFlow：

1
pip install tensorflow

三、搭建模型

使用 TensorFlow Keras API 构建一个简单的卷积神经网络（CNN）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import tensorflow as tf
from tensorflow.keras import layers, models

model = models.Sequential([
    layers.Rescaling(1./255, input_shape=(180, 180, 3)),
    layers.Conv2D(32, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Conv2D(64, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Conv2D(128, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(num_classes, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

四、加载和预处理数据

使用 TensorFlow 提供的 image_dataset_from_directory 方便加载图片：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    'dataset',
    validation_split=0.2,
    subset='training',
    seed=123,
    image_size=(180, 180),
    batch_size=32)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    'dataset',
    validation_split=0.2,
    subset='validation',
    seed=123,
    image_size=(180, 180),
    batch_size=32)

五、训练模型

1
model.fit(train_ds, validation_data=val_ds, epochs=10)

六、保存模型

1
model.save('my_image_classifier.h5')

七、推理和测试

加载模型并预测新图片类别：

1
2
3
4
5
6
7
8
9
10
11
12
import numpy as np
from tensorflow.keras.preprocessing import image

model = tf.keras.models.load_model('my_image_classifier.h5')

img = image.load_img('test.jpg', target_size=(180, 180))
img_array = image.img_to_array(img) / 255.0
img_array = np.expand_dims(img_array, axis=0)

predictions = model.predict(img_array)
predicted_class = tf.argmax(predictions[0]).numpy()
print(f"预测类别索引为: {predicted_class}")

上面是 TensorFlow 训练自定义图像分类器的基本步骤。如果有大佬有更高效的技巧，欢迎留言指点！