Spaces:
Paused
Paused
File size: 5,525 Bytes
1feed70 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """
src/training.py
Training pipeline for FirenetCNN model
"""
import tensorflow as tf
from pathlib import Path
from typing import Optional, Tuple, Dict, Any
from .model import FireNetModel
def train(
train_dir: str,
val_dir: str,
model_save_path: str = 'models/FirenetCNN.keras',
epochs: int = 100,
learning_rate: float = 0.0001,
batch_size: int = 32,
image_size: Tuple[int, int] = (224, 224),
fine_tune: bool = False,
fine_tune_epochs: int = 10,
fine_tune_lr: float = 1e-5,
callback_save_best: bool = True,
callback_early_stopping: bool = True,
early_stopping_patience: int = 10,
) -> Dict[str, Any]:
"""
Train the FirenetCNN model.
Args:
train_dir: Path to training data directory
val_dir: Path to validation data directory
model_save_path: Path to save the trained model
epochs: Number of training epochs
learning_rate: Learning rate for initial training
batch_size: Training batch size
image_size: Input image size (height, width)
fine_tune: Whether to fine-tune the base model after initial training
fine_tune_epochs: Number of fine-tuning epochs
fine_tune_lr: Learning rate for fine-tuning
callback_save_best: Save model with best validation accuracy
callback_early_stopping: Stop training if validation loss stops improving
early_stopping_patience: Patience for early stopping
Returns:
Dictionary with training history and results
"""
# Build model
model = FireNetModel.build_model(input_shape=(*image_size, 3))
# Create data generators
train_gen, val_gen = FireNetModel.create_data_generators(train_dir, val_dir)
# Setup callbacks
callbacks = []
save_path = Path(model_save_path)
save_path.parent.mkdir(parents=True, exist_ok=True)
if callback_save_best:
best_model_path = save_path.with_name(save_path.stem + '_best' + save_path.suffix)
callbacks.append(tf.keras.callbacks.ModelCheckpoint(
str(best_model_path),
monitor='val_accuracy',
save_best_only=True,
mode='max',
verbose=1
))
if callback_early_stopping:
callbacks.append(tf.keras.callbacks.EarlyStopping(
monitor='val_loss',
patience=early_stopping_patience,
restore_best_weights=True,
verbose=1
))
# Initial training (frozen base model)
print(f"Phase 1: Training with frozen base for {epochs} epochs...")
history = model.fit(
train_gen,
epochs=epochs,
validation_data=val_gen,
callbacks=callbacks,
verbose=1
)
# Fine-tuning phase
if fine_tune:
print(f"\nPhase 2: Fine-tuning for {fine_tune_epochs} epochs...")
# Unfreeze the base model
base_model = None
for layer in model.layers:
if hasattr(layer, 'layers') and len(layer.layers) > 50:
base_model = layer
break
if base_model is not None:
base_model.trainable = True
# Recompile with lower learning rate
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=fine_tune_lr),
loss='categorical_crossentropy',
metrics=['accuracy']
)
# Continue training
history_fine = model.fit(
train_gen,
epochs=fine_tune_epochs,
validation_data=val_gen,
callbacks=callbacks,
verbose=1
)
# Merge histories
for key in history_fine.history:
history.history[key].extend(history_fine.history[key])
# Save final model
model.save(str(save_path))
print(f"\nModel saved to: {save_path}")
return {
'history': history.history,
'epochs_completed': len(history.history['accuracy']),
'final_train_acc': history.history['accuracy'][-1],
'final_val_acc': history.history['val_accuracy'][-1],
'model_path': str(save_path)
}
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='Train FirenetCNN model')
parser.add_argument('--train-dir', default='data/forestfire-classifier-dataset/train',
help='Training data directory')
parser.add_argument('--val-dir', default='data/forestfire-classifier-dataset/val',
help='Validation data directory')
parser.add_argument('--model-path', default='models/FirenetCNN.keras',
help='Path to save trained model')
parser.add_argument('--epochs', type=int, default=100,
help='Number of training epochs')
parser.add_argument('--fine-tune', action='store_true',
help='Enable fine-tuning phase')
parser.add_argument('--fine-tune-epochs', type=int, default=10,
help='Number of fine-tuning epochs')
args = parser.parse_args()
results = train(
train_dir=args.train_dir,
val_dir=args.val_dir,
model_save_path=args.model_path,
epochs=args.epochs,
fine_tune=args.fine_tune,
fine_tune_epochs=args.fine_tune_epochs
)
print(f"\nTraining complete!")
print(f"Final train accuracy: {results['final_train_acc']:.4f}")
print(f"Final val accuracy: {results['final_val_acc']:.4f}")
|