dsipts.models.vva.vqvae module

class dsipts.models.vva.vqvae.VectorQuantizer(num_embeddings, embedding_dim, commitment_cost)

Bases: Module

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

__init__(num_embeddings, embedding_dim, commitment_cost)

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

forward(inputs)

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.vva.vqvae.VectorQuantizerEMA(num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5)

Bases: Module

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

__init__(num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5)

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

forward(inputs)

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.vva.vqvae.Residual(in_channels, hidden_channels, num_residual_hiddens)

Bases: Module

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

__init__(in_channels, hidden_channels, num_residual_hiddens)

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.

class dsipts.models.vva.vqvae.ResidualStack(in_channels, hidden_channels, num_residual_layers, num_residual_hiddens)

Bases: Module

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

__init__(in_channels, hidden_channels, num_residual_layers, num_residual_hiddens)

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.

class dsipts.models.vva.vqvae.Encoder(in_channels, hidden_channels, num_residual_layers=3)

Bases: Module

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

__init__(in_channels, hidden_channels, num_residual_layers=3)

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

forward(inputs)

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.vva.vqvae.Jitter(probability=0.12)

Bases: Module

Jitter implementation from [Chorowski et al., 2019]. During training, each latent vector can replace either one or both of its neighbors. As in dropout, this prevents the model from relying on consistency across groups of tokens. Additionally, this regularization also promotes latent representation stability over time: a latent vector extracted at time step t must strive to also be useful at time steps t − 1 or t + 1.

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

__init__(probability=0.12)

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

forward(quantized)

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.vva.vqvae.Decoder(in_channels, hidden_channels, out_channels, num_residual_layers=3)

Bases: Module

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

__init__(in_channels, hidden_channels, out_channels, num_residual_layers=3)

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

forward(x, is_training=True)

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.vva.vqvae.VQVAE(in_channels, hidden_channels, out_channels, num_embeddings, embedding_dim, commitment_cost, decay)

Bases: Module

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

__init__(in_channels, hidden_channels, out_channels, num_embeddings, embedding_dim, commitment_cost, decay)

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

forward(x, is_training=True)

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.