Update
This commit is contained in:
zhutoutoutousan
@@ -0,0 +1,455 @@
|
||||
"""
|
||||
ONNX Model Training Script for BTCUSD 1-Minute Data
|
||||
Uses yfinance (Yahoo Finance) for historical data
|
||||
|
||||
This script trains a neural network model for BTCUSD price prediction on 1-minute timeframe
|
||||
and exports it to ONNX format.
|
||||
|
||||
Usage:
|
||||
python main.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import yfinance as yf
|
||||
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
|
||||
import pickle
|
||||
|
||||
|
||||
class BTCUSD1MinTrainer:
|
||||
"""
|
||||
Trainer class for creating ONNX models from BTCUSD 1-minute MT5 data.
|
||||
"""
|
||||
|
||||
def __init__(self, symbol: str = "BTC-USD", lookback: int = 60,
|
||||
prediction_horizon: int = 1):
|
||||
"""
|
||||
Initialize the trainer.
|
||||
|
||||
Args:
|
||||
symbol: Trading symbol (default: 'BTC-USD' for Yahoo Finance)
|
||||
lookback: Number of bars to look back for prediction (default: 60)
|
||||
prediction_horizon: Number of bars ahead to predict (default: 1)
|
||||
"""
|
||||
self.symbol = symbol
|
||||
self.lookback = lookback
|
||||
self.prediction_horizon = prediction_horizon
|
||||
|
||||
self.scaler = MinMaxScaler()
|
||||
self.model = None
|
||||
|
||||
print(f"Using yfinance for data source. Symbol: {self.symbol}")
|
||||
|
||||
def fetch_data(self, start_date: datetime, end_date: datetime) -> pd.DataFrame:
|
||||
"""
|
||||
Fetch historical data from Yahoo Finance using yfinance.
|
||||
|
||||
Args:
|
||||
start_date: Start date for data
|
||||
end_date: End date for data
|
||||
|
||||
Returns:
|
||||
DataFrame with OHLCV data
|
||||
"""
|
||||
print(f"\nFetching {self.symbol} 1-minute data from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}...")
|
||||
|
||||
# yfinance can only fetch 7 days of 1-minute data at a time
|
||||
# For longer periods, we need to fetch in chunks
|
||||
all_data = []
|
||||
current_start = start_date
|
||||
|
||||
# For 1-minute data, yfinance limits to last 7 days
|
||||
# So we'll fetch the most recent 7 days available
|
||||
print("Note: yfinance 1-minute data is limited to last 7 days")
|
||||
print("Fetching most recent available 1-minute data...")
|
||||
|
||||
# Get ticker
|
||||
ticker = yf.Ticker(self.symbol)
|
||||
|
||||
# Try to fetch 1-minute data (limited to 7 days)
|
||||
# If we need more data, we'll use daily data and resample
|
||||
try:
|
||||
# Fetch 1-minute data (max 7 days)
|
||||
df = ticker.history(start=start_date, end=end_date, interval='1m')
|
||||
|
||||
if df is None or len(df) == 0:
|
||||
print("Warning: No 1-minute data available, trying daily data...")
|
||||
# Fall back to daily data
|
||||
df = ticker.history(start=start_date, end=end_date, interval='1d')
|
||||
if df is None or len(df) == 0:
|
||||
raise ValueError(f"No data available for {self.symbol}")
|
||||
print(f"Using daily data instead (will resample to 1-minute for training)")
|
||||
except Exception as e:
|
||||
print(f"Error fetching 1-minute data: {e}")
|
||||
print("Falling back to daily data...")
|
||||
df = ticker.history(start=start_date, end=end_date, interval='1d')
|
||||
if df is None or len(df) == 0:
|
||||
raise ValueError(f"No data available for {self.symbol}: {e}")
|
||||
|
||||
# Rename columns to match expected format
|
||||
df.columns = [col.lower().replace(' ', '_') for col in df.columns]
|
||||
|
||||
# Ensure we have the required columns
|
||||
required_cols = ['open', 'high', 'low', 'close', 'volume']
|
||||
missing_cols = [col for col in required_cols if col not in df.columns]
|
||||
if missing_cols:
|
||||
raise ValueError(f"Missing required columns: {missing_cols}")
|
||||
|
||||
# Rename 'volume' to 'tick_volume' for consistency
|
||||
if 'volume' in df.columns:
|
||||
df['tick_volume'] = df['volume']
|
||||
df = df.drop('volume', axis=1)
|
||||
|
||||
# Remove duplicates and sort
|
||||
df = df[~df.index.duplicated(keep='first')]
|
||||
df = df.sort_index()
|
||||
|
||||
print(f"Total fetched: {len(df)} bars")
|
||||
if len(df) > 0:
|
||||
print(f"Date range: {df.index[0]} to {df.index[-1]}")
|
||||
print(f"Timeframe: {df.index[1] - df.index[0] if len(df) > 1 else 'N/A'}")
|
||||
else:
|
||||
raise ValueError("DataFrame is empty after processing")
|
||||
|
||||
return df
|
||||
|
||||
def prepare_features(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
"""
|
||||
Prepare features for training.
|
||||
|
||||
Args:
|
||||
df: Raw OHLCV data
|
||||
|
||||
Returns:
|
||||
DataFrame with features
|
||||
"""
|
||||
print("\nPreparing features...")
|
||||
|
||||
feature_df = df[['open', 'high', 'low', 'close', 'tick_volume']].copy()
|
||||
|
||||
# Add technical indicators as features
|
||||
print(" Calculating RSI...")
|
||||
feature_df['rsi'] = self._calculate_rsi(df['close'], period=14)
|
||||
|
||||
print(" Calculating EMAs...")
|
||||
feature_df['ema_20'] = df['close'].ewm(span=20, adjust=False).mean()
|
||||
feature_df['ema_50'] = df['close'].ewm(span=50, adjust=False).mean()
|
||||
|
||||
print(" Calculating ATR...")
|
||||
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()
|
||||
|
||||
print(f" Features prepared: {len(feature_df)} samples, {len(feature_df.columns)} features")
|
||||
print(f" Features: {list(feature_df.columns)}")
|
||||
|
||||
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 (price change percentages)
|
||||
|
||||
Returns:
|
||||
Tuple of (X, y) sequences
|
||||
"""
|
||||
print("\nCreating sequences...")
|
||||
X, y = [], []
|
||||
|
||||
for i in tqdm(range(self.lookback, len(data) - self.prediction_horizon + 1), desc="Creating sequences"):
|
||||
X.append(data[i - self.lookback:i])
|
||||
y.append(target[i])
|
||||
|
||||
X = np.array(X)
|
||||
y = np.array(y)
|
||||
|
||||
print(f" Sequences created: X shape {X.shape}, y shape {y.shape}")
|
||||
|
||||
return X, 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
|
||||
"""
|
||||
print(f"\nBuilding model with input shape: {input_shape}")
|
||||
|
||||
model = keras.Sequential([
|
||||
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),
|
||||
layers.Dense(32, activation='relu'),
|
||||
layers.Dense(16, activation='relu'),
|
||||
layers.Dense(1) # Predict price change percentage
|
||||
])
|
||||
|
||||
model.compile(
|
||||
optimizer=keras.optimizers.Adam(learning_rate=0.0005),
|
||||
loss='mse',
|
||||
metrics=['mae']
|
||||
)
|
||||
|
||||
print(f" Model parameters: {model.count_params():,}")
|
||||
model.summary()
|
||||
|
||||
return model
|
||||
|
||||
def train(self, start_date: datetime, end_date: datetime,
|
||||
epochs: int = 50, batch_size: int = 32,
|
||||
validation_split: float = 0.2, verbose: int = 1):
|
||||
"""
|
||||
Train the model.
|
||||
|
||||
Args:
|
||||
start_date: Start date for training data
|
||||
end_date: End date for training data
|
||||
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
|
||||
df = self.fetch_data(start_date, end_date)
|
||||
feature_df = self.prepare_features(df)
|
||||
|
||||
# Prepare target: price change percentage for next bar
|
||||
# Calculate future price change: (next_close - current_close) / current_close
|
||||
close_prices = feature_df['close'].values
|
||||
target = []
|
||||
for i in range(len(close_prices)):
|
||||
if i + self.prediction_horizon < len(close_prices):
|
||||
current_price = close_prices[i]
|
||||
future_price = close_prices[i + self.prediction_horizon]
|
||||
price_change_pct = (future_price - current_price) / current_price if current_price > 0 else 0.0
|
||||
target.append(price_change_pct)
|
||||
else:
|
||||
target.append(0.0)
|
||||
target = pd.Series(target, index=feature_df.index)
|
||||
|
||||
# Keep 'close' in features - it's needed for the model
|
||||
|
||||
# Align target with features
|
||||
valid_idx = ~(target.isna() | feature_df.isna().any(axis=1))
|
||||
feature_df = feature_df[valid_idx]
|
||||
target = target[valid_idx]
|
||||
|
||||
print(f"\nValid samples after alignment: {len(feature_df)}")
|
||||
|
||||
# Normalize features
|
||||
print("\nNormalizing features...")
|
||||
feature_array = self.scaler.fit_transform(feature_df.values)
|
||||
|
||||
# Create sequences
|
||||
X, y = self.create_sequences(feature_array, target.values)
|
||||
|
||||
# Split into train and validation
|
||||
split_idx = int(len(X) * (1 - validation_split))
|
||||
X_train, X_val = X[:split_idx], X[split_idx:]
|
||||
y_train, y_val = y[:split_idx], y[split_idx:]
|
||||
|
||||
print(f"\nTrain set: {len(X_train)} samples")
|
||||
print(f"Validation set: {len(X_val)} samples")
|
||||
|
||||
# Build model
|
||||
input_shape = (self.lookback, feature_array.shape[1])
|
||||
self.model = self.build_model(input_shape)
|
||||
|
||||
# Train model
|
||||
print(f"\nTraining model for {epochs} epochs...")
|
||||
history = self.model.fit(
|
||||
X_train, y_train,
|
||||
batch_size=batch_size,
|
||||
epochs=epochs,
|
||||
validation_data=(X_val, y_val),
|
||||
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=1e-7
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Evaluate
|
||||
print("\nEvaluating model...")
|
||||
train_loss = self.model.evaluate(X_train, y_train, verbose=0)
|
||||
val_loss = self.model.evaluate(X_val, y_val, verbose=0)
|
||||
|
||||
print(f"Train Loss: {train_loss[0]:.6f}, MAE: {train_loss[1]:.6f}")
|
||||
print(f"Val Loss: {val_loss[0]:.6f}, MAE: {val_loss[1]:.6f}")
|
||||
|
||||
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 number of features
|
||||
num_features = self.model.input_shape[2] if len(self.model.input_shape) > 2 else self.model.input_shape[1]
|
||||
|
||||
print(f"Using {num_features} features for ONNX export")
|
||||
|
||||
# Create input signature
|
||||
input_shape = (None, self.lookback, num_features)
|
||||
spec = (tf.TensorSpec(input_shape, tf.float32, name="input"),)
|
||||
|
||||
# Fix output_names for Sequential model
|
||||
if not hasattr(self.model, 'output_names'):
|
||||
if hasattr(self.model, 'outputs') and self.model.outputs:
|
||||
self.model.output_names = [f'output_{i}' for i in range(len(self.model.outputs))]
|
||||
else:
|
||||
self.model.output_names = ['output']
|
||||
|
||||
# Convert to ONNX
|
||||
onnx_model, _ = tf2onnx.convert.from_keras(
|
||||
self.model,
|
||||
input_signature=spec,
|
||||
opset=13
|
||||
)
|
||||
|
||||
# Save ONNX model
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
onnx.save_model(onnx_model, output_path)
|
||||
|
||||
print(f"ONNX model saved to: {output_path}")
|
||||
|
||||
# Save scaler
|
||||
scaler_path = output_path.replace('.onnx', '_scaler.pkl')
|
||||
with open(scaler_path, 'wb') as f:
|
||||
pickle.dump(self.scaler, f)
|
||||
print(f"Scaler saved to: {scaler_path}")
|
||||
|
||||
def cleanup(self):
|
||||
"""Clean up (no-op for yfinance)."""
|
||||
pass
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function."""
|
||||
print("="*60)
|
||||
print("BTCUSD 1-Minute ONNX Model Training")
|
||||
print("="*60)
|
||||
|
||||
# Training parameters
|
||||
symbol = "BTC-USD" # Yahoo Finance symbol
|
||||
lookback = 60 # 60 minutes of history
|
||||
epochs = 50
|
||||
batch_size = 64 # Larger batch for 1-minute data
|
||||
|
||||
# Date range: Use recent data (yfinance 1m data limited to 7 days)
|
||||
# For longer training, we'll use the most recent available data
|
||||
end_date = datetime.now()
|
||||
start_date = end_date - timedelta(days=7) # Last 7 days for 1-minute data
|
||||
|
||||
print(f"Note: yfinance 1-minute data is limited to last 7 days")
|
||||
print(f"Using date range: {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}")
|
||||
|
||||
# Output paths
|
||||
output_dir = "models"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
model_path = os.path.join(output_dir, f"{symbol}_M1_model.onnx")
|
||||
|
||||
trainer = None
|
||||
try:
|
||||
# Create trainer
|
||||
trainer = BTCUSD1MinTrainer(
|
||||
symbol=symbol,
|
||||
lookback=lookback
|
||||
)
|
||||
|
||||
# Train model
|
||||
history = trainer.train(
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
epochs=epochs,
|
||||
batch_size=batch_size,
|
||||
validation_split=0.2,
|
||||
verbose=1
|
||||
)
|
||||
|
||||
# Export to ONNX
|
||||
trainer.export_to_onnx(model_path)
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("Training completed successfully!")
|
||||
print("="*60)
|
||||
print(f"Model saved to: {model_path}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"\nERROR: Training failed: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
finally:
|
||||
if trainer:
|
||||
trainer.cleanup()
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user