Spaces:
Runtime error
Runtime error
File size: 6,034 Bytes
c012514 | 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 167 | #!/usr/bin/env python3
"""
Model Optimization Script for SmokeSignal AI
Compresses the Keras model to reduce deployment size and improve load times
"""
import os
import sys
import tensorflow as tf
from tensorflow.keras.models import load_model, save_model
import numpy as np
def compress_model(input_path, output_path=None, compression_factor=0.8):
"""
Compress a Keras model by reducing precision and optimizing weights
Args:
input_path: Path to the input model
output_path: Path to save the compressed model (optional)
compression_factor: Factor to reduce model size (0.1-1.0)
"""
if output_path is None:
# Create compressed version in same directory
base_name = os.path.splitext(input_path)[0]
output_path = f"{base_name}_compressed.keras"
print(f"🔧 Loading model from {input_path}...")
try:
# Load the model
model = load_model(input_path)
# Get original model size
original_size = os.path.getsize(input_path) / (1024 * 1024)
print(f"📊 Original model size: {original_size:.2f} MB")
# Convert to float16 for compression
print("🔄 Converting model to float16 precision...")
# Create a new model with float16 weights
compressed_model = tf.keras.models.clone_model(model)
compressed_model.set_weights(model.get_weights())
# Convert weights to float16
for layer in compressed_model.layers:
if hasattr(layer, 'weights') and layer.weights:
weights = layer.get_weights()
compressed_weights = []
for weight in weights:
# Convert to float16 and back to float32 for compatibility
compressed_weight = tf.cast(weight, tf.float16)
compressed_weight = tf.cast(compressed_weight, tf.float32)
compressed_weights.append(compressed_weight)
layer.set_weights(compressed_weights)
# Save compressed model
print(f"💾 Saving compressed model to {output_path}...")
save_model(compressed_model, output_path, include_optimizer=False)
# Get compressed model size
compressed_size = os.path.getsize(output_path) / (1024 * 1024)
compression_ratio = (1 - compressed_size / original_size) * 100
print(f"✅ Compression complete!")
print(f"📊 Compressed model size: {compressed_size:.2f} MB")
print(f"📈 Compression ratio: {compression_ratio:.1f}%")
print(f"💾 Size reduction: {original_size - compressed_size:.2f} MB")
return output_path, compressed_size, compression_ratio
except Exception as e:
print(f"❌ Error compressing model: {str(e)}")
return None, None, None
def optimize_model_for_deployment():
"""Main function to optimize the model for deployment"""
model_path = "model/wildfire_detector_model.keras"
if not os.path.exists(model_path):
print(f"❌ Model not found at {model_path}")
return False
print("🚀 SmokeSignal AI Model Optimization")
print("=" * 40)
# Compress the model
compressed_path, compressed_size, compression_ratio = compress_model(model_path)
if compressed_path:
print(f"\n💡 Optimization Summary:")
print(f" - Original: {os.path.getsize(model_path) / (1024*1024):.2f} MB")
print(f" - Compressed: {compressed_size:.2f} MB")
print(f" - Savings: {compression_ratio:.1f}%")
# Create backup of original
backup_path = model_path.replace('.keras', '_backup.keras')
if not os.path.exists(backup_path):
import shutil
shutil.copy2(model_path, backup_path)
print(f" - Backup created: {backup_path}")
# Option to replace original
replace = input("\n🤔 Replace original model with compressed version? (y/n): ").lower()
if replace == 'y':
import shutil
shutil.move(compressed_path, model_path)
print(f"✅ Original model replaced with compressed version")
else:
print(f"📁 Compressed model saved as: {compressed_path}")
print(f"💡 Update your app.py to use the compressed model if desired")
return True
else:
print("❌ Model optimization failed")
return False
def test_model_performance(model_path):
"""Test model performance with sample data"""
print(f"\n🧪 Testing model performance...")
try:
model = load_model(model_path)
# Create dummy input
input_shape = model.input_shape[1:]
dummy_input = np.random.random((1,) + input_shape).astype(np.float32)
# Test inference time
import time
start_time = time.time()
for _ in range(10): # Run 10 predictions
prediction = model.predict(dummy_input, verbose=0)
end_time = time.time()
avg_time = (end_time - start_time) / 10
print(f"✅ Model loaded successfully")
print(f"📊 Average inference time: {avg_time:.3f} seconds")
print(f"📈 Prediction shape: {prediction.shape}")
return True
except Exception as e:
print(f"❌ Model test failed: {str(e)}")
return False
if __name__ == "__main__":
# Check if TensorFlow is available
try:
import tensorflow as tf
print(f"✅ TensorFlow version: {tf.__version__}")
except ImportError:
print("❌ TensorFlow not installed. Please install tensorflow-cpu")
sys.exit(1)
# Run optimization
success = optimize_model_for_deployment()
if success:
# Test the model
test_model_performance("model/wildfire_detector_model.keras")
print("\n🎉 Optimization complete!")
|