Update
This commit is contained in:
zhutoutoutousan
@@ -0,0 +1,364 @@
|
||||
"""
|
||||
ONNX Model Training Script for MetaTrader 5
|
||||
|
||||
This script trains a neural network model for price prediction and exports it to ONNX format.
|
||||
Based on MQL5 ONNX documentation: https://www.mql5.com/en/docs/onnx/onnx_prepare
|
||||
|
||||
Usage:
|
||||
python train_onnx_model.py --symbol XAUUSD --timeframe H1 --lookback 60 --epochs 50
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import MetaTrader5 as mt5
|
||||
import tensorflow as tf
|
||||
from tensorflow import keras
|
||||
from tensorflow.keras import layers
|
||||
from sklearn.preprocessing import MinMaxScaler
|
||||
from sklearn.model_selection import train_test_split
|
||||
import tf2onnx
|
||||
import onnx
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
class ONNXModelTrainer:
|
||||
"""
|
||||
Trainer class for creating ONNX models from MT5 data.
|
||||
"""
|
||||
|
||||
def __init__(self, symbol: str, timeframe: int, lookback: int = 60,
|
||||
prediction_horizon: int = 1, features: list = None):
|
||||
"""
|
||||
Initialize the trainer.
|
||||
|
||||
Args:
|
||||
symbol: Trading symbol (e.g., 'XAUUSD', 'EURUSD')
|
||||
timeframe: MT5 timeframe constant
|
||||
lookback: Number of bars to look back for prediction
|
||||
prediction_horizon: Number of bars ahead to predict
|
||||
features: List of features to use (default: OHLC + volume)
|
||||
"""
|
||||
self.symbol = symbol
|
||||
self.timeframe = timeframe
|
||||
self.lookback = lookback
|
||||
self.prediction_horizon = prediction_horizon
|
||||
self.features = features or ['open', 'high', 'low', 'close', 'tick_volume']
|
||||
|
||||
self.scaler = MinMaxScaler()
|
||||
self.model = None
|
||||
|
||||
# Initialize MT5
|
||||
if not mt5.initialize():
|
||||
raise RuntimeError(f"MT5 initialization failed: {mt5.last_error()}")
|
||||
|
||||
def fetch_data(self, start_date: datetime, end_date: datetime) -> pd.DataFrame:
|
||||
"""
|
||||
Fetch historical data from MT5.
|
||||
|
||||
Args:
|
||||
start_date: Start date for data
|
||||
end_date: End date for data
|
||||
|
||||
Returns:
|
||||
DataFrame with OHLCV data
|
||||
"""
|
||||
print(f"Fetching data for {self.symbol} from {start_date} to {end_date}...")
|
||||
|
||||
rates = mt5.copy_rates_range(self.symbol, self.timeframe, start_date, end_date)
|
||||
|
||||
if rates is None or len(rates) == 0:
|
||||
raise ValueError(f"No data available for {self.symbol} in the specified date range")
|
||||
|
||||
df = pd.DataFrame(rates)
|
||||
df['time'] = pd.to_datetime(df['time'], unit='s')
|
||||
df.set_index('time', inplace=True)
|
||||
|
||||
print(f"Fetched {len(df)} bars")
|
||||
return df
|
||||
|
||||
def prepare_features(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
"""
|
||||
Prepare features for training.
|
||||
|
||||
Args:
|
||||
df: Raw OHLCV data
|
||||
|
||||
Returns:
|
||||
DataFrame with features
|
||||
"""
|
||||
feature_df = df[self.features].copy()
|
||||
|
||||
# Add technical indicators as features
|
||||
feature_df['rsi'] = self._calculate_rsi(df['close'], period=14)
|
||||
feature_df['ema_20'] = df['close'].ewm(span=20).mean()
|
||||
feature_df['ema_50'] = df['close'].ewm(span=50).mean()
|
||||
feature_df['atr'] = self._calculate_atr(df, period=14)
|
||||
|
||||
# Price changes
|
||||
feature_df['price_change'] = df['close'].pct_change()
|
||||
feature_df['high_low_ratio'] = df['high'] / df['low']
|
||||
|
||||
# Volume features
|
||||
feature_df['volume_ma'] = df['tick_volume'].rolling(window=20).mean()
|
||||
feature_df['volume_ratio'] = df['tick_volume'] / feature_df['volume_ma']
|
||||
|
||||
# Drop NaN values
|
||||
feature_df = feature_df.dropna()
|
||||
|
||||
return feature_df
|
||||
|
||||
def _calculate_rsi(self, prices: pd.Series, period: int = 14) -> pd.Series:
|
||||
"""Calculate RSI indicator."""
|
||||
delta = prices.diff()
|
||||
gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
|
||||
loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
|
||||
rs = gain / loss
|
||||
rsi = 100 - (100 / (1 + rs))
|
||||
return rsi
|
||||
|
||||
def _calculate_atr(self, df: pd.DataFrame, period: int = 14) -> pd.Series:
|
||||
"""Calculate ATR indicator."""
|
||||
high_low = df['high'] - df['low']
|
||||
high_close = np.abs(df['high'] - df['close'].shift())
|
||||
low_close = np.abs(df['low'] - df['close'].shift())
|
||||
tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
|
||||
atr = tr.rolling(window=period).mean()
|
||||
return atr
|
||||
|
||||
def create_sequences(self, data: np.ndarray, target: np.ndarray) -> tuple:
|
||||
"""
|
||||
Create sequences for LSTM/RNN training.
|
||||
|
||||
Args:
|
||||
data: Feature data
|
||||
target: Target values
|
||||
|
||||
Returns:
|
||||
Tuple of (X, y) sequences
|
||||
"""
|
||||
X, y = [], []
|
||||
|
||||
for i in range(self.lookback, len(data) - self.prediction_horizon + 1):
|
||||
X.append(data[i - self.lookback:i])
|
||||
y.append(target[i + self.prediction_horizon - 1])
|
||||
|
||||
return np.array(X), np.array(y)
|
||||
|
||||
def build_model(self, input_shape: tuple) -> keras.Model:
|
||||
"""
|
||||
Build the neural network model.
|
||||
|
||||
Args:
|
||||
input_shape: Shape of input data (lookback, features)
|
||||
|
||||
Returns:
|
||||
Compiled Keras model
|
||||
"""
|
||||
model = keras.Sequential([
|
||||
layers.LSTM(128, return_sequences=True, input_shape=input_shape),
|
||||
layers.Dropout(0.2),
|
||||
layers.LSTM(64, return_sequences=True),
|
||||
layers.Dropout(0.2),
|
||||
layers.LSTM(32),
|
||||
layers.Dropout(0.2),
|
||||
layers.Dense(16, activation='relu'),
|
||||
layers.Dense(1) # Predict next close price
|
||||
])
|
||||
|
||||
model.compile(
|
||||
optimizer=keras.optimizers.Adam(learning_rate=0.001),
|
||||
loss='mse',
|
||||
metrics=['mae']
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
def train(self, epochs: int = 50, batch_size: int = 32,
|
||||
validation_split: float = 0.2, verbose: int = 1):
|
||||
"""
|
||||
Train the model.
|
||||
|
||||
Args:
|
||||
epochs: Number of training epochs
|
||||
batch_size: Batch size for training
|
||||
validation_split: Fraction of data to use for validation
|
||||
verbose: Verbosity level
|
||||
"""
|
||||
# Fetch data (last 2 years)
|
||||
end_date = datetime.now()
|
||||
start_date = end_date - timedelta(days=730)
|
||||
|
||||
df = self.fetch_data(start_date, end_date)
|
||||
feature_df = self.prepare_features(df)
|
||||
|
||||
# Prepare data
|
||||
feature_data = feature_df.values
|
||||
target_data = df['close'].values[feature_df.index]
|
||||
|
||||
# Scale features
|
||||
feature_data_scaled = self.scaler.fit_transform(feature_data)
|
||||
|
||||
# Create sequences
|
||||
X, y = self.create_sequences(feature_data_scaled, target_data)
|
||||
|
||||
# Split data
|
||||
X_train, X_test, y_train, y_test = train_test_split(
|
||||
X, y, test_size=validation_split, shuffle=False
|
||||
)
|
||||
|
||||
print(f"\nTraining data shape: {X_train.shape}")
|
||||
print(f"Validation data shape: {X_test.shape}")
|
||||
|
||||
# Build model
|
||||
self.model = self.build_model((X_train.shape[1], X_train.shape[2]))
|
||||
|
||||
print("\nModel architecture:")
|
||||
self.model.summary()
|
||||
|
||||
# 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,
|
||||
callbacks=[
|
||||
keras.callbacks.EarlyStopping(
|
||||
monitor='val_loss',
|
||||
patience=10,
|
||||
restore_best_weights=True
|
||||
),
|
||||
keras.callbacks.ReduceLROnPlateau(
|
||||
monitor='val_loss',
|
||||
factor=0.5,
|
||||
patience=5,
|
||||
min_lr=0.0001
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Evaluate
|
||||
train_loss = self.model.evaluate(X_train, y_train, verbose=0)
|
||||
test_loss = self.model.evaluate(X_test, y_test, verbose=0)
|
||||
|
||||
print(f"\nTraining Loss: {train_loss[0]:.4f}, MAE: {train_loss[1]:.4f}")
|
||||
print(f"Validation Loss: {test_loss[0]:.4f}, MAE: {test_loss[1]:.4f}")
|
||||
|
||||
return history
|
||||
|
||||
def export_to_onnx(self, output_path: str):
|
||||
"""
|
||||
Export the trained model to ONNX format.
|
||||
|
||||
Args:
|
||||
output_path: Path to save ONNX model
|
||||
"""
|
||||
if self.model is None:
|
||||
raise ValueError("Model must be trained before exporting")
|
||||
|
||||
print(f"\nExporting model to ONNX format: {output_path}")
|
||||
|
||||
# Get input shape
|
||||
input_shape = (1, self.lookback, len(self.features) + 7) # +7 for added features
|
||||
|
||||
# Create dummy input
|
||||
dummy_input = np.random.randn(*input_shape).astype(np.float32)
|
||||
|
||||
# Convert to ONNX
|
||||
spec = (tf.TensorSpec((None, self.lookback, len(self.features) + 7), tf.float32, name="input"),)
|
||||
output_path_onnx = tf2onnx.convert.from_keras(
|
||||
self.model,
|
||||
input_signature=spec,
|
||||
opset=13,
|
||||
output_path=output_path
|
||||
)
|
||||
|
||||
print(f"✓ ONNX model saved to: {output_path}")
|
||||
|
||||
# Verify ONNX model
|
||||
try:
|
||||
onnx_model = onnx.load(output_path)
|
||||
onnx.checker.check_model(onnx_model)
|
||||
print("✓ ONNX model validation passed")
|
||||
except Exception as e:
|
||||
print(f"⚠ ONNX model validation warning: {e}")
|
||||
|
||||
def cleanup(self):
|
||||
"""Clean up MT5 connection."""
|
||||
mt5.shutdown()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function."""
|
||||
parser = argparse.ArgumentParser(description='Train ONNX model for MT5 price prediction')
|
||||
parser.add_argument('--symbol', type=str, default='XAUUSD', help='Trading symbol')
|
||||
parser.add_argument('--timeframe', type=str, default='H1',
|
||||
choices=['M1', 'M5', 'M15', 'M30', 'H1', 'H4', 'D1'],
|
||||
help='Timeframe')
|
||||
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()
|
||||
|
||||
# Convert timeframe string to MT5 constant
|
||||
timeframe_map = {
|
||||
'M1': mt5.TIMEFRAME_M1,
|
||||
'M5': mt5.TIMEFRAME_M5,
|
||||
'M15': mt5.TIMEFRAME_M15,
|
||||
'M30': mt5.TIMEFRAME_M30,
|
||||
'H1': mt5.TIMEFRAME_H1,
|
||||
'H4': mt5.TIMEFRAME_H4,
|
||||
'D1': mt5.TIMEFRAME_D1
|
||||
}
|
||||
timeframe = timeframe_map[args.timeframe]
|
||||
|
||||
# Create output directory
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
|
||||
# Create trainer
|
||||
trainer = ONNXModelTrainer(
|
||||
symbol=args.symbol,
|
||||
timeframe=timeframe,
|
||||
lookback=args.lookback
|
||||
)
|
||||
|
||||
try:
|
||||
# Train model
|
||||
trainer.train(epochs=args.epochs, batch_size=args.batch_size)
|
||||
|
||||
# Export to ONNX
|
||||
model_name = f"{args.symbol}_{args.timeframe}_model.onnx"
|
||||
output_path = os.path.join(args.output, model_name)
|
||||
trainer.export_to_onnx(output_path)
|
||||
|
||||
# Save scaler for consistent normalization
|
||||
scaler_path = os.path.join(args.output, f"{args.symbol}_{args.timeframe}_scaler.pkl")
|
||||
import pickle
|
||||
with open(scaler_path, 'wb') as f:
|
||||
pickle.dump(trainer.scaler, f)
|
||||
print(f"✓ Scaler saved to: {scaler_path}")
|
||||
print(f" Use this with predict_with_onnx.py for consistent normalization")
|
||||
|
||||
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}")
|
||||
sys.exit(1)
|
||||
finally:
|
||||
trainer.cleanup()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user