训练一个火灾检测模型完整代码
时间: 2023-11-19 16:06:57 浏览: 28
下面是一个简单的火灾检测模型的训练代码,使用的是 TensorFlow 框架:
```python
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# 设置训练和验证数据集路径
train_dir = "/path/to/train/dataset"
validation_dir = "/path/to/validation/dataset"
# 设置图像增强器
train_image_generator = ImageDataGenerator(rescale=1./255,
rotation_range=45,
zoom_range=0.5,
horizontal_flip=True,
vertical_flip=True)
validation_image_generator = ImageDataGenerator(rescale=1./255)
# 定义模型
model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(16, (3,3), activation='relu', input_shape=(150, 150, 3)),
tf.keras.layers.MaxPooling2D(2, 2),
tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
tf.keras.layers.MaxPooling2D(2,2),
tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
tf.keras.layers.MaxPooling2D(2,2),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid')
])
# 编译模型
model.compile(loss='binary_crossentropy',
optimizer=tf.keras.optimizers.Adam(lr=0.001),
metrics=['accuracy'])
# 训练模型
batch_size = 16
epochs = 50
train_data_gen = train_image_generator.flow_from_directory(batch_size=batch_size,
directory=train_dir,
shuffle=True,
target_size=(150, 150),
class_mode='binary')
val_data_gen = validation_image_generator.flow_from_directory(batch_size=batch_size,
directory=validation_dir,
target_size=(150, 150),
class_mode='binary')
history = model.fit_generator(train_data_gen,
steps_per_epoch=train_data_gen.samples // batch_size,
epochs=epochs,
validation_data=val_data_gen,
validation_steps=val_data_gen.samples // batch_size)
# 保存模型
model.save("fire_detection_model.h5")
```
请注意,在上面的代码中,我们使用了 TensorFlow 的 `ImageDataGenerator` 类来增强训练数据集中的图像。我们还使用了一个简单的卷积神经网络模型来训练我们的数据集,并使用了二元交叉熵作为损失函数和 Adam 优化器来编译模型。最后,我们使用 `fit_generator` 函数来训练模型并保存模型。