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",
)

Bases: Sequential

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

Source code in src/autoden/models/msd.py 
13
14
15
16
17
18
def __init__(self, in_ch: int, out_ch: int, dilation: int = 1, pad_mode: str = "replicate") -> None:
    super().__init__(
        nn.Conv2d(in_ch, out_ch, 3, padding=dilation, dilation=dilation, padding_mode=pad_mode),
        nn.BatchNorm2d(out_ch),
        nn.LeakyReLU(0.2, inplace=True),
    )

MSDDilBlock 

MSDDilBlock(
    n_channels_in: int,
    n_features: int,
    n_layers: int,
    dilations: Sequence[int],
)

Bases: Module

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

Source code in src/autoden/models/msd.py 
73
74
75
76
77
78
79
80
81
82
83
def __init__(self, n_channels_in: int, n_features: int, n_layers: int, dilations: Sequence[int]) -> None:
    super().__init__()
    self.n_features = n_features
    self.n_layers = n_layers
    self.dilations = dilations
    convs = [
        DilatedConvBlock(n_channels_in + n_features * ii, n_features, dilation=self._layer_dilation(ii))
        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],
)

Bases: Module

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

Source code in src/autoden/models/msd.py 
38
39
40
41
42
43
44
45
46
47
48
def __init__(self, n_channels_in: int, n_features: int, n_layers: int, dilations: Sequence[int]) -> None:
    super().__init__()
    self.n_features = n_features
    self.n_layers = n_layers
    self.dilations = dilations
    convs = [
        SamplingConvBlock(n_channels_in + n_features * ii, n_features, samp_factor=self._layer_sampling(ii))
        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,
    dilations: Sequence[int] = [1, 2, 3, 4],
    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 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def __init__(
    self,
    n_channels_in: int = 1,
    n_channels_out: int = 1,
    n_layers: int = 12,
    n_features: int = 1,
    dilations: Sequence[int] = [1, 2, 3, 4],
    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)
    else:
        self.msd_block = MSDSampBlock(n_channels_in, n_features, n_layers, dilations)
    self.outc = nn.Conv2d(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
)

Bases: Sequential

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

Source code in src/autoden/models/msd.py 
24
25
26
27
28
29
30
31
32
def __init__(self, in_ch: int, out_ch: int, samp_factor: int = 1) -> None:
    if samp_factor > 1:
        pre = [nn.AvgPool2d(samp_factor)]
        post = [nn.Upsample(scale_factor=samp_factor, mode="bilinear", align_corners=True)]
    else:
        pre = post = []
    super().__init__(
        *pre, nn.Conv2d(in_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.LeakyReLU(0.2, inplace=True), *post
    )