direct.nn.varnet package

direct.nn.varnet package#

Submodules#

direct.nn.varnet.config module#

class direct.nn.varnet.config.EndToEndVarNetConfig(model_name='???', engine_name=None, num_layers=8, regularizer_num_filters=18, regularizer_num_pull_layers=4, regularizer_dropout=0.0)[source]#

Bases: ModelConfig

num_layers = 8#
regularizer_num_filters = 18#
regularizer_num_pull_layers = 4#
regularizer_dropout = 0.0#
__init__(model_name='???', engine_name=None, num_layers=8, regularizer_num_filters=18, regularizer_num_pull_layers=4, regularizer_dropout=0.0)#

direct.nn.varnet.varnet module#

class direct.nn.varnet.varnet.EndToEndVarNet(forward_operator, backward_operator, num_layers, regularizer_num_filters=18, regularizer_num_pull_layers=4, regularizer_dropout=0.0, in_channels=2, **kwargs)[source]#

Bases: Module

End-to-End Variational Network based on [1].

References:

__init__(forward_operator, backward_operator, num_layers, regularizer_num_filters=18, regularizer_num_pull_layers=4, regularizer_dropout=0.0, in_channels=2, **kwargs)[source]#

Inits EndToEndVarNet.

Parameters:

  • forward_operator (Callable) – Forward Operator.

  • backward_operator (Callable) – Backward Operator.

  • num_layers (int) – Number of cascades.

  • regularizer_num_filters (int) – Regularizer model number of filters. Default: 18.

  • regularizer_num_pull_layers (int) – Regularizer model number of pulling layers. Default: 4.

  • regularizer_dropout (float) – Regularizer model dropout probability. Default: 0.0.

  • in_channels (int) – Number of input channels. Default: 2.

  • **kwargs – Additional keyword arguments.

forward(masked_kspace, sampling_mask, sensitivity_map)[source]#

Performs the forward pass of EndToEndVarNet.

Parameters:

  • masked_kspace (Tensor) – Masked k-space of shape (N, coil, height, width, complex=2).

  • sampling_mask (Tensor) – Sampling mask of shape (N, 1, height, width, 1).

  • sensitivity_map (Tensor) – Sensitivity map of shape (N, coil, height, width, complex=2).

Return type:

Tensor

Returns:

K-space prediction of shape (N, coil, height, width, complex=2).

class direct.nn.varnet.varnet.EndToEndVarNetBlock(forward_operator, backward_operator, regularizer_model)[source]#

Bases: Module

End-to-End Variational Network block.

__init__(forward_operator, backward_operator, regularizer_model)[source]#

Inits EndToEndVarNetBlock.

Parameters:

  • forward_operator (Callable) – Forward Operator.

  • backward_operator (Callable) – Backward Operator.

  • regularizer_model (Module) – Regularizer model.

forward(current_kspace, masked_kspace, sampling_mask, sensitivity_map)[source]#

Performs the forward pass of EndToEndVarNetBlock.

Parameters:

  • current_kspace (Tensor) – Current k-space prediction of shape (N, coil, height, width, complex=2).

  • masked_kspace (Tensor) – Masked k-space of shape (N, coil, height, width, complex=2).

  • sampling_mask (Tensor) – Sampling mask of shape (N, 1, height, width, 1).

  • sensitivity_map (Tensor) – Sensitivity map of shape (N, coil, height, width, complex=2).

Return type:

Tensor

Returns:

Next k-space prediction of shape (N, coil, height, width, complex=2).

direct.nn.varnet.varnet_engine module#

Engines for End-to-End Variational Network model.

Includes supervised, self-supervised and joint supervised and self-supervised learning engines.

class direct.nn.varnet.varnet_engine.EndToEndVarNetEngine(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)[source]#

Bases: MRIModelEngine

End-to-End Variational Network Engine.

Parameters:

  • cfg (BaseConfig) – Configuration file.

  • model (Module) – Model.

  • device (str) – Device. Can be “cuda: {idx}” or “cpu”.

  • forward_operator (Optional[Callable]) – The forward operator. Default: None.

  • backward_operator (Optional[Callable]) – The backward operator. Default: None.

  • mixed_precision (bool) – Use mixed precision. Default: False.

  • **models (Module) – Additional models.

__init__(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)[source]#

Inits EndToEndVarNetEngine.

Parameters:

  • cfg (BaseConfig) – Configuration file.

  • model (Module) – Model.

  • device (str) – Device. Can be “cuda: {idx}” or “cpu”.

  • forward_operator (Optional[Callable]) – The forward operator. Default: None.

  • backward_operator (Optional[Callable]) – The backward operator. Default: None.

  • mixed_precision (bool) – Use mixed precision. Default: False.

  • **models (Module) – Additional models.

forward_function(data)[source]#

This method performs the model’s forward method given data which contains all tensor inputs.

Must be implemented by child classes.

Return type:

tuple[Tensor, Tensor]

class direct.nn.varnet.varnet_engine.EndToEndVarNetSSLEngine(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)[source]#

Bases: SSLMRIModelEngine

Self-supervised Learning End-to-End Variational Network Engine.

Used for supplementary experiments for End-to-End Variational Network model with SLL in the JSSL paper [1].

Parameters:

  • cfg (BaseConfig) – Configuration file.

  • model (Module) – Model.

  • device (str) – Device. Can be “cuda: {idx}” or “cpu”.

  • forward_operator (Optional[Callable]) – The forward operator. Default: None.

  • backward_operator (Optional[Callable]) – The backward operator. Default: None.

  • mixed_precision (bool) – Use mixed precision. Default: False.

  • **models (Module) – Additional models.

  • References

  • Yiasemis (.. [1]) – JSSL: Joint Supervised and

  • G. – JSSL: Joint Supervised and

  • Moriakov – JSSL: Joint Supervised and

  • N. – JSSL: Joint Supervised and

  • Sánchez – JSSL: Joint Supervised and

  • C.I. – JSSL: Joint Supervised and

  • Sonke – JSSL: Joint Supervised and

  • J.-J. – JSSL: Joint Supervised and

  • Teuwen – JSSL: Joint Supervised and

  • J. – JSSL: Joint Supervised and

  • Reconstruction (Self-supervised Learning for MRI) – //arxiv.org/abs/2311.15856, (2023).

  • http – //arxiv.org/abs/2311.15856, (2023).

  • https – //doi.org/10.48550/arXiv.2311.15856.

__init__(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)[source]#

Inits EndToEndVarNetSSLEngine.

Parameters:

  • cfg (BaseConfig) – Configuration file.

  • model (Module) – Model.

  • device (str) – Device. Can be “cuda: {idx}” or “cpu”.

  • forward_operator (Optional[Callable]) – The forward operator. Default: None.

  • backward_operator (Optional[Callable]) – The backward operator. Default: None.

  • mixed_precision (bool) – Use mixed precision. Default: False.

  • **models (Module) – Additional models.

forward_function(data)[source]#

Forward function for EndToEndVarNetSSLEngine.

Parameters:

  • data (dict[str, Any]) – Data dictionary. Should contain the following keys:

  • training (- "input_sampling_mask" if)

  • inference ("sampling_mask" if)

  • training

  • inference

  • "sensitivity_map" (-)

Return type:

tuple[None, Tensor]

Returns:

None for image and output k-space.

class direct.nn.varnet.varnet_engine.EndToEndVarNetJSSLEngine(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)[source]#

Bases: JSSLMRIModelEngine

Joint Supervised and Self-supervised Learning End-to-End Variational Network Engine.

Used for supplementary experiments for End-to-End Variational Network model with JSLL in the JSSL paper [1].

Parameters:

  • cfg (BaseConfig) – Configuration file.

  • model (Module) – Model.

  • device (str) – Device. Can be “cuda: {idx}” or “cpu”.

  • forward_operator (Optional[Callable]) – The forward operator. Default: None.

  • backward_operator (Optional[Callable]) – The backward operator. Default: None.

  • mixed_precision (bool) – Use mixed precision. Default: False.

  • **models (Module) – Additional models.

  • References

  • Yiasemis (.. [1]) – JSSL: Joint Supervised and

  • G. – JSSL: Joint Supervised and

  • Moriakov – JSSL: Joint Supervised and

  • N. – JSSL: Joint Supervised and

  • Sánchez – JSSL: Joint Supervised and

  • C.I. – JSSL: Joint Supervised and

  • Sonke – JSSL: Joint Supervised and

  • J.-J. – JSSL: Joint Supervised and

  • Teuwen – JSSL: Joint Supervised and

  • J. – JSSL: Joint Supervised and

  • Reconstruction (Self-supervised Learning for MRI) – //arxiv.org/abs/2311.15856, (2023).

  • http – //arxiv.org/abs/2311.15856, (2023).

  • https – //doi.org/10.48550/arXiv.2311.15856.

__init__(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)[source]#

Inits EndToEndVarNetJSSLEngine.

Parameters:

  • cfg (BaseConfig) – Configuration file.

  • model (Module) – Model.

  • device (str) – Device. Can be “cuda: {idx}” or “cpu”.

  • forward_operator (Optional[Callable]) – The forward operator. Default: None.

  • backward_operator (Optional[Callable]) – The backward operator. Default: None.

  • mixed_precision (bool) – Use mixed precision. Default: False.

  • **models (Module) – Additional models.

forward_function(data)[source]#

Forward function for EndToEndVarNetJSSLEngine.

Parameters:

  • data (dict[str, Any]) – Data dictionary. Should contain the following keys:

  • SSL (- "input_sampling_mask" if training and training is)

  • SSL

  • inference ("sampling_mask" if)

  • SSL

  • inference

  • "sensitivity_map" (-)

Return type:

tuple[None, Tensor]

Returns:

None for image and output k-space.

Module contents#

Contents