dsipts.models.duet package

Submodules

dsipts.models.duet.layers module

class dsipts.models.duet.layers.Linear_extractor(seq_len, d_model, enc_in, CI, moving_avg, individual=False)

Bases: Module

individual: Bool, whether shared model among different variates.

encoder(x)

forecast(x_enc)

forward(x_enc)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.duet.layers.Linear_extractor_cluster(noisy_gating, num_experts, seq_len, k, d_model, enc_in, CI, moving_avg, hidden_size)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

cv_squared(x)

The squared coefficient of variation of a sample. Useful as a loss to encourage a positive distribution to be more uniform. Epsilons added for numerical stability. Returns 0 for an empty Tensor. Args: x: a Tensor. Returns: a Scalar.

forward(x, loss_coef=1)

Args: x: tensor shape [batch_size, input_size] train: a boolean scalar. loss_coef: a scalar - multiplier on load-balancing losses

Returns: y: a tensor with shape [batch_size, output_size]. extra_training_loss: a scalar. This should be added into the overall training loss of the model. The backpropagation of this loss encourages all experts to be approximately equally used across a batch.

noisy_top_k_gating(x, train, noise_epsilon=0.01)

Noisy top-k gating. See paper: https://arxiv.org/abs/1701.06538. :param x: input Tensor with shape [batch_size, input_size] :param train: a boolean - we only add noise at training time. :param noise_epsilon: a float

Returns:

a Tensor with shape [batch_size, num_experts] load: a Tensor with shape [num_experts]

Return type:

gates

class dsipts.models.duet.layers.RevIN(num_features: int, eps=1e-05, affine=True, subtract_last=False)

Bases: Module

Parameters:

• num_features – the number of features or channels

• eps – a value added for numerical stability

• affine – if True, RevIN has learnable affine parameters

forward(x, mode: str)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.duet.layers.SparseDispatcher(num_experts, gates)

Bases: object

Create a SparseDispatcher.

combine(expert_out, multiply_by_gates=True)

Sum together the expert output, weighted by the gates. The slice corresponding to a particular batch element b is computed as the sum over all experts i of the expert output, weighted by the corresponding gate values. If multiply_by_gates is set to False, the gate values are ignored. :param expert_out: a list of num_experts `Tensor`s, each with shape

[expert_batch_size_i, <extra_output_dims>].

Parameters:

multiply_by_gates – a boolean

Returns:

a Tensor with shape [batch_size, <extra_output_dims>].

dispatch(inp)

Create one input Tensor for each expert. The Tensor for an expert i contains the slices of inp corresponding to the batch elements b where gates[b, i] > 0. :param inp: a Tensor of shape “[batch_size, <extra_input_dims>]`

Returns:

a list of num_experts `Tensor`s with shapes

[expert_batch_size_i, <extra_input_dims>].

expert_to_gates()

Gate values corresponding to the examples in the per-expert `Tensor`s. :returns:

a list of num_experts one-dimensional Tensor`s with type `tf.float32

and shapes [expert_batch_size_i]

class dsipts.models.duet.layers.encoder(seq_len, num_experts, hidden_size)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dsipts.models.duet.masked module

class dsipts.models.duet.masked.AttentionLayer(attention, d_model, n_heads, d_keys=None, d_values=None)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(queries, keys, values, attn_mask, tau=None, delta=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.duet.masked.Encoder(attn_layers, conv_layers=None, norm_layer=None)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, attn_mask=None, tau=None, delta=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.duet.masked.EncoderLayer(attention, d_model, d_ff=None, dropout=0.1, activation='relu')

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, attn_mask=None, tau=None, delta=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.duet.masked.FullAttention(mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(queries, keys, values, attn_mask, tau=None, delta=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.duet.masked.Mahalanobis_mask(input_size)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

bernoulli_gumbel_rsample(distribution_matrix)

calculate_prob_distance(X)

forward(X)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Module contents