zhutoutoutousan
338 lines
12 KiB
Python
338 lines
12 KiB
Python
"""
|
|
ONNX Model Training Script for RSI Divergence Classification
|
|
Trains a neural network to identify genuine RSI divergences and exports to ONNX format.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import pandas as pd
|
|
import tensorflow as tf
|
|
from tensorflow import keras
|
|
from tensorflow.keras import layers
|
|
from sklearn.preprocessing import MinMaxScaler, LabelEncoder
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.metrics import classification_report, confusion_matrix
|
|
import tf2onnx
|
|
import onnx
|
|
import pickle
|
|
from tqdm import tqdm
|
|
|
|
|
|
class RSIDivergenceTrainer:
|
|
"""
|
|
Trainer class for creating ONNX models to classify RSI divergences.
|
|
"""
|
|
|
|
def __init__(self, lookback: int = 60, num_classes: int = 5):
|
|
"""
|
|
Initialize the trainer.
|
|
|
|
Args:
|
|
lookback: Number of bars to look back for prediction
|
|
num_classes: Number of divergence classes (5: NONE + 4 divergence types)
|
|
"""
|
|
self.lookback = lookback
|
|
self.num_classes = num_classes
|
|
self.scaler = MinMaxScaler()
|
|
self.label_encoder = LabelEncoder()
|
|
self.model = None
|
|
|
|
def load_data(self, data_path: str) -> tuple:
|
|
"""
|
|
Load labeled data from CSV file.
|
|
|
|
Args:
|
|
data_path: Path to labeled CSV file
|
|
|
|
Returns:
|
|
Tuple of (X, y) where X is features and y is labels
|
|
"""
|
|
print(f"Loading data from {data_path}...")
|
|
df = pd.read_csv(data_path, index_col=0, parse_dates=True)
|
|
|
|
# Exclude label columns from features
|
|
exclude_cols = ['divergence_type', 'divergence_confidence', 'divergence_strength']
|
|
feature_cols = [col for col in df.columns if col not in exclude_cols]
|
|
|
|
# Remove any remaining non-numeric columns
|
|
feature_cols = [col for col in feature_cols if df[col].dtype in [np.float64, np.int64, np.float32, np.int32]]
|
|
|
|
print(f"Using {len(feature_cols)} features: {feature_cols[:10]}...")
|
|
|
|
# Prepare sequences
|
|
X, y = [], []
|
|
|
|
for i in range(self.lookback, len(df)):
|
|
# Get feature sequence
|
|
X.append(df[feature_cols].iloc[i - self.lookback:i].values)
|
|
|
|
# Get label (divergence type at current bar)
|
|
y.append(int(df['divergence_type'].iloc[i]))
|
|
|
|
X = np.array(X)
|
|
y = np.array(y)
|
|
|
|
print(f"Created {len(X)} sequences")
|
|
print(f"Label distribution: {np.bincount(y)}")
|
|
|
|
return X, y, feature_cols
|
|
|
|
def prepare_data(self, X: np.ndarray, y: np.ndarray) -> tuple:
|
|
"""
|
|
Prepare and scale data for training.
|
|
|
|
Args:
|
|
X: Feature sequences
|
|
y: Labels
|
|
|
|
Returns:
|
|
Tuple of (X_scaled, y_encoded, X_train, X_test, y_train, y_test)
|
|
"""
|
|
# Scale features
|
|
print("Scaling features...")
|
|
original_shape = X.shape
|
|
X_reshaped = X.reshape(-1, X.shape[-1])
|
|
X_scaled = self.scaler.fit_transform(X_reshaped)
|
|
X_scaled = X_scaled.reshape(original_shape)
|
|
|
|
# Encode labels (already integers, but ensure they're 0-4)
|
|
y_encoded = y.astype(int)
|
|
|
|
# Split data (no shuffle to preserve temporal order)
|
|
X_train, X_test, y_train, y_test = train_test_split(
|
|
X_scaled, y_encoded, test_size=0.2, shuffle=False
|
|
)
|
|
|
|
print(f"Training samples: {len(X_train)}")
|
|
print(f"Test samples: {len(X_test)}")
|
|
|
|
return X_scaled, y_encoded, X_train, X_test, y_train, y_test
|
|
|
|
def build_model(self, input_shape: tuple) -> keras.Model:
|
|
"""
|
|
Build the neural network model for classification.
|
|
|
|
Args:
|
|
input_shape: Shape of input data (lookback, features)
|
|
|
|
Returns:
|
|
Compiled Keras model
|
|
"""
|
|
model = keras.Sequential([
|
|
# LSTM layers for sequence learning
|
|
layers.LSTM(128, return_sequences=True, input_shape=input_shape),
|
|
layers.Dropout(0.3),
|
|
layers.LSTM(64, return_sequences=True),
|
|
layers.Dropout(0.3),
|
|
layers.LSTM(32),
|
|
layers.Dropout(0.3),
|
|
|
|
# Dense layers for classification
|
|
layers.Dense(64, activation='relu'),
|
|
layers.Dropout(0.2),
|
|
layers.Dense(32, activation='relu'),
|
|
layers.Dropout(0.2),
|
|
layers.Dense(self.num_classes, activation='softmax') # Multi-class classification
|
|
])
|
|
|
|
model.compile(
|
|
optimizer=keras.optimizers.Adam(learning_rate=0.001),
|
|
loss='sparse_categorical_crossentropy',
|
|
metrics=['accuracy']
|
|
)
|
|
|
|
return model
|
|
|
|
def train(self, X_train: np.ndarray, y_train: np.ndarray,
|
|
X_test: np.ndarray, y_test: np.ndarray,
|
|
epochs: int = 50, batch_size: int = 32, verbose: int = 1):
|
|
"""
|
|
Train the model.
|
|
|
|
Args:
|
|
X_train: Training features
|
|
y_train: Training labels
|
|
X_test: Test features
|
|
y_test: Test labels
|
|
epochs: Number of training epochs
|
|
batch_size: Batch size for training
|
|
verbose: Verbosity level
|
|
"""
|
|
# Build model
|
|
self.model = self.build_model((X_train.shape[1], X_train.shape[2]))
|
|
|
|
print("\nModel architecture:")
|
|
self.model.summary()
|
|
|
|
# Handle class imbalance with class weights
|
|
from sklearn.utils.class_weight import compute_class_weight
|
|
class_weights = compute_class_weight(
|
|
'balanced',
|
|
classes=np.unique(y_train),
|
|
y=y_train
|
|
)
|
|
class_weight_dict = {i: weight for i, weight in enumerate(class_weights)}
|
|
|
|
print(f"\nClass weights: {class_weight_dict}")
|
|
|
|
# Train model
|
|
print("\nTraining model...")
|
|
history = self.model.fit(
|
|
X_train, y_train,
|
|
batch_size=batch_size,
|
|
epochs=epochs,
|
|
validation_data=(X_test, y_test),
|
|
verbose=verbose,
|
|
class_weight=class_weight_dict,
|
|
callbacks=[
|
|
keras.callbacks.EarlyStopping(
|
|
monitor='val_loss',
|
|
patience=15,
|
|
restore_best_weights=True,
|
|
verbose=1
|
|
),
|
|
keras.callbacks.ReduceLROnPlateau(
|
|
monitor='val_loss',
|
|
factor=0.5,
|
|
patience=5,
|
|
min_lr=0.0001,
|
|
verbose=1
|
|
)
|
|
]
|
|
)
|
|
|
|
# Evaluate
|
|
train_loss, train_acc = self.model.evaluate(X_train, y_train, verbose=0)
|
|
test_loss, test_acc = self.model.evaluate(X_test, y_test, verbose=0)
|
|
|
|
print(f"\nTraining - Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}")
|
|
print(f"Test - Loss: {test_loss:.4f}, Accuracy: {test_acc:.4f}")
|
|
|
|
# Classification report
|
|
y_pred = self.model.predict(X_test, verbose=0)
|
|
y_pred_classes = np.argmax(y_pred, axis=1)
|
|
|
|
print("\nClassification Report:")
|
|
print(classification_report(y_test, y_pred_classes,
|
|
target_names=['NONE', 'REGULAR_BULLISH', 'REGULAR_BEARISH',
|
|
'HIDDEN_BULLISH', 'HIDDEN_BEARISH']))
|
|
|
|
return history
|
|
|
|
def export_to_onnx(self, output_path: str, num_features: int):
|
|
"""
|
|
Export the trained model to ONNX format.
|
|
|
|
Args:
|
|
output_path: Path to save ONNX model
|
|
num_features: Number of input features
|
|
"""
|
|
if self.model is None:
|
|
raise ValueError("Model must be trained before exporting")
|
|
|
|
print(f"\nExporting model to ONNX format: {output_path}")
|
|
|
|
# Create functional model from Sequential
|
|
input_layer = keras.Input(shape=(self.lookback, num_features), name="input")
|
|
x = input_layer
|
|
|
|
# Rebuild model as functional
|
|
for layer in self.model.layers:
|
|
x = layer(x)
|
|
|
|
functional_model = keras.Model(inputs=input_layer, outputs=x)
|
|
|
|
# Convert to ONNX
|
|
spec = (tf.TensorSpec((None, self.lookback, num_features), tf.float32, name="input"),)
|
|
|
|
try:
|
|
onnx_model_proto, _ = tf2onnx.convert.from_keras(
|
|
functional_model,
|
|
input_signature=spec,
|
|
opset=13
|
|
)
|
|
|
|
onnx.save_model(onnx_model_proto, output_path)
|
|
print(f"ONNX model saved to: {output_path}")
|
|
|
|
# Verify ONNX model
|
|
onnx_model = onnx.load(output_path)
|
|
onnx.checker.check_model(onnx_model)
|
|
print("ONNX model validation passed")
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to export ONNX model: {str(e)}")
|
|
|
|
def save_scaler(self, output_path: str):
|
|
"""Save the scaler for consistent normalization."""
|
|
with open(output_path, 'wb') as f:
|
|
pickle.dump(self.scaler, f)
|
|
print(f"Scaler saved to: {output_path}")
|
|
|
|
|
|
def main():
|
|
"""Main function."""
|
|
parser = argparse.ArgumentParser(description='Train ONNX model for RSI divergence classification')
|
|
parser.add_argument('--data', type=str, required=True,
|
|
help='Path to labeled CSV data file')
|
|
parser.add_argument('--lookback', type=int, default=60,
|
|
help='Number of bars to look back')
|
|
parser.add_argument('--epochs', type=int, default=50, help='Training epochs')
|
|
parser.add_argument('--batch-size', type=int, default=32, help='Batch size')
|
|
parser.add_argument('--output', type=str, default='models',
|
|
help='Output directory for ONNX model')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Create output directory
|
|
os.makedirs(args.output, exist_ok=True)
|
|
|
|
# Create trainer
|
|
trainer = RSIDivergenceTrainer(lookback=args.lookback)
|
|
|
|
try:
|
|
# Load data
|
|
X, y, feature_cols = trainer.load_data(args.data)
|
|
|
|
# Prepare data
|
|
X_scaled, y_encoded, X_train, X_test, y_train, y_test = trainer.prepare_data(X, y)
|
|
|
|
# Train model
|
|
trainer.train(X_train, y_train, X_test, y_test,
|
|
epochs=args.epochs, batch_size=args.batch_size)
|
|
|
|
# Export to ONNX
|
|
num_features = len(feature_cols)
|
|
model_name = "BTCUSD_H1_rsi_divergence_model.onnx"
|
|
output_path = os.path.join(args.output, model_name)
|
|
trainer.export_to_onnx(output_path, num_features)
|
|
|
|
# Save scaler
|
|
scaler_name = "BTCUSD_H1_rsi_divergence_scaler.pkl"
|
|
scaler_path = os.path.join(args.output, scaler_name)
|
|
trainer.save_scaler(scaler_path)
|
|
|
|
# Save feature list
|
|
features_name = "BTCUSD_H1_rsi_divergence_features.pkl"
|
|
features_path = os.path.join(args.output, features_name)
|
|
with open(features_path, 'wb') as f:
|
|
pickle.dump(feature_cols, f)
|
|
print(f"Feature list saved to: {features_path}")
|
|
|
|
print(f"\n{'='*60}")
|
|
print("Training completed successfully!")
|
|
print(f"ONNX model saved to: {output_path}")
|
|
print(f"{'='*60}\n")
|
|
|
|
except Exception as e:
|
|
print(f"\nError: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|