src.training.torch_model

🦉 AthenaAI is an AI-powered ensemble machine-learning trading bot for PocketOption, built on BinaryOptionsToolsV2 → Get it on GitLab
💬 Need help? Join the Chipa Discord · ⚡ Go no-code with ChipaEditor

Source: src/training/torch_model.py. TorchMLPClassifier is a binary classifier (input → hidden ReLU+dropout layers → 2 logits) with an sklearn-like fit/predict/predict_proba API so it can sit inside EnsemblePredictor._batch_models next to GBM and RF. Its defining design goals:

• Picklable across machines: weights are stored as a CPU state dict; the live torch.nn network is never pickled and is rebuilt lazily on first predict. A brain trained on a GPU machine loads fine on a CPU-only machine.
• Graceful degradation: if PyTorch is missing at load/predict time, predict_proba returns a neutral 0.5 vote for every sample instead of breaking the whole brain.

PyTorch is imported lazily inside methods (_torch() helper) — importing this module does not require torch.

class TorchMLPClassifier

__init__(self, hidden=(128, 64), epochs=30, batch_size=1024, lr=1e-3, weight_decay=1e-5, dropout=0.2, device="cpu")

State: n_features_ = None, _state_dict = None (CPU weights, picklable), _net = None (live network, never pickled).

fit(self, X: np.ndarray, y: np.ndarray) -> self

Trains the network. X should already be scaler-transformed (float32 tensors are built from it); y holds int labels 0/1. Uses AdamW, CrossEntropyLoss, and a shuffled DataLoader. Logs on the first epoch and every 5th:

   MLP epoch 5/30  loss=0.6821  acc=55.3%

After training it switches to eval mode, snapshots self._state_dict = {k: v.detach().cpu() ...}, and keeps the live net for immediate use. Raises ImportError if torch is absent (the trainer checks availability first).

predict_proba(self, X: np.ndarray) -> np.ndarray

Returns an (n, 2) softmax probability array (columns: P(class 0 = PUT-win), P(class 1 = CALL-win)). If _ensure_net() fails — no state dict yet, or torch not installed — returns np.full((len(X), 2), 0.5): a neutral vote that leaves the rest of the ensemble unaffected. Runs under torch.no_grad() on whatever device the rebuilt net lives on.

predict(self, X: np.ndarray) -> np.ndarray

predict_proba(X).argmax(axis=1) — hard 0/1 labels.

_ensure_net(self) -> bool (internal)

Lazy rebuild: returns True if a live net exists; returns False if there are no saved weights or torch can’t be imported. Otherwise rebuilds the architecture from hidden and n_features_, loads the CPU state dict, and moves it to whatever this machine has ("cuda" if available, else "cpu" — note MPS is not chosen here), updates self.device, sets eval mode.

Pickling — __getstate__ / __setstate__

__getstate__ copies __dict__ with _net forced to None (only the CPU _state_dict and hyperparameters are serialized). __setstate__ restores the dict and leaves _net = None for lazy rebuild. This is why a brain pickle containing an MLP is portable and safe to load without a GPU.

Usage example

import numpy as np
from src.training.torch_model import TorchMLPClassifier
from src.training.device import detect_device

device, _ = detect_device()
mlp = TorchMLPClassifier(hidden=(128, 64), epochs=30,
                         batch_size=2048, lr=1e-3, device=device)
mlp.fit(X_scaled_train, y_train)
proba = mlp.predict_proba(X_scaled_test)   # (n, 2)

In the pipeline, Trainer._train_nn() fits it on scaler-transformed data and inserts it into the ensemble as batch model "MLP" with its holdout accuracy as vote weight (weighted ×0.75 at prediction time like all batch models — see models.md).

See also

• GPU guide, device.md, trainer.md

💬 Unsure how this interacts with the rest of the configuration? Ask in the Chipa Discord — or prototype strategies no-code with ChipaEditor.