neural_de.external.maxim_tf.maxim package

Subpackages

• neural_de.external.maxim_tf.maxim.blocks package
  • Submodules
  • neural_de.external.maxim_tf.maxim.blocks.attentions module
    • CALayer()
    • RCAB()
    • RDCAB()
    • SAM()
  • neural_de.external.maxim_tf.maxim.blocks.block_gating module
    • BlockGatingUnit()
    • BlockGmlpLayer()
  • neural_de.external.maxim_tf.maxim.blocks.bottleneck module
    • BottleneckBlock()
  • neural_de.external.maxim_tf.maxim.blocks.grid_gating module
    • GridGatingUnit()
    • GridGmlpLayer()
  • neural_de.external.maxim_tf.maxim.blocks.misc_gating module
    • CrossGatingBlock()
    • GetSpatialGatingWeights()
    • ResidualSplitHeadMultiAxisGmlpLayer()
  • neural_de.external.maxim_tf.maxim.blocks.others module
    • MlpBlock()
    • UpSampleRatio()
  • neural_de.external.maxim_tf.maxim.blocks.unet module
    • UNetDecoderBlock()
    • UNetEncoderBlock()
  • Module contents

Submodules

neural_de.external.maxim_tf.maxim.configs module

Configs based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py

neural_de.external.maxim_tf.maxim.layers module

Layers based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py

class neural_de.external.maxim_tf.maxim.layers.BlockImages(*args, **kwargs)

Bases: Layer

call(x, patch_size)

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state in __init__(), or the build() method that is called automatically before call() executes the first time.

Parameters:

• inputs –

  Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

  arguments, and inputs cannot be provided via the default value of a keyword argument.

  • NumPy array or Python scalar values in inputs get cast as tensors.

  • Keras mask metadata is only collected from inputs.

  • Layers are built (build(input_shape) method) using shape info from inputs only.

  • input_spec compatibility is only checked against inputs.

  • Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.

  • The SavedModel input specification is generated using inputs only.

  • Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.

• *args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.

• **kwargs –

  Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

  whether the call is meant for training or inference.

  • mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a Keras layer with masking support).

Returns:

A tensor or list/tuple of tensors.

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns:

Python dictionary.

class neural_de.external.maxim_tf.maxim.layers.Resizing(*args, **kwargs)

Bases: Layer

call(x)

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state in __init__(), or the build() method that is called automatically before call() executes the first time.

Parameters:

• inputs –

  Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

  arguments, and inputs cannot be provided via the default value of a keyword argument.

  • NumPy array or Python scalar values in inputs get cast as tensors.

  • Keras mask metadata is only collected from inputs.

  • Layers are built (build(input_shape) method) using shape info from inputs only.

  • input_spec compatibility is only checked against inputs.

  • Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.

  • The SavedModel input specification is generated using inputs only.

  • Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.

• *args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.

• **kwargs –

  Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

  whether the call is meant for training or inference.

  • mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a Keras layer with masking support).

Returns:

A tensor or list/tuple of tensors.

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns:

Python dictionary.

class neural_de.external.maxim_tf.maxim.layers.SwapAxes(*args, **kwargs)

Bases: Layer

call(x, axis_one, axis_two)

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state in __init__(), or the build() method that is called automatically before call() executes the first time.

Parameters:

• inputs –

  Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

  arguments, and inputs cannot be provided via the default value of a keyword argument.

  • NumPy array or Python scalar values in inputs get cast as tensors.

  • Keras mask metadata is only collected from inputs.

  • Layers are built (build(input_shape) method) using shape info from inputs only.

  • input_spec compatibility is only checked against inputs.

  • Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.

  • The SavedModel input specification is generated using inputs only.

  • Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.

• *args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.

• **kwargs –

  Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

  whether the call is meant for training or inference.

  • mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a Keras layer with masking support).

Returns:

A tensor or list/tuple of tensors.

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns:

Python dictionary.

class neural_de.external.maxim_tf.maxim.layers.UnblockImages(*args, **kwargs)

Bases: Layer

call(x, grid_size, patch_size)

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state in __init__(), or the build() method that is called automatically before call() executes the first time.

Parameters:

• inputs –

  Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

  arguments, and inputs cannot be provided via the default value of a keyword argument.

  • NumPy array or Python scalar values in inputs get cast as tensors.

  • Keras mask metadata is only collected from inputs.

  • Layers are built (build(input_shape) method) using shape info from inputs only.

  • input_spec compatibility is only checked against inputs.

  • Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.

  • The SavedModel input specification is generated using inputs only.

  • Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.

• *args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.

• **kwargs –

  Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

  whether the call is meant for training or inference.

  • mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a Keras layer with masking support).

Returns:

A tensor or list/tuple of tensors.

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns:

Python dictionary.

neural_de.external.maxim_tf.maxim.maxim module

MAXIM based on https://github.com/google-research/maxim/blob/main/maxim/models/maxim.py

neural_de.external.maxim_tf.maxim.maxim.MAXIM(features=64, depth=3, num_stages=2, num_groups=1, use_bias=True, num_supervision_scales=1, lrelu_slope=0.2, use_global_mlp=True, use_cross_gating=True, high_res_stages=2, block_size_hr=(16, 16), block_size_lr=(8, 8), grid_size_hr=(16, 16), grid_size_lr=(8, 8), num_bottleneck_blocks=1, block_gmlp_factor=2, grid_gmlp_factor=2, input_proj_factor=2, channels_reduction=4, num_outputs=3, dropout_rate=0.0)

The MAXIM model function with multi-stage and multi-scale supervision.

For more model details, please check the CVPR paper: MAXIM: MUlti-Axis MLP for Image Processing (https://arxiv.org/abs/2201.02973)

neural_de.external.maxim_tf.maxim.maxim.features

initial hidden dimension for the input resolution.

neural_de.external.maxim_tf.maxim.maxim.depth

the number of downsampling depth for the model.

neural_de.external.maxim_tf.maxim.maxim.num_stages

how many stages to use. It will also affects the output list.

neural_de.external.maxim_tf.maxim.maxim.num_groups

how many blocks each stage contains.

neural_de.external.maxim_tf.maxim.maxim.use_bias

whether to use bias in all the conv/mlp layers.

neural_de.external.maxim_tf.maxim.maxim.num_supervision_scales

the number of desired supervision scales.

neural_de.external.maxim_tf.maxim.maxim.lrelu_slope

the negative slope parameter in leaky_relu layers.

neural_de.external.maxim_tf.maxim.maxim.use_global_mlp

whether to use the multi-axis gated MLP block (MAB) in each layer.

neural_de.external.maxim_tf.maxim.maxim.use_cross_gating

whether to use the cross-gating MLP block (CGB) in the skip connections and multi-stage feature fusion layers.

neural_de.external.maxim_tf.maxim.maxim.high_res_stages

how many stages are specificied as high-res stages. The rest (depth - high_res_stages) are called low_res_stages.

neural_de.external.maxim_tf.maxim.maxim.block_size_hr

the block_size parameter for high-res stages.

neural_de.external.maxim_tf.maxim.maxim.block_size_lr

the block_size parameter for low-res stages.

neural_de.external.maxim_tf.maxim.maxim.grid_size_hr

the grid_size parameter for high-res stages.

neural_de.external.maxim_tf.maxim.maxim.grid_size_lr

the grid_size parameter for low-res stages.

neural_de.external.maxim_tf.maxim.maxim.num_bottleneck_blocks

how many bottleneck blocks.

neural_de.external.maxim_tf.maxim.maxim.block_gmlp_factor

the input projection factor for block_gMLP layers.

neural_de.external.maxim_tf.maxim.maxim.grid_gmlp_factor

the input projection factor for grid_gMLP layers.

neural_de.external.maxim_tf.maxim.maxim.input_proj_factor

the input projection factor for the MAB block.

neural_de.external.maxim_tf.maxim.maxim.channels_reduction

the channel reduction factor for SE layer.

neural_de.external.maxim_tf.maxim.maxim.num_outputs

the output channels.

neural_de.external.maxim_tf.maxim.maxim.dropout_rate

Dropout rate.

Returns:

The output contains a list of arrays consisting of multi-stage multi-scale outputs. For example, if num_stages = num_supervision_scales = 3 (the model used in the paper), the output specs are: outputs = [[output_stage1_scale1, output_stage1_scale2, output_stage1_scale3],

[output_stage2_scale1, output_stage2_scale2, output_stage2_scale3], [output_stage3_scale1, output_stage3_scale2, output_stage3_scale3],]

The final output can be retrieved by outputs[-1][-1].

Module contents