from __future__ import annotations import argparse import json import math import sys from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) try: import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset except ImportError as exc: # pragma: no cover - exercised on machines without training deps. raise SystemExit( "PyTorch is not installed. Install local training deps with: " "python -m pip install torch --index-url https://download.pytorch.org/whl/cpu" ) from exc from crypto_spot_bot.bybit import BybitClient from crypto_spot_bot.config import load_settings from crypto_spot_bot.time_series import _log_returns @dataclass(slots=True) class PreparedData: train_x: torch.Tensor train_y: torch.Tensor validation_x: torch.Tensor validation_y: torch.Tensor validation_returns: list[float] mean: float scale: float train_samples: int validation_samples: int class RecurrentReturnModel(nn.Module): def __init__( self, *, architecture: str, hidden_size: int, num_layers: int, dropout: float, ) -> None: super().__init__() recurrent_cls = nn.LSTM if architecture == "lstm" else nn.GRU self.rnn = recurrent_cls( input_size=1, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout if num_layers > 1 else 0.0, batch_first=True, ) self.head = nn.Linear(hidden_size, 1) def forward(self, values: torch.Tensor) -> torch.Tensor: output, _state = self.rnn(values) return self.head(output[:, -1, :]).squeeze(-1) def main() -> None: args = _parse_args() if args.threads > 0: torch.set_num_threads(args.threads) _seed(args.seed) settings = load_settings(args.env) client = BybitClient(settings) symbols = _symbols(args.symbols, settings, client) interval = args.interval or settings.base_interval output = Path(args.output) if args.output else settings.time_series_lstm_model_path device = _device(args.device) artifact: dict[str, Any] = { "version": 2, "type": "pytorch_recurrent_forecaster", "created_at": datetime.now(timezone.utc).isoformat(), "trainer": Path(__file__).name, "interval": interval, "limit": args.limit, "validation_window": args.validation_window, "device": str(device), "symbols": {}, } for symbol in symbols: result = _train_symbol( client=client, symbol=symbol, interval=interval, limit=args.limit, validation_window=args.validation_window, architectures=_strings(args.architectures), lookbacks=_ints(args.lookbacks), hidden_sizes=_ints(args.hidden_sizes), layers_values=_ints(args.layers), dropouts=_floats(args.dropouts), epochs=args.epochs, patience=args.patience, batch_size=args.batch_size, learning_rate=args.learning_rate, weight_decay=args.weight_decay, clip=args.clip, device=device, seed=args.seed, ) if result is None: print(f"{symbol}: skipped, not enough candles or train/validation samples") continue artifact["symbols"][symbol] = result print( f"{symbol}: model={result['model']} lookback={result['lookback']} " f"hidden={result['hidden_size']} layers={result['num_layers']} " f"mae={result['validation_mae_percent']:.5f}% " f"baseline={result['baseline_mae_percent']:.5f}% skill={result['skill']:.4f}" ) output.parent.mkdir(parents=True, exist_ok=True) tmp_output = output.with_name(f"{output.name}.tmp") tmp_output.write_text(json.dumps(artifact, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") tmp_output.replace(output) print(f"saved {output}") def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train PyTorch LSTM/GRU forecast models on Bybit spot candles.") parser.add_argument("--env", default=None, help="Path to .env file.") parser.add_argument("--symbols", default="", help="Comma-separated symbols. Defaults to configured or popular pairs.") parser.add_argument("--interval", default="", help="Bybit kline interval. Defaults to BASE_INTERVAL.") parser.add_argument("--limit", type=int, default=1000, help="Kline limit per symbol.") parser.add_argument("--validation-window", type=int, default=120, help="Held-out tail returns used for validation.") parser.add_argument("--architectures", default="lstm,gru", help="Comma-separated recurrent types: lstm,gru.") parser.add_argument("--lookbacks", default="32,64", help="Comma-separated sequence lengths.") parser.add_argument("--hidden-sizes", default="16,32", help="Comma-separated hidden sizes.") parser.add_argument("--layers", default="1", help="Comma-separated recurrent layer counts.") parser.add_argument("--dropouts", default="0.0", help="Comma-separated dropout values; only used with layers > 1.") parser.add_argument("--epochs", type=int, default=60, help="Maximum epochs per hyperparameter candidate.") parser.add_argument("--patience", type=int, default=10, help="Early stopping patience in epochs.") parser.add_argument("--batch-size", type=int, default=64, help="Training batch size.") parser.add_argument("--learning-rate", type=float, default=0.001, help="AdamW learning rate.") parser.add_argument("--weight-decay", type=float, default=0.0001, help="AdamW weight decay.") parser.add_argument("--clip", type=float, default=8.0, help="Clamp normalized returns and predictions to this range.") parser.add_argument("--seed", type=int, default=7, help="Random seed.") parser.add_argument("--threads", type=int, default=0, help="Torch CPU threads; 0 keeps torch default.") parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or mps.") parser.add_argument("--output", default="", help="Output JSON path. Defaults to TIME_SERIES_LSTM_MODEL_PATH.") return parser.parse_args() def _symbols(raw: str, settings: Any, client: BybitClient) -> list[str]: if raw.strip(): return [item.strip().upper() for item in raw.split(",") if item.strip()] if settings.symbols: return list(settings.symbols) return client.popular_spot_symbols(settings.top_symbols_count) def _train_symbol( *, client: BybitClient, symbol: str, interval: str, limit: int, validation_window: int, architectures: list[str], lookbacks: list[int], hidden_sizes: list[int], layers_values: list[int], dropouts: list[float], epochs: int, patience: int, batch_size: int, learning_rate: float, weight_decay: float, clip: float, device: torch.device, seed: int, ) -> dict[str, Any] | None: candles = client.klines(symbol, interval, limit) closes = [float(candle.close) for candle in candles if candle.close > 0] returns = _log_returns(closes) if len(returns) < max(100, validation_window + 80): return None best: dict[str, Any] | None = None for lookback in lookbacks: prepared = _prepare_data(returns, lookback, validation_window, clip, device) if prepared is None: continue baseline_mae = sum(abs(value) for value in prepared.validation_returns) / len(prepared.validation_returns) for architecture in architectures: if architecture not in {"lstm", "gru"}: continue for hidden_size in hidden_sizes: for num_layers in layers_values: for dropout in dropouts: candidate = _fit_candidate( prepared=prepared, architecture=architecture, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout, epochs=epochs, patience=patience, batch_size=batch_size, learning_rate=learning_rate, weight_decay=weight_decay, clip=clip, device=device, seed=seed, ) validation_mae = float(candidate["validation_mae"]) skill = (baseline_mae - validation_mae) / baseline_mae if baseline_mae > 0 else 0.0 row = { **candidate, "model": f"torch_{architecture}", "architecture": architecture, "lookback": lookback, "hidden_size": hidden_size, "num_layers": num_layers, "dropout": dropout, "mean": prepared.mean, "scale": prepared.scale, "clip": clip, "validation_mae_percent": validation_mae * 100, "baseline_mae_percent": baseline_mae * 100, "skill": skill, "candles": len(candles), "returns": len(returns), "train_samples": prepared.train_samples, "validation_samples": prepared.validation_samples, } if best is None or validation_mae < float(best["validation_mae"]): best = row if best is None: return None best.pop("validation_mae", None) return best def _prepare_data( returns: list[float], lookback: int, validation_window: int, clip: float, device: torch.device, ) -> PreparedData | None: validation_window = min(max(16, validation_window), max(16, len(returns) // 3)) split = len(returns) - validation_window if split <= lookback + 16: return None train_returns = returns[:split] mean = sum(train_returns) / len(train_returns) scale = _return_scale(train_returns) normalized = [_clamp((value - mean) / scale, -clip, clip) for value in returns] train_x: list[list[list[float]]] = [] train_y: list[float] = [] validation_x: list[list[list[float]]] = [] validation_y: list[float] = [] validation_returns: list[float] = [] for target_index in range(lookback, len(returns)): row = [[value] for value in normalized[target_index - lookback : target_index]] target = normalized[target_index] if target_index < split: train_x.append(row) train_y.append(target) else: validation_x.append(row) validation_y.append(target) validation_returns.append(returns[target_index]) if len(train_x) < 24 or len(validation_x) < 8: return None return PreparedData( train_x=torch.tensor(train_x, dtype=torch.float32, device=device), train_y=torch.tensor(train_y, dtype=torch.float32, device=device), validation_x=torch.tensor(validation_x, dtype=torch.float32, device=device), validation_y=torch.tensor(validation_y, dtype=torch.float32, device=device), validation_returns=validation_returns, mean=mean, scale=scale, train_samples=len(train_x), validation_samples=len(validation_x), ) def _fit_candidate( *, prepared: PreparedData, architecture: str, hidden_size: int, num_layers: int, dropout: float, epochs: int, patience: int, batch_size: int, learning_rate: float, weight_decay: float, clip: float, device: torch.device, seed: int, ) -> dict[str, Any]: _seed(seed) model = RecurrentReturnModel( architecture=architecture, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout, ).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay) criterion = nn.SmoothL1Loss(beta=0.5) generator = torch.Generator(device="cpu").manual_seed(seed) loader = DataLoader( TensorDataset(prepared.train_x, prepared.train_y), batch_size=max(1, batch_size), shuffle=True, generator=generator, ) best_state: dict[str, torch.Tensor] | None = None best_mae = math.inf best_epoch = 0 stale_epochs = 0 for epoch in range(1, max(1, epochs) + 1): model.train() for batch_x, batch_y in loader: optimizer.zero_grad(set_to_none=True) loss = criterion(model(batch_x), batch_y) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() validation_mae = _validation_mae(model, prepared, clip) if validation_mae + 1e-12 < best_mae: best_mae = validation_mae best_epoch = epoch best_state = {key: value.detach().cpu().clone() for key, value in model.state_dict().items()} stale_epochs = 0 else: stale_epochs += 1 if stale_epochs >= max(1, patience): break if best_state: model.load_state_dict(best_state) return { "validation_mae": best_mae, "best_epoch": best_epoch, "epochs_trained": best_epoch + stale_epochs, "state_dict": _export_recurrent_state(model), "head_weight": _round_list(model.head.weight.detach().cpu().squeeze(0).tolist()), "head_bias": round(float(model.head.bias.detach().cpu().item()), 10), } def _validation_mae(model: nn.Module, prepared: PreparedData, clip: float) -> float: model.eval() with torch.no_grad(): normalized_predictions = model(prepared.validation_x).detach().cpu().tolist() errors = [] for prediction, actual in zip(normalized_predictions, prepared.validation_returns): raw_prediction = _clamp(float(prediction), -clip, clip) * prepared.scale + prepared.mean errors.append(abs(raw_prediction - actual)) return sum(errors) / len(errors) if errors else math.inf def _export_recurrent_state(model: RecurrentReturnModel) -> dict[str, Any]: return { key: _round_nested(value.detach().cpu().tolist()) for key, value in model.rnn.state_dict().items() } def _device(raw: str) -> torch.device: value = raw.strip().lower() if value == "auto": if torch.cuda.is_available(): return torch.device("cuda") if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") return torch.device(value) def _seed(seed: int) -> None: torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def _return_scale(returns: list[float]) -> float: values = sorted(abs(value) for value in returns if math.isfinite(value)) if not values: return 0.0005 median = values[len(values) // 2] mean = sum(values) / len(values) return max(max(median, mean * 0.5), 1e-5) def _clamp(value: float, low: float, high: float) -> float: return max(low, min(high, value)) def _round_nested(value: Any) -> Any: if isinstance(value, list): return [_round_nested(item) for item in value] return round(float(value), 10) def _round_list(values: list[float]) -> list[float]: return [round(float(value), 10) for value in values] def _ints(raw: str) -> list[int]: return [int(item.strip()) for item in raw.split(",") if item.strip()] def _floats(raw: str) -> list[float]: return [float(item.strip()) for item in raw.split(",") if item.strip()] def _strings(raw: str) -> list[str]: return [item.strip().lower() for item in raw.split(",") if item.strip()] if __name__ == "__main__": main()