simulation.utils.machine_learning.models package

Subpackages

• simulation.utils.machine_learning.models.test package
  • Submodules
  • simulation.utils.machine_learning.models.test.test_helper module
    • test_get_norm_layer()
    • test_get_scheduler()
    • test_init_weights()
    • test_set_requires_grad()
    • main()
  • simulation.utils.machine_learning.models.test.test_resnet_block module
    • test_creating_resnet_block()
    • main()
  • simulation.utils.machine_learning.models.test.test_resnet_generator module
    • test_resnet_generator()
    • main()
  • simulation.utils.machine_learning.models.test.test_wgan_critic module
    • test_weight_clipping()
    • test_forward()
    • test_optimization_step()
    • main()
  • Module contents

Submodules

simulation.utils.machine_learning.models.helper module

Functions:

get_norm_layer([norm_type])	Return a normalization layer.
get_scheduler(optimizer, lr_policy, ...)	Return a learning rate scheduler.
init_weights(net[, init_type, init_gain])	Initialize network weights.
init_net(net[, init_type, init_gain, device])	Initialize a network.
set_requires_grad(nets[, requires_grad])	Set requires_grad=False for all the networks to avoid unnecessary computations.

get_norm_layer(norm_type: str = 'instance') → Type[Module]

Return a normalization layer.

For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.

Parameters:

norm_type – Name of the normalization layer: batch | instance | none

get_scheduler(optimizer: Optimizer, lr_policy: str, lr_decay_iters: int, n_epochs: int, lr_step_factor: float) → LambdaLR | StepLR | ReduceLROnPlateau

Return a learning rate scheduler.

For ‘linear’, we keep the same learning rate for the first <n_epochs> epochs and linearly decay the rate to zero over the next <n_epochs_decay> epochs. For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. See https://pytorch.org/docs/stable/optim.html for more details.

Parameters:

• optimizer – Optimizer of the network’s parameters

• lr_policy – Learning rate policy. [linear | step | plateau | cosine]

• lr_decay_iters – Multiply by a gamma every lr_decay_iters iterations

• n_epochs – Number of epochs with the initial learning rate

• lr_step_factor – Multiplication factor at every step in the step scheduler

init_weights(net: Module, init_type: str = 'normal', init_gain: float = 0.02) → None

Initialize network weights.

We use ‘normal’ in the original pix2pix and CycleGAN paper. But xavier and kaiming might work better for some applications. Feel free to try yourself.

Parameters:

• net – Network to be initialized

• init_type – Name of an initialization method: normal | xavier | kaiming | orthogonal

• init_gain – Scaling factor for normal, xavier and orthogonal.

init_net(net: Module, init_type: str = 'normal', init_gain: float = 0.02, device: device = device(type='cpu')) → Module

Initialize a network.

1. register CPU/GPU device;

2. initialize the network weights

Return an initialized network.

Parameters:

• net – Network to be initialized

• init_type – Name of an initialization method: normal | xavier | kaiming | orthogonal

• init_gain – Scaling factor for normal, xavier and orthogonal.

• device – Device to the net run

set_requires_grad(nets: List[Module] | Module, requires_grad: bool = False)

Set requires_grad=False for all the networks to avoid unnecessary computations.

Parameters:

• nets – A single network or a list of networks

• requires_grad – Enable or disable grads

simulation.utils.machine_learning.models.resnet_block module

Classes:

ResnetBlock(dim, padding_type, norm_layer, ...)

Define a Resnet block.

class ResnetBlock(dim: int, padding_type: str, norm_layer: Type[Module], use_dropout: bool, use_bias: bool, n_conv_layers: int = 2, dilations: List[int] | None = None)

Bases: Module

Define a Resnet block.

Methods:

forward(x)

Standard forward with skip connection.

Attributes:

training
_is_full_backward_hook

forward(x: Tensor) → Tensor

Standard forward with skip connection.

Parameters:

x – Input tensor

training: bool

_is_full_backward_hook: bool | None

simulation.utils.machine_learning.models.resnet_generator module

Classes:

ResnetGenerator(input_nc, output_nc, ngf, ...)

Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.

class ResnetGenerator(input_nc: int, output_nc: int, ngf: int = 64, norm_layer: ~typing.Type[~torch.nn.modules.module.Module] = <class 'torch.nn.modules.batchnorm.BatchNorm2d'>, use_dropout: bool = False, n_blocks: int = 6, padding_type: str = 'reflect', activation: ~torch.nn.modules.module.Module = Tanh(), conv_layers_in_block: int = 2, dilations: ~typing.List[int] | None = None)

Bases: Module, InitOptions

Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.

We adapt Torch code and idea from Justin Johnson’s neural style transfer project( https://github.com/jcjohnson/fast-neural-style)

Methods:

forward(input)

Standard forward.

Attributes:

training
_is_full_backward_hook

forward(input: Tensor) → Tensor

Standard forward.

Parameters:

input – The input tensor

training: bool

_is_full_backward_hook: bool | None

simulation.utils.machine_learning.models.unet_block module

Classes:

UnetSkipConnectionBlock(outer_nc, inner_nc, ...)

Defines the Unet submodule with skip connection.

class UnetSkipConnectionBlock(outer_nc: int, inner_nc: int, input_nc: int | None = None, submodule: ~torch.nn.modules.module.Module | None = None, outermost: bool = False, innermost: bool = False, norm_layer: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.batchnorm.BatchNorm2d'>, use_dropout: bool = False)

Bases: Module

Defines the Unet submodule with skip connection. X.

——————-identity———————- |-- downsampling -- |submodule| – upsampling –|

Methods:

forward(x)

Forward with skip connection, if this is not the outermost.

Attributes:

training
_is_full_backward_hook

forward(x: Tensor) → Tensor

Forward with skip connection, if this is not the outermost.

Parameters:

x (torch.Tensor) – the input tensor

training: bool

_is_full_backward_hook: bool | None

simulation.utils.machine_learning.models.unet_generator module

Classes:

UnetGenerator(input_nc, output_nc, ...)

Create a Unet-based generator.

class UnetGenerator(input_nc: int, output_nc: int, num_downs: int, ngf: int = 64, norm_layer: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.batchnorm.BatchNorm2d'>, use_dropout: bool = False)

Bases: Module

Create a Unet-based generator.

Methods:

forward(input)

Standard forward.

Attributes:

training
_is_full_backward_hook

forward(input: Tensor) → Tensor

Standard forward.

Parameters:

input (Tensor) – the input tensor

training: bool

_is_full_backward_hook: bool | None

simulation.utils.machine_learning.models.wasserstein_critic module

Classes:

WassersteinCritic(input_nc[, n_blocks, ...])

class WassersteinCritic(input_nc: int, n_blocks: int = 3, norm: str = 'instance', ndf=32, height=256, width=256, use_dropout: bool = False, padding_type: str = 'reflect', conv_layers_in_block: int = 2, dilations: List[int] | None = None)

Bases: Module, InitOptions

Methods:

forward(input)	Defines the computation performed at every call.
_clip_weights([bounds])	Clip weights to given bounds.
perform_optimization_step(generator, ...[, ...])	Do one iteration to update the parameters.

Attributes:

training
_is_full_backward_hook

forward(input: Tensor)

Defines 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.

_clip_weights(bounds: Tuple[float, float] = (-0.01, 0.01))

Clip weights to given bounds.

perform_optimization_step(generator: Module, optimizer: Optimizer, batch_critic: Tensor, batch_generator: Tensor, weight_clips: Tuple[float, float] | None = None) → float

Do one iteration to update the parameters.

Parameters:

• generator – Generation network

• optimizer – Optimizer for the critic’s weights

• batch_critic – A batch of inputs for the critic

• batch_generator – A batch of inputs for the generator

• weight_clips – Optional weight bounds for the critic’s weights

Returns:

Current wasserstein distance estimated by critic.

training: bool

_is_full_backward_hook: bool | None

Module contents