quantlab.model.predefined.gats.GATsNet

class quantlab.model.predefined.gats.GATsNet(num_features, num_labels, hidden_size=64, num_layers=2, dropout=0.0, base_model='LSTM')

Bases: Module

Qlib’s GATModel with an output of any width.

Each symbol’s window [N, F] goes through an LSTM or GRU, and the hidden state of the last bar h ([S_t, H]) is kept. With x = W h + b (transformation), the attention score of symbol i on symbol j is LeakyReLU(a[:H]·x_j + a[H:]·x_i) with slope 0.01, normalised by a softmax over every j of the bar, self included. The output is fc_out(LeakyReLU(fc(h + A h))): the untransformed h is aggregated, as in Qlib. The parameter names are Qlib’s, so a Qlib GATModel state dict loads into a one-label net.

The scores are computed as a sum of two [S_t, 1] projections rather than Qlib’s [S_t, S_t, 2H] concatenation, which is the same value without the memory that grows with the square of the cross-section.

Parameters:

num_featuresint

Features per bar, F.

num_labelsint

Outputs per symbol, L; Qlib’s model has one.

hidden_sizeint, default 64

H, the encoder’s hidden size.

num_layersint, default 2

Encoder layers.

dropoutfloat, default 0.0

Dropout between encoder layers.

base_model{“LSTM”, “GRU”}, default “LSTM”

The encoder.

Raises:

ValueError

If base_model is neither "LSTM" nor "GRU".

Examples

>>> net = GATsNet(num_features=4, num_labels=2, hidden_size=8)
>>> net(torch.randn(5, 3, 4)).shape
torch.Size([5, 2])

attention(hidden)

Return the [S_t, S_t] attention weights, each row summing to 1.

Examples

>>> weights = net.attention(torch.randn(5, 8))
>>> weights.shape, bool(torch.allclose(weights.sum(dim=1), torch.ones(5)))
(torch.Size([5, 5]), True)

forward(x)

Map one bar’s [S_t, N, F] windows to [S_t, L] outputs.

Examples

>>> net.forward(torch.randn(3, 6, 4)).shape
torch.Size([3, 2])