quantlab.model.torch_model.TorchModel

class quantlab.model.torch_model.TorchModel(config)

Bases: TrainingTargetMixin, BaseModel

Torch variant: standard PyTorch datasets and loaders, every learning choice a hook.

The base class assembles the collected panel as a TrainingPanel of torch tensors and computes the training target once per fit. The head’s _dataset hook turns it into a PyTorch Dataset and _dataloader batches it. The default is one cross-section per step (ADR 0006): one bar’s present symbols, each with its last window_bars bars, so the network sees [S_t, N, F]; S_t changes from bar to bar, so it must not depend on the order or the number of symbols, and a symbol the model never saw in training still gets a prediction. Every sample carries its where (timestamp and symbol index), through which the base puts predictions back into the panel for any sample shape.

The base class owns every data contract: the warm-up, the training target and its mask, moving batches to the device, the epoch loop, evaluation under no_grad in eval mode, the {split}_loss of each split, .pth checkpoints and prediction; the other metrics come from the model’s evaluation after training (see BaseModel._evaluate). A head writes three things:

window_bars

N, the bars in each symbol’s window.

_init_model(num_features, num_labels, hyperparameters)

The nn.Module; several networks go in an nn.ModuleDict.

_loss(output, batch)

The loss of one Batch from the network’s raw output; count only batch.mask samples.

and may override any of these, each of which has a working default:

_dataset(panel, bars, training)

The PyTorch Dataset over bars. Default: CrossSectionDataset, one item per bar; SymbolSequenceDataset gives Qlib-style (bar, symbol) samples batched to [B, N, F].

_dataloader(dataset, training)

The DataLoader. Default: batch_size and num_workers from the hyperparameters (None and 0, one item per step), shuffled only in training with a generator seeded from random_seed, the last batch never dropped.

_transform_feature(x)

A batch’s raw x, NaN where a value or a bar is missing, to the network’s input. Default: clip to ±3, NaN to 0.

_transform_target(y, training)

One bar’s raw [S_t, L] labels to (target, keep), computed once per fit; keep (or None) removes symbols from the loss only. Default: (y, None).

_init_optim(model)

Default: Adam at hyperparameters["lr"] (1e-3), kept on self.optim; a head may return anything its own _train_one_batch understands, such as a dict of optimizers.

_train_one_batch(epoch, batch)

One optimisation step; returns the loss. Default: forward, _loss, backward, gradient values clipped to grad_clip_value (3.0; None disables), step.

_val_one_batch(epoch, batch)

The evaluation loss of one batch. Default: _loss.

_test_one_batch(epoch, batch)

Called on every test batch after each epoch. Default: nothing.

_forward(x)

The prediction, batch.mask.shape + (L,), from a transformed x, used for metrics and predict_panel. Default: self.model(x); override it when the network returns more than the prediction.

_on_fit_start(), _should_stop(epoch, train_loss, val_loss), _on_fit_end()

When to stop and which weights to keep. Default: run epochs epochs, keep the last weights.

epochs, the epoch cap, is read from the hyperparameters (default 100) and must be a positive integer; see RESERVED_HYPERPARAMETERS for the other keys the base reads.

The model’s warm-up is N - 1 bars: every feature request, in training and in a backtest, starts that many bars earlier, counted on each factor’s own calendar (the bars of its output panel, Factor.bar_before: days for a daily factor computed from one-minute bars, not its dataset’s minutes), and every factor starts at the earliest of those bars, so the first requested bar has a full window.

Examples

A minimal head: a linear map of each symbol’s latest bar:

>>> class LastBarHead(TorchModel):
...     window_bars = 5
...     def _init_model(self, num_features, num_labels, hyperparameters):
...         return LastBarLinear(num_features, num_labels)
...     def _loss(self, output, batch):
...         return masked_mse(output, batch.y, batch.mask)
>>> head = LastBarHead(ModelConfig(
...     factors=[factor], labels=[label], model_save_dir="checkpoints",
...     factor_data_strategy="read", label_data_strategy="read",
...     hyperparameters={"epochs": 2},
...     train_start="2024-01-01", train_end="2024-01-30",
...     test_start="2024-01-31", test_end="2024-02-09",
... ))
>>> head.collect().train().suffix
'.pth'

where LastBarLinear is nn.Linear(num_features, num_labels) applied to x[:, -1].

PANEL_DEVICES = ('auto', 'cuda', 'cpu')

Accepted hyperparameters["panel_device"] values.

PANEL_DTYPES = {'float16': torch.float16, 'float32': torch.float32}

Accepted hyperparameters["panel_dtype"] values and the feature dtype each stores.

PANEL_GPU_BUDGET = 0.5

Largest share of free GPU memory panel_device="auto" gives the panel.

check_hyperparameters()

Refuse an invalid epochs, prefetch, pin_memory, panel_device or panel_dtype.

Raises:

ValueError

If one of them is invalid.

Examples

>>> model.check_hyperparameters()  # the defaults are valid

property device

"cuda" when a CUDA device is available, else "cpu".

Examples

>>> head.device
'cpu'

property epochs

The epoch cap, hyperparameters["epochs"], 100 when unset.

Raises:

ValueError

If the value is not a positive integer (a bool is refused too).

Examples

>>> head.epochs
100

grad_clip_value = 3.0

Gradient value clip of the default _train_one_batch; None disables.

reserved_hyperparameters = frozenset({'batch_size', 'epochs', 'lr', 'num_workers', 'panel_device', 'panel_dtype', 'pin_memory', 'prefetch'})

The hyperparameters keys this variant reads itself, which head_hyperparameters leaves out.

property warmup_bars

window_bars - 1: bars requested before the start of every feature panel.

Examples

>>> head.warmup_bars
4

abstract property window_bars

N, the number of bars in each symbol’s input window.

Examples

>>> head.window_bars
5