Eye-Movement-Recognition / model_building_sequences.py
shayan5422's picture
Upload 12 files
1d4559c verified
# model_building_sequences.py
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, TimeDistributed, LSTM, BatchNormalization
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
import pickle
def build_cnn_lstm_model(input_shape, num_classes):
"""
Builds a CNN-LSTM model for sequence classification.
Args:
input_shape (tuple): Shape of input sequences (frames, height, width, channels).
num_classes (int): Number of output classes.
Returns:
tensorflow.keras.Model: Compiled model.
"""
model = Sequential()
# Apply Conv2D to each frame in the sequence
model.add(TimeDistributed(Conv2D(32, (3,3), activation='relu'), input_shape=input_shape))
model.add(TimeDistributed(MaxPooling2D((2,2))))
model.add(TimeDistributed(BatchNormalization()))
# Additional Conv2D layers
model.add(TimeDistributed(Conv2D(64, (3,3), activation='relu')))
model.add(TimeDistributed(MaxPooling2D((2,2))))
model.add(TimeDistributed(BatchNormalization()))
# Flatten the output from Conv layers
model.add(TimeDistributed(Flatten()))
# LSTM layer to capture temporal dependencies
model.add(LSTM(128, return_sequences=False))
model.add(Dropout(0.5))
# Fully connected layer
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
# Output layer with softmax activation for classification
model.add(Dense(num_classes, activation='softmax'))
# Compile the model with Adam optimizer and categorical cross-entropy loss
model.compile(optimizer=Adam(learning_rate=1e-4),
loss='categorical_crossentropy',
metrics=['accuracy'])
return model
def load_dataset_pickle(pickle_path='dataset_sequences.pkl'):
"""
Loads the dataset from a pickle file.
Args:
pickle_path (str): Path to the pickle file.
Returns:
tuple: Split data and label mapping.
"""
with open(pickle_path, 'rb') as f:
data = pickle.load(f)
return data['X_train'], data['X_test'], data['y_train'], data['y_test'], data['label_map']
def main():
# Load the dataset
X_train, X_test, y_train, y_test, label_map = load_dataset_pickle('dataset_sequences.pkl')
num_classes = y_train.shape[1]
input_shape = X_train.shape[1:] # (frames, height, width, channels)
# Build the CNN-LSTM model
model = build_cnn_lstm_model(input_shape, num_classes)
model.summary()
# Define callbacks with updated filepath (.keras)
checkpoint = ModelCheckpoint(
'best_model_sequences.keras', # Changed from .h5 to .keras
monitor='val_accuracy',
save_best_only=True,
mode='max'
)
early_stop = EarlyStopping(
monitor='val_accuracy',
patience=10,
restore_best_weights=True
)
# Train the model using GPU
with tf.device('/GPU:0'):
history = model.fit(
X_train, y_train,
epochs=50,
batch_size=128, # Adjust based on your system's memory
validation_data=(X_test, y_test)
)
# Save the final trained model with .keras extension
model.save('final_model_sequences.keras') # Changed from .h5 to .keras
print("Model training completed and saved as 'final_model_sequences.keras'.")
# Save training history for future reference
with open('history_sequences.pkl', 'wb') as f:
pickle.dump(history.history, f)
print("Training history saved as 'history_sequences.pkl'.")
if __name__ == "__main__":
main()