nn package

nn.tmGlow module

class nn.tmGlow.Encoder(in_features, enc_block_layers, growth_rate=4, init_features=48, output_features=8, cond_features=8, cglow_upscale=1, bn_size=8, drop_rate=0.0, bottleneck=False)

Bases: torch.nn.modules.module.Module

Convolutional dense encoder for conditioning the generative model. See Fig. 4 of paper: https://arxiv.org/abs/2006.04731

Parameters

  • in_features (int) – Number of input feature channels

  • enc_block_layers (list) – list of the number of dense layers in each desired dense block

  • growth_rate (int, optional) – Growth rate of channels in dense block, defaults to 4

  • init_features (int, optional) – Number of channels after te first convolution, defaults to 48

  • output_features (int, optional) – Number of output channel features (will be twice this value since final value will be turned into mu and stds for latent variables), defaults to 8

  • cond_features (int, optional) – Number of conditional feature channels, defaults to 8

  • cglow_upscale (int, optional) – The factor to upscale the features from the dense encode before passing them to the generative model, defaults to 1

  • bn_size (int, optional) – number of features after bottleneck if enabled, defaults to 8

  • drop_rate (float, optional) – Dropout rate, defaults to 0.

  • bottleneck (bool, optional) – Enable bottle next to reduce the number of feature channels, defaults to False

enconding_transition(in_features, out_features, drop_rate=0, padding=1)

Encoding transition convolution placed between dense blocks, halves feature map size

Parameters

  • in_features (int) – Number of input feature channels

  • output_features (int, optional) – Number of output channels

  • drop_rate (float, optional) – Dropout rate, defaults to 0.

  • padding (int, optional) – convolutional padding, defaults to 1

Returns

  • encoder_trans: Transition PyTorch module

Return type

nn.Module

first_encoding(in_features, init_features, drop_rate=0.0)

First transition encodering block, halves feature map size

Parameters

  • in_features (int) – Number of input feature channels

  • init_features (int, optional) – Number of channels after te first convolution, defaults to 48

  • drop_rate (float, optional) – Dropout rate, defaults to 0.

Returns

  • first_encoder: Encoding PyTorch module

Return type

nn.Module

forward(x)

Encoder forward pass

Parameters

x (torch.Tensor) – [B, in_features, H, W] input feature tensor

Returns

  • out: Final encoder output

  • c_out: List of conditional output states to conditional the generative model

Return type

(torch.Tensor, list)

training: bool
class nn.tmGlow.LSTMCFlowDecoder(in_features, glow_block_layers, cond_features=8, rec_features=8, squeeze_factor=2, conv_ksize=3, LUdecompose=False, train_sampling=True, squeeze_type=0)

Bases: torch.nn.modules.module.Module

LSTM Conditional generative normalizing flow decoder model. See Fig. 4 of paper: https://arxiv.org/abs/2006.04731

Parameters

  • in_features (int) – Number of input feature channels

  • glow_block_layers (list) – list of the number of affine layers in each flow block

  • cond_features (int, optional) – Number of conditional feature channels, defaults to 8

  • rec_features (int, optional) – Number of recurrent feature channels in LSTM, defaults to 8

  • squeeze_factor (int, optional) – Squeeze factor to reduce the dimensionality of the feature map, defaults to 2

  • LUdecompose (bool, optional) – Use the LU decomposition for 1x1 convolution, defaults to False

  • train_sampling (bool, optional) – Train by sampling (1x1 conv. is inverted for the forward pass), defaults to True

  • squeeze_type (int, optional) – Type of squeeze method, see nn.modules.flowLSTMBlock.LSTMFLowBlock for more details, defaults to 0

forward(x, c_in, h_in, return_eps=False)

Forward pass flow model \(x \rightarrow z\)

Parameters

  • x (torch.Tensor) – [B, in_features, H, W] input feature tensor

  • c_in (list) – [len(flow_blocks)] List of conditional features for each flow block

  • h_in (list) – [len(flow_blocks)] List of recurrent features for each flow block

  • return_eps (bool, optional) – Return samples from latent densities, defaults to False

Returns

  • z: Final encoded latent variables

  • log_det: [B] log determinate + the log-likelihood of the latent variables

  • s_out: [len(flow_blocks)] List of output recurrent features from each flow block

  • eps: list of sampled latent variables from each split

Return type

(torch.Tensor, torch.Tensor, list, list)

reverse(z, c_in, h_in, eps)

Backward pass flow model \(x \leftarrow z\)

Parameters

  • z – input latent tensor

  • c_in (list) – [len(flow_blocks)] List of conditional features for each flow block

  • h_in (list) – [len(flow_blocks)] List of recurrent features for each flow block

  • eps (torch.Tensor, optional) – list sampled latent variables, defaults to None

Returns

  • x: [B, in_features, H, W] input feature tensor

  • log_det: [B] log determinate + the log-likelihood of the latent variables

  • s_out: [len(flow_blocks)] List of output recurrent features from each flow block

Return type

(torch.Tensor, torch.Tensor, list)

training: bool
class nn.tmGlow.TMGlow(in_features, out_features, enc_blocks, glow_blocks, cond_features=8, cglow_upscale=1, growth_rate=4, init_features=48, rec_features=8, bn_size=8, drop_rate=0, bottleneck=False)

Bases: torch.nn.modules.module.Module

Transient Multi-Fidelity Glow Model. Consists of a dense encoder model: nn.tmGlow.Encoder and a generative flow decoder: nn.tmGlow.LSTMCFlowDecoder. See Fig. 4 of paper: https://arxiv.org/abs/2006.04731

Parameters

  • in_features (int) – Number of input feature channels

  • out_features (int) – Number of output feature channels

  • enc_blocks (list) – list of the number of dense layers in each desired dense block

  • glow_blocks (list) – list of the number of affine layers in each flow block

  • cond_features (int, optional) – Number of conditional feature channels, defaults to 8

  • cglow_upscale (int, optional) – The factor to upscale the features from the dense encode before passing them to the generative model, defaults to 1

  • growth_rate (int, optional) – Growth rate of channels in dense block, defaults to 4

  • init_features (int, optional) – Number of channels after te first convolution, defaults to 48

  • rec_features (int, optional) – Number of recurrent feature channels in LSTM, defaults to 8

  • bn_size (int, optional) – number of features after bottleneck if enabled, defaults to 8

  • drop_rate (float, optional) – Dropout rate, defaults to 0.

  • bottleneck (bool, optional) – Enable bottle next to reduce the number of feature channels, defaults to False

forward(x, y, h_in=None, return_eps=False)

Forward pass of the encoder and flow model

Parameters

  • x (torch.Tensor) – [B, in_features, H, W] Input feature tensor

  • y (torch.Tensor) – [B, out_features, H2, W2] Output feature tensor

  • h_in (list, optional) – [len(glow_blocks)] List of recurrent features for each flow block, defaults to None

  • return_eps (bool, optional) – Return samples from latent densities, defaults to False

Returns

  • z: [B, C, H, W] latent feature tensor

  • log_det: [B] log determinate + the log-likelihood of the latent variables

  • h_out: [len(flow_blocks)] List of output recurrent features from each flow block

  • eps: list of sampled latent variables from each split

Return type

(torch.Tensor, torch.Tensor, list, list)

initLSTMStates(seeds, input_dim)

Intialized the LSTM states using random generators with the specified seeds. Generates the LSTM hidden states by a U[-1,1] and cell states by N(0,1).

Parameters

  • seeds (int) – Random generator seed

  • input_dim (torch.size) – input_feature dimensions

Returns

  • a_states: list of initialized recurrent states

Return type

list

reconstruct(x, h_in, eps)

Backward pass; Reconstruct specific y

Parameters

  • x (torch.Tensor) – [B, in_features, H, W] Input feature tensor

  • h_in (list, optional) – [len(glow_blocks)] List of recurrent features for each flow block, defaults to None

  • eps (torch.Tensor, optional) – list sampled latent variables, defaults to None

Returns

  • z: [B, C, H, W] latent feature tensor

  • log_det: [B] log determinate + the log-likelihood of the latent variables

  • h_out: [len(flow_blocks)] List of output recurrent features from each flow block

Return type

(torch.Tensor, torch.Tensor, list)

sample(x, h_in=None)

Backward pass; Conditional generation!

Parameters

  • x (torch.Tensor) – [B, in_features, H, W] Input feature tensor

  • h_in (list, optional) – [len(glow_blocks)] List of recurrent features for each flow block, defaults to None

Returns

  • z: [B, C, H, W] latent feature tensor

  • log_det: [B] log determinate + the log-likelihood of the latent variables

  • h_out: [len(flow_blocks)] List of output recurrent features from each flow block

Return type

(torch.Tensor, torch.Tensor, list)

training: bool

nn.trainFlowParallel module

class nn.trainFlowParallel.TMGLowLoss(args, model, log=None)

Bases: torch.nn.modules.module.Module

A torch.nn module for executing loss calculations on multiple GPUs. Args:

args (argparse): object with programs arguments model (torch.nn): Glow-TM model

calcVDiv(yPred)

Calculates the divergence residual term. Args:

yPred (Tensor): [b x 3 x d1 x d2]

calcVPres(yPred)

Calculates the pressure residual term. Args:

yPred (Tensor): [b x 3 x d1 x d2]

forward(yPred, logp, target, target_mean, target_rms)

The forward of this class is called through the parallel wrapper Thus this supports the extention of the loss calculation to different GPUs Args:

yPred (torch.Tensor): [batch x time-steps x 3 x nx x ny] Glow-TM prediction logp (torch.Tensor): The log probability of the likelihood (entropy term) yPred (torch.Tensor): [batch x time-steps x 3 x nx x ny] Glow-TM prediction target_mean (torch.Tensor): [batch x 3 x nx x ny] mean flow field to use in RMS predictions target_rms (torch.Tensor): [batch x 3 x nx x ny] target RMS value of flucation states

training: bool
class nn.trainFlowParallel.TMGLowPredictionItem(yPred: Optional[tensorList] = None, yTarget: Optional[tensorList] = None, logp: Optional[tensorList] = None, tback: int = 1)

Bases: object

Class used to store model predictions overtime

add(yPred0, logp0, yTarget0)

Adds provided prediction and target tensors to class Time-steps are stored in dimension 1

clear()

Clears stored prediction and target variables. Should be called after backprop

concat(tensor1, tensor2, dim=1)

Concats the new tensor 2 onto existing series tensor 1

getOutputs()

Combines yPred and logp into a single tuple for loss evaluation

getTargets(*newTargets)

Returns target tensors and allows for additional target tensors to be added into a tuple

logp: Optional[tensorList] = None
tback: int = 1
unsqueeze(tensor, dim=1)

Unsqueezes the tensor or list of tensors in the specified dimension

yPred: Optional[tensorList] = None
yTarget: Optional[tensorList] = None
class nn.trainFlowParallel.TrainFlow(args, model, train_loader, test_loader, log=None)

Bases: object

Trains recursive encoder attention-driven decoder Args:

args (argparse): object with programs arguments model (torch.nn): Glow-TM model train_loader (torch.dataloader): dataloader with training cases test_loader (torch.dataloader): dataloader with training cases log (Log): class for logging console outputs

test(model, samples=1, epoch=0, plot=True)

Tests the model Args:

model (torch.nn.Module): PyTorch model to test stride (int): The stride the low-fidelity input takes compared to output samples (int): Number of prediction to sample from the model epoch (int): current epoch plot (boolean): If to plot two of the predictions or not

Returns:

mse (float): mean-squared-error between the predictive mean and targets

trainParallel(model, optimizer, tback=1, epoch=0, **kwargs)

Trains the model for a single epoch Args:

model (torch.nn.Module): PyTorch model to train optimizer (torch.optim): PyTorch optimizer to update the models parameters tback (int): number of time-steps to back propagate through in time stride (int): The stride the low-fidelity input takes compared to output epoch (int): current epoch

Returns:

total_loss (float): current loss