direct.nn.cirim package

Submodules

direct.nn.cirim.cirim module

class direct.nn.cirim.cirim.ConvRNNStack(convs, recurrent)

Bases: Module

A stack of convolutional RNNs.

Takes as input a sequence of recurrent and convolutional layers.

__init__(convs, recurrent)

Inits ConvRNNStack.

Parameters:

• convs – List of convolutional layers.

• recurrent – Recurrent layer.

forward(_input, hidden)

Forward pass of ConvRNNStack.

Parameters:

• _input – Input tensor of shape (batch_size, seq_len, input_size).

• hidden – Hidden state of shape (num_layers * num_directions, batch_size, hidden_size).

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).

class direct.nn.cirim.cirim.ConvNonlinear(input_size, features, kernel_size, dilation, bias)

Bases: Module

A convolutional layer with nonlinearity.

__init__(input_size, features, kernel_size, dilation, bias)

Inits ConvNonlinear.

Parameters:

• input_size – Size of the input.

• features – Number of features.

• kernel_size – Size of the kernel.

• dilation – Dilation of the kernel.

• bias – Whether to use bias.

reset_parameters()

Resets the parameters of the convolutional layer.

forward(_input)

Forward pass of the convolutional layer.

Parameters:

_input – Input tensor of shape (batch_size, seq_len, input_size).

Returns:

Output tensor of shape (batch_size, seq_len, features).

class direct.nn.cirim.cirim.IndRNNCell(in_channels, hidden_channels, kernel_size=1, dilation=1, bias=True)

Bases: Module

Base class for Independently RNN cells as presented in [1]_.

References:

[1]

Li, S. et al. (2018) ‘Independently Recurrent Neural Network (IndRNN): Building A Longer and Deeper RNN’, Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, (1), pp. 5457–5466. doi: 10.1109/CVPR.2018.00572.

__init__(in_channels, hidden_channels, kernel_size=1, dilation=1, bias=True)

Parameters:

• in_channels (int) – Number of input channels.

• hidden_channels (int) – Number of hidden channels.

• kernel_size (int) – Kernel size. Default: 1.

• dilation (int) – Dilation size. Default: 1.

• bias (bool) – Whether to use bias. Default: True.

reset_parameters()

Reset the parameters.

static orthotogonalize_weights(weights, chunks=1)

Orthogonalize weights.

Parameters:

• weights – The weights to orthogonalize.

• chunks – Number of chunks. Default: 1.

Returns:

The orthogonalized weights.

forward(_input, hx)

Forward pass of the cell.

Parameters:

• _input – Input tensor of shape (batch_size, seq_len, input_size), containing input features.

• hx – Hidden state of shape (batch_size, hidden_channels, 1, 1), containing hidden state features.

Returns:

Output tensor of shape (batch_size, seq_len, hidden_channels), containing the next hidden state.

class direct.nn.cirim.cirim.CIRIM(forward_operator, backward_operator, depth=2, in_channels=2, time_steps=8, recurrent_hidden_channels=64, num_cascades=8, no_parameter_sharing=True, **kwargs)

Bases: Module

Cascades of Independently Recurrent Inference Machines implementation as presented in [1]_.

References:

[1]

Karkalousos, D. et al. (2021) ‘Assessment of Data Consistency through Cascades of Independently Recurrent Inference Machines for fast and robust accelerated MRI reconstruction’. Available at: https://arxiv.org/abs/2111.15498v1

__init__(forward_operator, backward_operator, depth=2, in_channels=2, time_steps=8, recurrent_hidden_channels=64, num_cascades=8, no_parameter_sharing=True, **kwargs)

Parameters:

• forward_operator (Callable) – Forward Operator.

• backward_operator (Callable) – Backward Operator.

• depth (int) – Number of layers.

• time_steps (int) – Number of iterations: math:T.

• in_channels (int) – Input channel number. Default is 2 for complex data.

• recurrent_hidden_channels (int) – Hidden channels number for the recurrent unit of the CIRIM Blocks. Default: 64.

• recurrent_num_layers – Number of layers for the recurrent unit of the CIRIM Block (: math:n_l). Default: 4.

• no_parameter_sharing (bool) – If False, the same CIRIM Block is used for all time_steps. Default: True.

forward(masked_kspace, sampling_mask, sensitivity_map)

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) – Coil sensitivities of shape (N, coil, height, width, complex=2).

Return type:

List[List[Union[Tensor, Any]]]

Returns:

Imspace prediction.

class direct.nn.cirim.cirim.RIMBlock(forward_operator, backward_operator, depth=2, in_channels=2, hidden_channels=64, time_steps=4, no_parameter_sharing=False)

Bases: Module

Recurrent Inference Machines block as presented in [1]_.

References:

[1]

Karkalousos, D. et al. (2021) ‘Assessment of Data Consistency through Cascades of Independently Recurrent Inference Machines for fast and robust accelerated MRI reconstruction’. Available at: https://arxiv.org/abs/2111.15498v1

__init__(forward_operator, backward_operator, depth=2, in_channels=2, hidden_channels=64, time_steps=4, no_parameter_sharing=False)

Parameters:

• forward_operator (Callable) – Forward Fourier Transform.

• backward_operator (Callable) – Backward Fourier Transform.

• depth (int) – Number of layers in the RIM block.

• in_channels (int) – Input channel number. Default is 2 for complex data.

• hidden_channels (int) – Hidden channels. Default: 64.

• time_steps (int) – Number of layers of: math:n_l recurrent unit. Default: 4.

• data_consistency – If False, the DC component is removed from the input.

forward(current_prediction, masked_kspace, sampling_mask, sensitivity_map, hidden_state, parameter_sharing=False, coil_dim=1, spatial_dims=(2, 3))

Parameters:

• current_prediction (Tensor) – Current k-space.

• 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) – Coil sensitivities of shape (N, coil, height, width, complex=2).

• hidden_state (Optional[Tensor]) – IndRNN hidden state of shape (N, hidden_channels, height, width, num_layers) if not None. Optional.

• parameter_sharing (bool) – If True, the weights of the convolutional layers are shared between the forward and backward pass.

• coil_dim (int) – Coil dimension. Default: 1.

• spatial_dims (Tuple[int, int]) – Spatial dimensions. Default: (2, 3).

Return type:

Union[Tuple[List, None], Tuple[List, Union[List, Tensor]]]

Returns:

New k-space prediction of shape (N, coil, height, width, complex=2) and next hidden state of shape (N, hidden_channels, height, width, num_layers) if parameter_sharing else None.

direct.nn.cirim.cirim_engine module

class direct.nn.cirim.cirim_engine.CIRIMEngine(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)

Bases: MRIModelEngine

Cascades of Independently Recurrent Inference Machines Engine.

__init__(cfg, model, device, forward_operator=None, backward_operator=None, mixed_precision=False, **models)

Inits :class:`CIRIMEngine.

direct.nn.cirim.config module

class direct.nn.cirim.config.CIRIMConfig(model_name='???', engine_name=None, time_steps=8, depth=2, recurrent_hidden_channels=64, num_cascades=8, no_parameter_sharing=True)

Bases: ModelConfig

time_steps = 8

depth = 2

recurrent_hidden_channels = 64

num_cascades = 8

no_parameter_sharing = True

__init__(model_name='???', engine_name=None, time_steps=8, depth=2, recurrent_hidden_channels=64, num_cascades=8, no_parameter_sharing=True)

Module contents