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)