def _create_block(inc: int, outc: int, channels: Sequence[int],
strides: Sequence[int],
is_top: bool) -> nn.Sequential:
"""
Builds the UNet structure from the bottom up by recursing down to the bottom block, then creating sequential
blocks containing the downsample path, a skip connection around the previous block, and the upsample path.
Args:
inc: number of input channels.
outc: number of output channels.
channels: sequence of channels. Top block first.
strides: convolution stride.
is_top: True if this is the top block.
"""
c = channels[0]
s = strides[0]
subblock: Union[nn.Sequential, ResidualUnit, Convolution]
if len(channels) > 2:
subblock = _create_block(c, c, channels[1:], strides[1:],
False) # continue recursion down
upc = c * 2
else:
# the next layer is the bottom so stop recursion, create the bottom layer as the sublock for this layer
subblock = self._get_bottom_layer(c, channels[1])
upc = c + channels[1]
down = self._get_down_layer(
inc, c, s, is_top) # create layer in downsampling path
up = self._get_up_layer(upc, outc, s,
is_top) # create layer in upsampling path
return nn.Sequential(down, SkipConnection(subblock), up)
def _get_connection_block(self, down_path: nn.Module, up_path: nn.Module, subblock: nn.Module) -> nn.Module:
"""
Returns the block object defining a layer of the UNet structure including the implementation of the skip
between encoding (down) and and decoding (up) sides of the network.
Args:
down_path: encoding half of the layer
up_path: decoding half of the layer
subblock: block defining the next layer in the network.
Returns: block for this layer: `nn.Sequential(down_path, SkipConnection(subblock), up_path)`
"""
return nn.Sequential(down_path, SkipConnection(subblock), up_path)
def _create_block(inc, outc, channels, strides, kernel_sizes,
sample_kernel_sizes, is_top):
"""
Builds the UNet2d5_spvPA structure from the bottom up by recursing down to the bottom block, then creating sequential
blocks containing the downsample path, a skip connection around the previous block, and the upsample path.
"""
c = channels[0]
s = strides[0]
k = kernel_sizes[0]
sk = sample_kernel_sizes[0]
# create layer in downsampling path
down = self._get_down_layer(in_channels=inc,
out_channels=c,
kernel_size=k)
downsample = self._get_downsample_layer(in_channels=c,
out_channels=c,
strides=s,
kernel_size=sk)
if len(channels) > 2:
# continue recursion down
subblock = _create_block(c,
channels[1],
channels[1:],
strides[1:],
kernel_sizes[1:],
sample_kernel_sizes[1:],
is_top=False)
else:
# the next layer is the bottom so stop recursion, create the bottom layer as the sublock for this layer
subblock = self._get_bottom_layer(
in_channels=c,
out_channels=channels[1],
kernel_size=kernel_sizes[1],
)
upsample = self._get_upsample_layer(in_channels=channels[1],
out_channels=c,
strides=s,
up_kernel_size=sk)
subblock_with_resampling = nn.Sequential(downsample, subblock,
upsample)
# create layer in upsampling path
up = self._get_up_layer(in_channels=2 * c,
out_channels=outc,
kernel_size=k,
is_top=is_top)
return nn.Sequential(down,
SkipConnection(subblock_with_resampling), up)
def _create_block(inc: int, outc: int, channels: Sequence[int], strides: Sequence[int], is_top: bool):
"""
Builds the UNet structure from the bottom up by recursing down to the bottom block, then creating sequential
blocks containing the downsample path, a skip connection around the previous block, and the upsample path.
"""
c = channels[0]
s = strides[0]
if len(channels) > 2:
subblock = _create_block(c, c, channels[1:], strides[1:], False) # continue recursion down
upc = c * 2
else:
# the next layer is the bottom so stop recursion, create the bottom layer as the sublock for this layer
subblock = self._get_bottom_layer(c, channels[1])
upc = c + channels[1]
down = self._get_down_layer(inc, c, s, is_top) # create layer in downsampling path
up = self._get_up_layer(upc, outc, s, is_top) # create layer in upsampling path
return nn.Sequential(down, SkipConnection(subblock), up)
