def __init__(
self,
config: LayerConfig,
channels: int,
*,
kernel_size: Optional[KernelSize] = None,
padding: Optional[Padding] = None,
padding_mode: Optional[str] = None,
base_block: ConvBlockType = BlockConvNormActivation):
super().__init__()
config = copy.copy(config)
conv_kwargs = copy.copy(config.conv_kwargs)
if kernel_size is not None:
conv_kwargs['kernel_size'] = kernel_size
if padding is not None:
conv_kwargs['padding'] = padding
if padding_mode is not None:
conv_kwargs['padding_mode'] = padding_mode
config.conv_kwargs = conv_kwargs
_posprocess_padding(conv_kwargs)
stride = 1
self.block_1 = base_block(
config, channels, channels,
kernel_size=kernel_size, padding=padding,
stride=stride, padding_mode=padding_mode)
self.activation = config.activation(**config.activation_kwargs)
config.activation = None
self.block_2 = base_block(
config, channels, channels,
kernel_size=kernel_size, padding=padding,
stride=stride, padding_mode=padding_mode)
def denses(
sizes: Sequence[int],
dropout_probability: float = None,
activation: Any = nn.ReLU,
normalization_type: NormType = NormType.BatchNorm,
last_layer_is_output: bool = False,
with_flatten: bool = True,
config: LayerConfig = default_layer_config(dimensionality=None)
) -> nn.Module:
"""
Args:
sizes: the size of the linear layers_legacy. The format is [linear1_input, linear1_output, ..., linearN_output]
dropout_probability: the probability of the dropout layer. If `None`, no dropout layer is added.
activation: the activation to be used
normalization_type: the normalization to be used between dense layers_legacy. If `None`, no normalization added
last_layer_is_output: This must be set to `True` if the last layer of dense is actually an output.
If the last layer is an output, we should not add batch norm, dropout or
activation of the last `nn.Linear`
with_flatten: if True, the input will be flattened
config: defines the available operations
Returns:
a nn.Module
"""
config = copy.copy(config)
if activation is not None:
config.activation = activation
config.norm_type = normalization_type
config.set_dim(1)
ops = []
if with_flatten:
ops.append(Flatten())
for n in range(len(sizes) - 1):
current = sizes[n]
next = sizes[n + 1]
ops.append(nn.Linear(current, next))
if n + 2 == len(sizes):
if not last_layer_is_output:
ops.append(activation(**config.activation_kwargs))
else:
if config.norm_type is not None:
ops.append(nn.BatchNorm1d(next, **config.norm_kwargs))
ops.append(activation(**config.activation_kwargs))
if dropout_probability is not None and config.dropout is not None:
ops.append(
config.dropout(p=dropout_probability,
**config.dropout_kwargs))
return nn.Sequential(*ops)
def __init__(
self,
config: LayerConfig,
input_channels: int,
output_channels: int,
*,
kernel_size: Optional[KernelSize] = None,
padding: Optional[Padding] = None,
output_padding: Optional[Union[int, Sequence[int]]] = None,
stride: Optional[Stride] = None,
padding_mode: Optional[str] = None):
super().__init__()
# local override of the default config
deconv_kwargs = copy.copy(config.deconv_kwargs)
if kernel_size is not None:
deconv_kwargs['kernel_size'] = kernel_size
if padding is not None:
deconv_kwargs['padding'] = padding
if stride is not None:
deconv_kwargs['stride'] = stride
if output_padding is not None:
deconv_kwargs['output_padding'] = output_padding
if padding_mode is not None:
deconv_kwargs['padding_mode'] = padding_mode
_posprocess_padding(deconv_kwargs)
ops = [
config.deconv(
in_channels=input_channels,
out_channels=output_channels,
**deconv_kwargs),
]
if config.norm is not None:
ops.append(config.norm(num_features=output_channels, **config.norm_kwargs))
if config.activation is not None:
ops.append(config.activation(**config.activation_kwargs))
self.ops = nn.Sequential(*ops)
def __init__(
self,
config: LayerConfig,
kernel_size: Optional[KernelSize] = 2):
super().__init__()
pool_kwargs = copy.copy(config.pool_kwargs)
if kernel_size is not None:
pool_kwargs['kernel_size'] = kernel_size
self.op = config.pool(**pool_kwargs)
def __init__(self,
layer_config: LayerConfig,
bloc_level: int,
input_channels: int,
output_channels: int,
latent_channels: Optional[int] = None,
block: nn.Module = BlockConvNormActivation,
**block_kwargs):
super().__init__()
self.latent_channels = latent_channels
self.input_channels = input_channels
self.output_channels = output_channels
if latent_channels is None:
latent_channels = 0
# local copy for local configuration
layer_config = copy.copy(layer_config)
for key, value in block_kwargs.items():
layer_config.conv_kwargs[key] = value
self.dim = layer_config.ops.dim
self.ops = block(layer_config, input_channels + latent_channels,
output_channels)
def __init__(
self,
dimensionality: int,
input_channels: int,
base_model: ModuleWithIntermediate,
deconv_filters: Sequence[int],
convolution_kernels: Union[int, Sequence[int]],
strides: Union[int, Sequence[int]],
activation=nn.ReLU,
nb_classes: Optional[int] = None,
concat_mode: str = 'add',
conv_filters: Optional[Sequence[int]] = None,
norm_type: NormType = NormType.BatchNorm,
norm_kwargs: Dict = {},
activation_kwargs: Dict = {},
deconv_block_fn: nn.Module = BlockDeconvNormActivation,
config: LayerConfig = default_layer_config(dimensionality=None)):
"""
Args:
base_model: a base model. Must be an instance of :class:`trw.layers_legacy.ModuleWithIntermediate`. This
model will return intermediate results of convolutional groups
deconv_filters: the number of filters of the transposed convolution layers. Specified from the model output
back to the input. It must have "intermediate results" filters.
convolution_kernels (list, int): the kernel sizes used by the transposed convolution. Can be an int or a list
strides (list, int): the strides used by the transposed convolution. Can be an int or a list.
concat_mode: one of 'add' or 'concatenate'. If 'add', the skip connection will be added to the
corresponding transposed convolution (i.e., similar to the FCNN paper). If 'concatenate',
the channels will be concatenated instead of added (i.e., similar to UNET)
conv_filters: if concat_mode == 'concatenate', we MUST have the shapes of the intermediates
"""
super().__init__()
# update the configuration locally
config = copy.copy(config)
if norm_type is not None:
config.norm_type = norm_type
if activation is not None:
config.activation = activation
config.set_dim(dimensionality)
config.norm_kwargs = {**norm_kwargs, **config.norm_kwargs}
config.activation_kwargs = {
**activation_kwargs,
**config.activation_kwargs
}
self.base_model = base_model
self.deconv_filters = deconv_filters
self.nb_classes = nb_classes
self.concat_mode = concat_mode
self.conv_filters = conv_filters
# normalize the arguments
nb_convs = len(deconv_filters)
if not isinstance(convolution_kernels, collections.Sequence):
convolution_kernels = [convolution_kernels] * nb_convs
if not isinstance(strides, collections.Sequence):
strides = [strides] * nb_convs
assert isinstance(
base_model, ModuleWithIntermediate
), 'it must be a model that returns intermediate results!'
assert concat_mode in ('add', 'concatenate')
# generic 2D/3D or nD once pytorch supports it
assert len(convolution_kernels) == len(strides)
assert len(convolution_kernels) == len(deconv_filters), 'the last `deconv_filters` should be the ' \
'number of output classes'
groups_deconv = nn.ModuleList()
prev_filter = input_channels
for layer_n in range(len(deconv_filters)):
current_filter = deconv_filters[layer_n]
kernel = convolution_kernels[layer_n]
stride = strides[layer_n]
output_padding = stride - 1
padding = div_shape(kernel, 2)
ops = deconv_block_fn(config=config,
input_channels=prev_filter,
output_channels=current_filter,
kernel_size=kernel,
stride=stride,
padding=padding,
output_padding=output_padding)
groups_deconv.append(ops)
if concat_mode == 'concatenate':
assert conv_filters is not None
if layer_n + 1 < len(conv_filters):
current_filter += conv_filters[::-1][layer_n + 1]
prev_filter = current_filter
self.deconvolutions = groups_deconv
if nb_classes is not None:
self.classifier = config.conv(deconv_filters[-1],
nb_classes,
kernel_size=1)
def __init__(
self,
dimensionality: int,
input_channels: int,
*,
channels: Sequence[int],
convolution_kernels: ConvKernels = 5,
strides: ConvStrides = 1,
pooling_size: Optional[PoolingSizes] = 2,
convolution_repeats: Union[int, List[int], NestedIntSequence] = 1,
activation: Optional[Activation] = nn.ReLU,
padding: Paddings = 'same',
with_flatten: bool = False,
dropout_probability: Optional[float] = None,
norm_type: Optional[NormType] = None,
norm_kwargs: Dict = {},
pool_kwargs: Dict = {},
activation_kwargs: Dict = {},
last_layer_is_output: bool = False,
conv_block_fn: ConvBlockType = BlockConvNormActivation,
config: LayerConfig = default_layer_config(dimensionality=None)):
"""
Args:
dimensionality: the dimension of the CNN (2 for 2D or 3 for 3D)
input_channels: the number of input channels
channels: the number of channels for each convolutional layer
convolution_kernels: for each convolution group, the kernel of the convolution
strides: for each convolution group, the stride of the convolution
pooling_size: the pooling size to be inserted after each convolution group
convolution_repeats: the number of repeats of a convolution. ``1`` means no repeat.
activation: the activation function
with_flatten: if True, the last output will be flattened
norm_kwargs: additional arguments to be used for the normalization layer
padding: 'same' will add padding so that convolution output as the same size as input
last_layer_is_output: if True, the last convolution will NOT have activation, dropout, batch norm, LRN
dropout_probability: dropout probability
"""
super().__init__()
# update the configuration locally
config = copy.copy(config)
if norm_type is not None:
config.norm_type = norm_type
if activation is not None:
config.activation = activation
config.set_dim(dimensionality)
config.pool_kwargs = {**pool_kwargs, **config.pool_kwargs}
config.norm_kwargs = {**norm_kwargs, **config.norm_kwargs}
config.activation_kwargs = {
**activation_kwargs,
**config.activation_kwargs
}
# normalize the arguments
nb_convs = len(channels)
if not isinstance(convolution_kernels, list):
convolution_kernels = [convolution_kernels] * nb_convs
if not isinstance(strides, list):
strides = [strides] * nb_convs
if not isinstance(pooling_size, list) and pooling_size is not None:
pooling_size = [pooling_size] * nb_convs
if isinstance(convolution_repeats, int):
convolution_repeats = [convolution_repeats] * nb_convs
if isinstance(padding, int):
padding = [padding] * nb_convs
elif isinstance(padding, str):
pass
else:
assert len(padding) == nb_convs
assert nb_convs == len(
convolution_kernels
), 'must be specified for each convolutional layer'
assert nb_convs == len(
strides), 'must be specified for each convolutional layer'
assert pooling_size is None or nb_convs == len(
pooling_size), 'must be specified for each convolutional layer'
assert nb_convs == len(convolution_repeats)
self.with_flatten = with_flatten
layers = nn.ModuleList()
prev = input_channels
for n in range(len(channels)):
is_last_layer = (n + 1) == len(channels)
current = channels[n]
if padding == 'same':
p = div_shape(convolution_kernels[n], 2)
else:
p = padding[n]
ops = []
for r in range(convolution_repeats[n]):
is_last_repetition = (r + 1) == convolution_repeats[n]
if is_last_repetition:
stride = strides[n]
else:
stride = 1
if last_layer_is_output and is_last_layer and is_last_repetition:
# Last layer layer should not have dropout/normalization/activation
config.activation = None
config.norm = None
config.dropout = None
ops.append(
conv_block_fn(config,
prev,
current,
kernel_size=convolution_kernels[n],
padding=p,
stride=stride))
prev = current
if pooling_size is not None:
ops.append(BlockPool(config, kernel_size=pooling_size[n]))
if not is_last_layer or not last_layer_is_output:
if dropout_probability is not None and config.dropout is not None:
ops.append(
config.dropout(p=dropout_probability,
**config.dropout_kwargs))
layers.append(nn.Sequential(*ops))
self.layers = layers
def __init__(
self,
dim: int,
input_channels: int,
channels: Sequence[int],
output_channels: int,
down_block_fn: DownType = Down,
up_block_fn: UpType = Up,
init_block_fn: ConvBlockType = BlockConvNormActivation,
middle_block_fn: MiddleType = LatentConv,
output_block_fn: ConvBlockType = BlockConvNormActivation,
init_block_channels: Optional[int] = None,
latent_channels: Optional[int] = None,
kernel_size: Optional[int] = 3,
strides: Union[int, Sequence[int]] = 2,
activation: Optional[Any] = None,
config: LayerConfig = default_layer_config(dimensionality=None)):
"""
The original UNet architecture is decomposed in 5 blocks:
- init_block_fn: will first be applied on the input
- down_block_fn: the encoder
- middle_block_fn: the junction between the encoder and decoder
- up_block_fn: the decoder
- output_block_fn: the output layer
Args:
dim: the dimensionality of the UNet (e.g., dim=2 for 2D images)
input_channels: the number of channels of the input
output_channels: the number of channels of the output
down_block_fn: a function taking (dim, in_channels, out_channels) and returning a nn.Module
up_block_fn: a function taking (dim, in_channels, out_channels, bilinear) and returning a nn.Module
init_block_channels: the number of channels to be used by the init block. If `None`, `channels[0] // 2`
will be used
"""
assert len(channels) >= 1
config = copy.copy(config)
config.set_dim(dim)
if kernel_size is not None:
config.conv_kwargs['kernel_size'] = kernel_size
config.deconv_kwargs['kernel_size'] = kernel_size
if activation is not None:
config.activation = activation
super().__init__()
self.dim = dim
self.input_channels = input_channels
self.channels = channels
self.init_block_channels = init_block_channels
self.output_channels = output_channels
self.latent_channels = latent_channels
if isinstance(strides, int):
strides = [strides] * len(channels)
assert len(strides) == len(channels), f'expected one stride per channel-layer.' \
f'Got={len(strides)} expected={len(channels)}'
if latent_channels is not None:
assert middle_block_fn is not None
self._build(config, init_block_fn, down_block_fn, up_block_fn,
middle_block_fn, output_block_fn, strides)