meddlr.modeling.meta_arch.VortexModel#

class meddlr.modeling.meta_arch.VortexModel(model: Module, augmentor: MRIReconAugmentor, use_supervised_consistency: bool = False, vis_period: int = -1)[source]#

VORTEX model.

This is the generalized model implementation for augmentation-based consistency. It differs from N2RModel and M2RModel in some ways:

  1. Generalizable augmentor: MRIReconAugmentor is used to perform augmentations.

  2. Faster augmentations: Augmentations are performed on the operating device (e.g. GPU) with large, but reproducible seeds.

  3. Spatial augmentations: Consistency with spatial augmentations are also supported. These augmentation are also used to transform the target image.

Reference:

A Desai, B Gunel, B Ozturkler, et al. VORTEX: Physics-Driven Data Augmentations Using Consistency Training for Robust Accelerated MRI Reconstruction. https://arxiv.org/abs/2111.02549.

__init__(model: Module, augmentor: MRIReconAugmentor, use_supervised_consistency: bool = False, vis_period: int = -1)[source]#
Parameters

  • model (nn.Module) – The base model.

  • augmentor (MRIReconAugmentor) – The augmentation module.

  • use_supervised_consistency (bool, optional) – If True, use consistency with supervised examples too.

  • vis_period (int, optional) – The period over which to visualize images. If <=0, it is ignored. Note if the model has a vis_period attribute, it will be overridden so that this class handles visualization.

Methods

__init__(model, augmentor[, ...])

param model

The base model.

add_module(name, module)

Adds a child module to the current module.

apply(fn)

Applies fn recursively to every submodule (as returned by .children()) as well as self.

augment(inputs, pred_base)

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

buffers([recurse])

Returns an iterator over module buffers.

children()

Returns an iterator over immediate children modules.

cpu()

Moves all model parameters and buffers to the CPU.

cuda([device])

Moves all model parameters and buffers to the GPU.

double()

Casts all floating point parameters and buffers to double datatype.

eval()

Sets the module in evaluation mode.

extra_repr()

Set the extra representation of the module

float()

Casts all floating point parameters and buffers to float datatype.

forward(inputs)

Defines the computation performed at every call.

from_config(cfg)

get_buffer(target)

Returns the buffer given by target if it exists, otherwise throws an error.

get_extra_state()

Returns any extra state to include in the module's state_dict.

get_parameter(target)

Returns the parameter given by target if it exists, otherwise throws an error.

get_submodule(target)

Returns the submodule given by target if it exists, otherwise throws an error.

half()

Casts all floating point parameters and buffers to half datatype.

ipu([device])

Moves all model parameters and buffers to the IPU.

load_state_dict(state_dict[, strict])

Copies parameters and buffers from state_dict into this module and its descendants.

modules()

Returns an iterator over all modules in the network.

named_buffers([prefix, recurse, ...])

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

named_children()

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

named_modules([memo, prefix, remove_duplicate])

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

named_parameters([prefix, recurse, ...])

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

parameters([recurse])

Returns an iterator over module parameters.

register_backward_hook(hook)

Registers a backward hook on the module.

register_buffer(name, tensor[, persistent])

Adds a buffer to the module.

register_forward_hook(hook, *[, prepend, ...])

Registers a forward hook on the module.

register_forward_pre_hook(hook, *[, ...])

Registers a forward pre-hook on the module.

register_full_backward_hook(hook[, prepend])

Registers a backward hook on the module.

register_full_backward_pre_hook(hook[, prepend])

Registers a backward pre-hook on the module.

register_load_state_dict_post_hook(hook)

Registers a post hook to be run after module's load_state_dict is called.

register_module(name, module)

Alias for add_module().

register_parameter(name, param)

Adds a parameter to the module.

register_state_dict_pre_hook(hook)

These hooks will be called with arguments: self, prefix, and keep_vars before calling state_dict on self.

requires_grad_([requires_grad])

Change if autograd should record operations on parameters in this module.

set_extra_state(state)

This function is called from load_state_dict() to handle any extra state found within the state_dict.

share_memory()

See torch.Tensor.share_memory_()

state_dict(*args[, destination, prefix, ...])

Returns a dictionary containing references to the whole state of the module.

to(*args, **kwargs)

Moves and/or casts the parameters and buffers.

to_empty(*, device)

Moves the parameters and buffers to the specified device without copying storage.

train([mode])

Sets the module in training mode.

type(dst_type)

Casts all parameters and buffers to dst_type.

visualize_aug_training(kspace, kspace_aug, ...)

Visualize training of augmented data.

xpu([device])

Moves all model parameters and buffers to the XPU.

zero_grad([set_to_none])

Sets gradients of all model parameters to zero.

Attributes

T_destination

alias of TypeVar('T_destination', bound=Dict[str, Any])

call_super_init

dump_patches