bead.src.models package

Submodules

bead.src.models.flows module

Collection of flow strategies

class bead.src.models.flows.CNN_Flow(dim, cnn_layers, kernel_size, test_mode=0, use_revert=True)

Bases: Module

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 bead.src.models.flows.IAF(z_size, num_flows=2, num_hidden=0, h_size=50, forget_bias=1.0, conv2d=False)

Bases: Module

forward(z, h_context)

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 bead.src.models.flows.NSF_AR(dim=15, K=64, B=3, hidden_dim=8, base_network=<class 'bead.src.models.layers.FCNN'>)

Bases: Module

Neural spline flow, auto-regressive. [Durkan et al. 2019]

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.

reset_parameters()

class bead.src.models.flows.Planar

Bases: Module

der_h(x)

Derivative of tanh

forward(zk, u, w, b)

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 bead.src.models.flows.Sylvester(num_ortho_vecs)

Bases: Module

Sylvester normalizing flow.

der_h(x)

der_tanh(x)

forward(zk, r1, r2, q_ortho, b, sum_ldj=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 bead.src.models.flows.TriangularSylvester(z_size)

Bases: Module

Sylvester normalizing flow with Q=P or Q=I.

der_h(x)

der_tanh(x)

forward(zk, r1, r2, q_ortho, b, sum_ldj=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.

bead.src.models.layers module

class bead.src.models.layers.CNN_Flow_Layer(dim, kernel_size, dilation, test_mode=0, rescale=True, skip=True)

Bases: Module

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 bead.src.models.layers.Dilation_Block(dim, kernel_size, test_mode=0)

Bases: Module

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 bead.src.models.layers.FCNN(in_dim, out_dim, hidden_dim)

Bases: Module

Simple fully connected neural network.

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 bead.src.models.layers.GraphConvolution(in_features, out_features, bias=True)

Bases: Module

Simple GCN layer, similar to https://arxiv.org/abs/1609.02907

forward(input, adj)

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.

reset_parameters()

class bead.src.models.layers.Identity

Bases: Module

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 bead.src.models.layers.MaskedConv2d(in_features, out_features, size_kernel=(3, 3), diagonal_zeros=False, bias=True)

Bases: Module

build_mask()

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.

reset_parameters()

class bead.src.models.layers.MaskedLinear(in_features, out_features, diagonal_zeros=False, bias=True)

Bases: Module

build_mask()

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.

reset_parameters()

bead.src.models.layers.RQS(inputs, unnormalized_widths, unnormalized_heights, unnormalized_derivatives, inverse=False, left=0.0, right=1.0, bottom=0.0, top=1.0, min_bin_width=0.001, min_bin_height=0.001, min_derivative=0.001)

bead.src.models.layers.searchsorted(bin_locations, inputs, eps=1e-06)

bead.src.models.layers.unconstrained_RQS(inputs, unnormalized_widths, unnormalized_heights, unnormalized_derivatives, inverse=False, tail_bound=1.0, min_bin_width=0.001, min_bin_height=0.001, min_derivative=0.001)

bead.src.models.models module

class bead.src.models.models.AE(in_shape, z_dim, *args, **kwargs)

Bases: Module

decode(z)

detach_hooks(hooks: list) → None

encode(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.

get_activations() → dict

get_hook(layer_name)

get_layers() → list

store_hooks() → list

class bead.src.models.models.AE_Dropout_BN(in_shape, z_dim, *args, **kwargs)

Bases: AE

dec_bn(z)

enc_bn(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.

class bead.src.models.models.ConvAE(in_shape, z_dim, *args, **kwargs)

Bases: Module

decode(z)

encode(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.

class bead.src.models.models.ConvFlow_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with convolutional flows in the decoder.

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 bead.src.models.models.ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvAE

decode(z)

encode(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.

reparameterize(mean, logvar)

class bead.src.models.models.HouseholderSylvester_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with householder sylvester flows in the decoder.

batch_construct_orthogonal(q)

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 bead.src.models.models.IAF_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with inverse autoregressive flows in the decoder.

encode(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.

class bead.src.models.models.NSFAR_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with auto-regressive neural spline flows in the decoder.

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 bead.src.models.models.OrthogonalSylvester_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with orthogonal flows in the decoder.

batch_construct_orthogonal(q)

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 bead.src.models.models.Planar_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with planar flows in the decoder.

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 bead.src.models.models.TransformerAE(in_dim, h_dim=256, n_heads=1, latent_size=50, activation=<built-in function gelu>)

Bases: Module

Autoencoder mixed with the Transformer Encoder layer

Parameters:

nn (_type_) – _description_

decoder(z: Tensor)

_summary_

Parameters:

z (_type_) – _description_

Returns:

_description_

Return type:

_type_

encoder(x: Tensor)

_summary_

Parameters:

x (_type_) – _description_

Returns:

_description_

Return type:

_type_

forward(x: Tensor)

_summary_

Parameters:

z (_type_) – _description_

Returns:

_description_

Return type:

_type_

class bead.src.models.models.TriangularSylvester_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with triangular sylvester flows in the decoder.

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