Skip to content

msd

Module implementing MS-D net.

Classes:

  • DilatedConvBlock –

    Dilated convolution block (dilated_conv => BN => ReLU).

  • MSDDilBlock –

    MS-D Block containing the sequence of dilated convolutional layers.

  • MSDSampBlock –

    MS-D Block containing the sequence of dilated convolutional layers.

  • MSDnet –

    Simple MS-D net implementation.

  • SamplingConvBlock –

    Down-sampling convolution module (down-samp => conv => BN => ReLU => up-samp).

DilatedConvBlock 

DilatedConvBlock(
    in_ch: int,
    out_ch: int,
    dilation: int = 1,
    pad_mode: str = "replicate",
    n_dims: int = 2,
)

Bases: Sequential

Dilated convolution block (dilated_conv => BN => ReLU).

Source code in src/autoden/models/msd.py 
26
27
28
29
30
31
def __init__(self, in_ch: int, out_ch: int, dilation: int = 1, pad_mode: str = "replicate", n_dims: int = 2) -> None:
    super().__init__(
        NDConv[n_dims](in_ch, out_ch, 3, padding=dilation, dilation=dilation, padding_mode=pad_mode),
        NDBatchNorm[n_dims](out_ch),
        nn.LeakyReLU(0.2, inplace=True),
    )

MSDDilBlock 

MSDDilBlock(
    n_channels_in: int,
    n_features: int,
    n_layers: int,
    dilations: Sequence[int] | NDArray,
    n_dims: int = 2,
)

Bases: Module

MS-D Block containing the sequence of dilated convolutional layers.

Source code in src/autoden/models/msd.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def __init__(
    self, n_channels_in: int, n_features: int, n_layers: int, dilations: Sequence[int] | NDArray, n_dims: int = 2
) -> None:
    super().__init__()
    self.n_features = n_features
    self.n_layers = n_layers
    self.dilations = dilations
    self.n_dims = n_dims
    convs = [
        DilatedConvBlock(n_channels_in + n_features * ii, n_features, dilation=self._layer_dilation(ii), n_dims=n_dims)
        for ii in range(n_layers)
    ]
    self.convs = nn.ModuleList(convs)
    self.n_ch_in = n_channels_in

MSDSampBlock 

MSDSampBlock(
    n_channels_in: int,
    n_features: int,
    n_layers: int,
    dilations: Sequence[int] | NDArray,
    n_dims: int = 2,
)

Bases: Module

MS-D Block containing the sequence of dilated convolutional layers.

Source code in src/autoden/models/msd.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def __init__(
    self, n_channels_in: int, n_features: int, n_layers: int, dilations: Sequence[int] | NDArray, n_dims: int = 2
) -> None:
    super().__init__()
    self.n_features = n_features
    self.n_layers = n_layers
    self.dilations = dilations
    self.n_dims = n_dims
    convs = [
        SamplingConvBlock(n_channels_in + n_features * ii, n_features, samp_factor=self._layer_sampling(ii), n_dims=n_dims)
        for ii in range(n_layers)
    ]
    self.convs = nn.ModuleList(convs)
    self.n_ch_in = n_channels_in

MSDnet 

MSDnet(
    n_channels_in: int = 1,
    n_channels_out: int = 1,
    n_layers: int = 12,
    n_features: int = 1,
    n_dims: int = 2,
    dilations: Sequence[int] | NDArray = tuple(range(1, 5)),
    device: str = "cuda" if is_available() else "cpu",
    use_dilations: bool = True,
)

Bases: Module

Simple MS-D net implementation.

Source code in src/autoden/models/msd.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
def __init__(
    self,
    n_channels_in: int = 1,
    n_channels_out: int = 1,
    n_layers: int = 12,
    n_features: int = 1,
    n_dims: int = 2,
    dilations: Sequence[int] | NDArray = tuple(range(1, 5)),
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    use_dilations: bool = True,
) -> None:
    init_params = locals()
    del init_params["self"]
    del init_params["__class__"]

    super().__init__()
    self.init_params = init_params
    self.device = device

    if use_dilations:
        self.msd_block = MSDDilBlock(n_channels_in, n_features, n_layers, dilations, n_dims=n_dims)
    else:
        self.msd_block = MSDSampBlock(n_channels_in, n_features, n_layers, dilations, n_dims=n_dims)
    self.outc = NDConv[n_dims](n_channels_in + n_features * n_layers, n_channels_out, kernel_size=1)

    self.to(self.device)

SamplingConvBlock 

SamplingConvBlock(
    in_ch: int,
    out_ch: int,
    samp_factor: int = 1,
    n_dims: int = 2,
)

Bases: Sequential

Down-sampling convolution module (down-samp => conv => BN => ReLU => up-samp).

Source code in src/autoden/models/msd.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
def __init__(self, in_ch: int, out_ch: int, samp_factor: int = 1, n_dims: int = 2) -> None:
    if samp_factor > 1:
        pre = [NDPool[n_dims](samp_factor)]
        post = [nn.Upsample(scale_factor=samp_factor, mode=NDUpsampling[n_dims], align_corners=True)]
    else:
        pre = post = []
    super().__init__(
        *pre,
        NDConv[n_dims](in_ch, out_ch, 3, padding=1),
        NDBatchNorm[n_dims](out_ch),
        nn.LeakyReLU(0.2, inplace=True),
        *post,
    )