def __init__(self, in_channels:int, num_filters:int, filter_lengths:Sequence[int], subsample_lengths:Sequence[int], dropout:Optional[float]=None, **config) -> NoReturn:
"""
Parameters:
-----------
in_channels: int,
number of features (channels) of the input
num_filters:int,
number of filters for the convolutional layers
filter_lengths: sequence of int,
filter length (kernel size) of each convolutional layer
subsample_lengths: sequence of int,
subsample length (stride) of each convolutional layer
dropout: float, optional,
if positive, a `Dropout` layer will be introduced with this dropout probability
config: dict,
other hyper-parameters, including
filter length (kernel size), activation choices, weight initializer, etc.
"""
super().__init__()
self.__num_convs = len(filter_lengths)
self.__in_channels = in_channels
self.__out_channels = num_filters
self.__dropout = dropout or 0.0
self.config = deepcopy(config)
if self.__DEBUG__:
print(f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}")
for idx, (kernel_size, stride) in enumerate(zip(filter_lengths[:-1], subsample_lengths[:-1])):
self.add_module(
f"baby_{idx+1}",
Conv_Bn_Activation(
self.__in_channels, self.__out_channels,
kernel_size=kernel_size,
stride=stride,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
batch_norm=self.config.batch_norm,
)
)
self.add_module(
"giant",
Conv_Bn_Activation(
self.__out_channels, self.__out_channels,
kernel_size=filter_lengths[-1],
stride=subsample_lengths[-1],
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
batch_norm=self.config.batch_norm,
)
)
if self.__dropout > 0:
self.add_module(
"dropout",
nn.Dropout(self.__dropout),
)
def __init__(self, in_channels: int, **config) -> NoReturn:
""" NOT finished, NOT checked,
Parameters:
-----------
in_channels: int,
number of channels in the input
config: dict,
other hyper-parameters of the Module, including
number of convolutional layers, number of filters for each layer, etc.
"""
super().__init__()
self.__in_channels = in_channels
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of ResNetStanford is as follows\n{dict_to_str(self.config)}"
)
self.add_module(
"cba_1",
Conv_Bn_Activation(
in_channels=self.__in_channels,
out_channels=self.config.num_filters_start,
kernel_size=self.config.filter_lengths,
stride=1,
batch_norm=True,
activation=self.config.block.activation,
kw_activation=self.config.block.kw_activation,
kernel_initializer=self.config.block.kernel_initializer,
kw_initializer=self.config.block.kw_initializer,
))
module_in_channels = self.config.num_filters_start
for idx, subsample_length in enumerate(self.config.subsample_lengths):
num_filters = self.get_num_filters_at_index(
idx, self.config.num_filters_start)
self.add_module(
f"resnet_block_{idx}",
ResNetStanfordBlock(
block_index=idx,
in_channels=module_in_channels,
num_filters=num_filters,
filter_length=self.config.filter_lengths,
subsample_length=subsample_length,
**(self.config.block),
))
module_in_channels = num_filters
def __init__(self, classes: Sequence[str], n_leads: int,
config: dict) -> NoReturn:
""" finished, checked,
Parameters:
-----------
classes: sequence of int,
name of the classes
n_leads: int,
number of input leads
config: dict,
other hyper-parameters, including kernel sizes, etc.
ref. the corresponding config file
"""
super().__init__()
self.classes = list(classes)
self.n_classes = len(classes) # final out_channels
self.__in_channels = n_leads
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
__debug_seq_len = 4000
self.init_conv = DoubleConv(
in_channels=self.__in_channels,
out_channels=self.n_classes,
filter_lengths=self.config.init_filter_length,
subsample_lengths=1,
groups=self.config.groups,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.init_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, init_conv output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.down_blocks = nn.ModuleDict()
in_channels = self.n_classes
for idx in range(self.config.down_up_block_num):
self.down_blocks[f'down_{idx}'] = \
DownDoubleConv(
down_scale=self.config.down_scales[idx],
in_channels=in_channels,
out_channels=self.config.down_num_filters[idx],
filter_lengths=self.config.down_filter_lengths[idx],
groups=self.config.groups,
mode=self.config.down_mode,
**(self.config.down_block)
)
in_channels = self.config.down_num_filters[idx]
if self.__DEBUG__:
__debug_output_shape = self.down_blocks[
f'down_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, down_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.up_blocks = nn.ModuleDict()
in_channels = self.config.down_num_filters[-1]
for idx in range(self.config.down_up_block_num):
self.up_blocks[f'up_{idx}'] = \
UpDoubleConv(
up_scale=self.config.up_scales[idx],
in_channels=in_channels,
out_channels=self.config.up_num_filters[idx],
filter_lengths=self.config.up_conv_filter_lengths[idx],
deconv_filter_length=self.config.up_deconv_filter_lengths[idx],
groups=self.config.groups,
mode=self.config.up_mode,
**(self.config.up_block)
)
in_channels = self.config.up_num_filters[idx]
if self.__DEBUG__:
__debug_output_shape = self.up_blocks[
f'up_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, up_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.out_conv = Conv_Bn_Activation(
in_channels=self.config.up_num_filters[-1],
out_channels=self.n_classes,
kernel_size=self.config.out_filter_length,
stride=1,
groups=self.config.groups,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.out_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, out_conv output shape = {__debug_output_shape}"
)
class ECG_UNET(nn.Module):
"""
"""
__DEBUG__ = True
__name__ = "ECG_UNET"
def __init__(self, classes: Sequence[str], n_leads: int,
config: dict) -> NoReturn:
""" finished, checked,
Parameters:
-----------
classes: sequence of int,
name of the classes
n_leads: int,
number of input leads
config: dict,
other hyper-parameters, including kernel sizes, etc.
ref. the corresponding config file
"""
super().__init__()
self.classes = list(classes)
self.n_classes = len(classes) # final out_channels
self.__in_channels = n_leads
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
__debug_seq_len = 4000
self.init_conv = DoubleConv(
in_channels=self.__in_channels,
out_channels=self.n_classes,
filter_lengths=self.config.init_filter_length,
subsample_lengths=1,
groups=self.config.groups,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.init_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, init_conv output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.down_blocks = nn.ModuleDict()
in_channels = self.n_classes
for idx in range(self.config.down_up_block_num):
self.down_blocks[f'down_{idx}'] = \
DownDoubleConv(
down_scale=self.config.down_scales[idx],
in_channels=in_channels,
out_channels=self.config.down_num_filters[idx],
filter_lengths=self.config.down_filter_lengths[idx],
groups=self.config.groups,
mode=self.config.down_mode,
**(self.config.down_block)
)
in_channels = self.config.down_num_filters[idx]
if self.__DEBUG__:
__debug_output_shape = self.down_blocks[
f'down_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, down_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.up_blocks = nn.ModuleDict()
in_channels = self.config.down_num_filters[-1]
for idx in range(self.config.down_up_block_num):
self.up_blocks[f'up_{idx}'] = \
UpDoubleConv(
up_scale=self.config.up_scales[idx],
in_channels=in_channels,
out_channels=self.config.up_num_filters[idx],
filter_lengths=self.config.up_conv_filter_lengths[idx],
deconv_filter_length=self.config.up_deconv_filter_lengths[idx],
groups=self.config.groups,
mode=self.config.up_mode,
**(self.config.up_block)
)
in_channels = self.config.up_num_filters[idx]
if self.__DEBUG__:
__debug_output_shape = self.up_blocks[
f'up_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, up_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.out_conv = Conv_Bn_Activation(
in_channels=self.config.up_num_filters[-1],
out_channels=self.n_classes,
kernel_size=self.config.out_filter_length,
stride=1,
groups=self.config.groups,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.out_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, out_conv output shape = {__debug_output_shape}"
)
def forward(self, input: Tensor) -> Tensor:
"""
Parameters:
-----------
input: Tensor,
of shape (batch_size, channels, seq_len)
"""
to_concat = [self.init_conv(input)]
if self.__DEBUG__:
print(f"shape of init conv block output = {to_concat[-1].shape}")
for idx in range(self.config.down_up_block_num):
to_concat.append(self.down_blocks[f"down_{idx}"](to_concat[-1]))
if self.__DEBUG__:
print(
f"shape of {idx}-th down block output = {to_concat[-1].shape}"
)
up_input = to_concat[-1]
to_concat = to_concat[-2::-1]
for idx in range(self.config.down_up_block_num):
up_output = self.up_blocks[f"up_{idx}"](up_input, to_concat[idx])
up_input = up_output
if self.__DEBUG__:
print(f"shape of {idx}-th up block output = {up_output.shape}")
output = self.out_conv(up_output)
if self.__DEBUG__:
print(f"shape of out_conv layer output = {output.shape}")
return output
def inference(self, input: Tensor) -> Tensor:
"""
"""
raise NotImplementedError
def compute_output_shape(self,
seq_len: int,
batch_size: Optional[int] = None
) -> Sequence[Union[int, type(None)]]:
""" finished, checked,
Parameters:
-----------
seq_len: int,
length of the 1d sequence
batch_size: int, optional,
the batch size, can be None
Returns:
--------
output_shape: sequence,
the output shape of this `ECG_UNET` layer, given `seq_len` and `batch_size`
"""
output_shape = (batch_size, self.n_classes, seq_len)
return output_shape
@property
def module_size(self) -> int:
"""
"""
module_parameters = filter(lambda p: p.requires_grad,
self.parameters())
n_params = sum([np.prod(p.size()) for p in module_parameters])
return n_params
def __init__(self,
num_convs: int,
in_channels: int,
out_channels: int,
groups: int = 1,
**config) -> NoReturn:
""" finished, checked,
Parameters:
-----------
num_convs: int,
number of convolutional layers of this block
in_channels: int,
number of channels in the input
out_channels: int,
number of channels produced by the convolutional layers
groups: int, default 1,
connection pattern (of channels) of the inputs and outputs
config: dict,
other parameters, including
filter length (kernel size), activation choices,
weight initializer, batch normalization choices, etc. for the convolutional layers;
and pool size for the pooling layer
"""
super().__init__()
self.__num_convs = num_convs
self.__in_channels = in_channels
self.__out_channels = out_channels
self.__groups = groups
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
self.add_module(
"cba_1",
Conv_Bn_Activation(
in_channels,
out_channels,
kernel_size=self.config.filter_length,
stride=self.config.subsample_length,
groups=self.__groups,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
batch_norm=self.config.batch_norm,
))
for idx in range(num_convs - 1):
self.add_module(
f"cba_{idx+2}",
Conv_Bn_Activation(
out_channels,
out_channels,
kernel_size=self.config.filter_length,
stride=self.config.subsample_length,
groups=self.__groups,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
batch_norm=self.config.batch_norm,
))
self.add_module(
"max_pool",
nn.MaxPool1d(self.config.pool_size, self.config.pool_stride))
def __init__(self, in_channels: int, **config) -> NoReturn:
""" finished, checked,
Parameters:
-----------
in_channels: int,
number of channels in the input
config: dict,
other hyper-parameters of the Module, ref. corresponding config file
"""
super().__init__()
self.__in_channels = in_channels
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
# self.__building_block = \
# ResNetBasicBlock if self.config.name == 'resnet' else ResNetBottleNeck
self.add_module(
"init_cba",
Conv_Bn_Activation(
in_channels=self.__in_channels,
out_channels=self.config.init_num_filters,
kernel_size=self.config.init_filter_length,
stride=self.config.init_conv_stride,
groups=self.config.groups,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
bias=self.config.bias,
))
if self.config.init_pool_size > 0:
self.add_module(
"init_pool",
nn.MaxPool1d(
kernel_size=self.config.init_pool_size,
stride=self.config.init_pool_stride,
padding=(self.config.init_pool_size - 1) // 2,
))
if isinstance(self.config.filter_lengths, int):
self.__filter_lengths = \
list(repeat(self.config.filter_lengths, len(self.config.num_blocks)))
else:
self.__filter_lengths = self.config.filter_lengths
assert len(self.__filter_lengths) == len(self.config.num_blocks), \
f"`config.filter_lengths` indicates {len(self.__filter_lengths)} macro blocks, while `config.num_blocks` indicates {len(self.config.num_blocks)}"
if isinstance(self.config.subsample_lengths, int):
self.__subsample_lengths = \
list(repeat(self.config.subsample_lengths, len(self.config.num_blocks)))
else:
self.__subsample_lengths = self.config.subsample_lengths
assert len(self.__subsample_lengths) == len(self.config.num_blocks), \
f"`config.subsample_lengths` indicates {len(self.__subsample_lengths)} macro blocks, while `config.num_blocks` indicates {len(self.config.num_blocks)}"
# grouped resnet (basic) blocks,
# number of channels are doubled at the first block of each macro-block
for macro_idx, nb in enumerate(self.config.num_blocks):
macro_in_channels = (2**macro_idx) * self.config.init_num_filters
macro_filter_lengths = self.__filter_lengths[macro_idx]
macro_subsample_lengths = self.__subsample_lengths[macro_idx]
block_in_channels = macro_in_channels
block_num_filters = 2 * block_in_channels
if isinstance(macro_filter_lengths, int):
block_filter_lengths = list(repeat(macro_filter_lengths, nb))
else:
block_filter_lengths = macro_filter_lengths
assert len(block_filter_lengths) == nb, \
f"at the {macro_idx}-th macro block, `macro_subsample_lengths` indicates {len(macro_subsample_lengths)} building blocks, while `config.num_blocks[{macro_idx}]` indicates {nb}"
if isinstance(macro_subsample_lengths, int):
block_subsample_lengths = list(repeat(1, nb))
block_subsample_lengths[-1] = macro_subsample_lengths
else:
block_subsample_lengths = macro_subsample_lengths
assert len(block_subsample_lengths) == nb, \
f"at the {macro_idx}-th macro block, `macro_subsample_lengths` indicates {len(macro_subsample_lengths)} building blocks, while `config.num_blocks[{macro_idx}]` indicates {nb}"
for block_idx in range(nb):
self.add_module(
f"block_{macro_idx}_{block_idx}",
self.building_block(
in_channels=block_in_channels,
num_filters=block_num_filters,
filter_length=block_filter_lengths[block_idx],
subsample_length=block_subsample_lengths[block_idx],
groups=self.config.groups,
dilation=1,
**(self.config.block)))
block_in_channels = block_num_filters
def __init__(self,
in_channels: int,
num_filters: int,
filter_length: int,
subsample_length: int,
groups: int = 1,
dilation: int = 1,
**config) -> NoReturn:
""" finished, checked,
Parameters:
-----------
in_channels: int,
number of features (channels) of the input
num_filters: int,
number of filters for the convolutional layers
filter_length: int,
length (size) of the filter kernels
subsample_lengths: int,
subsample length,
including pool size for short cut, and stride for the top convolutional layer
groups: int, default 1,
pattern of connections between inputs and outputs,
for more details, ref. `nn.Conv1d`
dilation: int, default 1,
not used
config: dict,
other hyper-parameters, including
filter length (kernel size), activation choices, weight initializer,
and short cut patterns, etc.
"""
super().__init__()
if dilation > 1:
raise NotImplementedError(
f"Dilation > 1 not supported in {self.__name__}")
self.__num_convs = 2
self.__in_channels = in_channels
self.__out_channels = num_filters
self.__kernel_size = filter_length
self.__down_scale = subsample_length
self.__stride = subsample_length
self.__groups = groups
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
if self.config.increase_channels_method.lower(
) == 'zero_padding' and self.__groups != 1:
raise ValueError(
"zero padding for increasing channels can not be used with groups != 1"
)
self.__increase_channels = (self.__out_channels > self.__in_channels)
self.shortcut = self._make_shortcut_layer()
self.main_stream = nn.Sequential()
conv_in_channels = self.__in_channels
for i in range(self.__num_convs):
conv_activation = (self.config.activation
if i < self.__num_convs - 1 else None)
self.main_stream.add_module(
f"cba_{i}",
Conv_Bn_Activation(
in_channels=conv_in_channels,
out_channels=self.__out_channels,
kernel_size=self.__kernel_size,
stride=(self.__stride if i == 0 else 1),
groups=self.__groups,
batch_norm=True,
activation=conv_activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
bias=self.config.bias,
))
conv_in_channels = self.__out_channels
if isinstance(self.config.activation, str):
self.out_activation = \
Activations[self.config.activation.lower()](**self.config.kw_activation)
else:
self.out_activation = \
self.config.activation(**self.config.kw_activation)
def __init__(self, classes: Sequence[str], n_leads: int,
config: dict) -> NoReturn:
""" finished, checked,
Parameters:
-----------
classes: sequence of int,
name of the classes
n_leads: int,
number of input leads
config: dict,
other hyper-parameters, including kernel sizes, etc.
ref. the corresponding config file
"""
super().__init__()
self.classes = list(classes)
self.n_classes = len(classes) # final out_channels if > 2
self.__out_channels = len(classes) if len(classes) > 2 else 1
self.__in_channels = n_leads
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
__debug_seq_len = 5000
# TODO: an init batch normalization?
if self.config.init_batch_norm:
self.init_bn = nn.BatchNorm1d(
num_features=self.__in_channels,
eps=1e-5, # default val
momentum=0.1, # default val
)
self.init_conv = TripleConv(
in_channels=self.__in_channels,
out_channels=self.config.init_num_filters,
filter_lengths=self.config.init_filter_length,
subsample_lengths=1,
groups=self.config.groups,
dropouts=self.config.init_dropouts,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.init_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, init_conv output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.down_blocks = nn.ModuleDict()
in_channels = self.config.init_num_filters
for idx in range(self.config.down_up_block_num - 1):
self.down_blocks[f'down_{idx}'] = \
DownTripleConv(
down_scale=self.config.down_scales[idx],
in_channels=in_channels,
out_channels=self.config.down_num_filters[idx],
filter_lengths=self.config.down_filter_lengths[idx],
groups=self.config.groups,
dropouts=self.config.down_dropouts[idx],
mode=self.config.down_mode,
**(self.config.down_block)
)
in_channels = self.config.down_num_filters[idx][-1]
if self.__DEBUG__:
__debug_output_shape = self.down_blocks[
f'down_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, down_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.bottom_block = DownBranchedDoubleConv(
down_scale=self.config.down_scales[-1],
in_channels=in_channels,
out_channels=self.config.bottom_num_filters,
filter_lengths=self.config.bottom_filter_lengths,
dilations=self.config.bottom_dilations,
groups=self.config.groups,
dropouts=self.config.bottom_dropouts,
mode=self.config.down_mode,
**(self.config.down_block))
if self.__DEBUG__:
__debug_output_shape = self.bottom_block.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, bottom_block output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.up_blocks = nn.ModuleDict()
in_channels = sum(
[branch[-1] for branch in self.config.bottom_num_filters])
for idx in range(self.config.down_up_block_num):
self.up_blocks[f'up_{idx}'] = \
UpTripleConv(
up_scale=self.config.up_scales[idx],
in_channels=in_channels,
out_channels=self.config.up_num_filters[idx],
filter_lengths=self.config.up_conv_filter_lengths[idx],
deconv_filter_length=self.config.up_deconv_filter_lengths[idx],
groups=self.config.groups,
mode=self.config.up_mode,
dropouts=self.config.up_dropouts[idx],
**(self.config.up_block)
)
in_channels = self.config.up_num_filters[idx][-1]
if self.__DEBUG__:
__debug_output_shape = self.up_blocks[
f'up_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, up_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.out_conv = Conv_Bn_Activation(
in_channels=self.config.up_num_filters[-1][-1],
out_channels=self.__out_channels,
kernel_size=self.config.out_filter_length,
stride=1,
groups=self.config.groups,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.out_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, out_conv output shape = {__debug_output_shape}"
)
class ECG_SUBTRACT_UNET(nn.Module):
"""
"""
__DEBUG__ = True
__name__ = "ECG_SUBTRACT_UNET"
def __init__(self, classes: Sequence[str], n_leads: int,
config: dict) -> NoReturn:
""" finished, checked,
Parameters:
-----------
classes: sequence of int,
name of the classes
n_leads: int,
number of input leads
config: dict,
other hyper-parameters, including kernel sizes, etc.
ref. the corresponding config file
"""
super().__init__()
self.classes = list(classes)
self.n_classes = len(classes) # final out_channels if > 2
self.__out_channels = len(classes) if len(classes) > 2 else 1
self.__in_channels = n_leads
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
__debug_seq_len = 5000
# TODO: an init batch normalization?
if self.config.init_batch_norm:
self.init_bn = nn.BatchNorm1d(
num_features=self.__in_channels,
eps=1e-5, # default val
momentum=0.1, # default val
)
self.init_conv = TripleConv(
in_channels=self.__in_channels,
out_channels=self.config.init_num_filters,
filter_lengths=self.config.init_filter_length,
subsample_lengths=1,
groups=self.config.groups,
dropouts=self.config.init_dropouts,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.init_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, init_conv output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.down_blocks = nn.ModuleDict()
in_channels = self.config.init_num_filters
for idx in range(self.config.down_up_block_num - 1):
self.down_blocks[f'down_{idx}'] = \
DownTripleConv(
down_scale=self.config.down_scales[idx],
in_channels=in_channels,
out_channels=self.config.down_num_filters[idx],
filter_lengths=self.config.down_filter_lengths[idx],
groups=self.config.groups,
dropouts=self.config.down_dropouts[idx],
mode=self.config.down_mode,
**(self.config.down_block)
)
in_channels = self.config.down_num_filters[idx][-1]
if self.__DEBUG__:
__debug_output_shape = self.down_blocks[
f'down_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, down_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.bottom_block = DownBranchedDoubleConv(
down_scale=self.config.down_scales[-1],
in_channels=in_channels,
out_channels=self.config.bottom_num_filters,
filter_lengths=self.config.bottom_filter_lengths,
dilations=self.config.bottom_dilations,
groups=self.config.groups,
dropouts=self.config.bottom_dropouts,
mode=self.config.down_mode,
**(self.config.down_block))
if self.__DEBUG__:
__debug_output_shape = self.bottom_block.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, bottom_block output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.up_blocks = nn.ModuleDict()
in_channels = sum(
[branch[-1] for branch in self.config.bottom_num_filters])
for idx in range(self.config.down_up_block_num):
self.up_blocks[f'up_{idx}'] = \
UpTripleConv(
up_scale=self.config.up_scales[idx],
in_channels=in_channels,
out_channels=self.config.up_num_filters[idx],
filter_lengths=self.config.up_conv_filter_lengths[idx],
deconv_filter_length=self.config.up_deconv_filter_lengths[idx],
groups=self.config.groups,
mode=self.config.up_mode,
dropouts=self.config.up_dropouts[idx],
**(self.config.up_block)
)
in_channels = self.config.up_num_filters[idx][-1]
if self.__DEBUG__:
__debug_output_shape = self.up_blocks[
f'up_{idx}'].compute_output_shape(__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, up_{idx} output shape = {__debug_output_shape}"
)
_, _, __debug_seq_len = __debug_output_shape
self.out_conv = Conv_Bn_Activation(
in_channels=self.config.up_num_filters[-1][-1],
out_channels=self.__out_channels,
kernel_size=self.config.out_filter_length,
stride=1,
groups=self.config.groups,
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
)
if self.__DEBUG__:
__debug_output_shape = self.out_conv.compute_output_shape(
__debug_seq_len)
print(
f"given seq_len = {__debug_seq_len}, out_conv output shape = {__debug_output_shape}"
)
def forward(self, input: Tensor) -> Tensor:
"""
"""
if self.config.init_batch_norm:
x = self.init_bn(input)
else:
x = input
# down
to_concat = [self.init_conv(x)]
if self.__DEBUG__:
print(f"shape of init conv block output = {to_concat[-1].shape}")
for idx in range(self.config.down_up_block_num - 1):
to_concat.append(self.down_blocks[f"down_{idx}"](to_concat[-1]))
if self.__DEBUG__:
print(
f"shape of {idx}-th down block output = {to_concat[-1].shape}"
)
to_concat.append(self.bottom_block(to_concat[-1]))
# up
up_input = to_concat[-1]
to_concat = to_concat[-2::-1]
for idx in range(self.config.down_up_block_num):
up_output = self.up_blocks[f"up_{idx}"](up_input, to_concat[idx])
up_input = up_output
if self.__DEBUG__:
print(f"shape of {idx}-th up block output = {up_output.shape}")
# output
output = self.out_conv(up_output)
if self.__DEBUG__:
print(f"shape of out_conv layer output = {output.shape}")
return output
def inference(self, input: Tensor) -> Tensor:
"""
"""
raise NotImplementedError
def compute_output_shape(self,
seq_len: int,
batch_size: Optional[int] = None
) -> Sequence[Union[int, type(None)]]:
""" finished, NOT checked,
Parameters:
-----------
seq_len: int,
length of the 1d sequence
batch_size: int, optional,
the batch size, can be None
Returns:
--------
output_shape: sequence,
the output shape of this `ECG_UNET` layer, given `seq_len` and `batch_size`
"""
output_shape = (batch_size, self.n_classes, seq_len)
return output_shape
def __init__(self,
block_index: int,
in_channels: int,
num_filters: int,
filter_length: int,
subsample_length: int,
dilation: int = 1,
**config) -> NoReturn:
""" finished, checked,
the main stream uses `subsample_length` as stride to perform down-sampling,
the short cut uses `subsample_length` as pool size to perform down-sampling,
Parameters:
-----------
block_index: int,
index of the block in the whole sequence of `ResNetStanford`
in_channels: int,
number of features (channels) of the input
num_filters: int,
number of filters for the convolutional layers
filter_length: int,
length (size) of the filter kernels
subsample_length: int,
subsample length,
including pool size for short cut, and stride for the top convolutional layer
config: dict,
other hyper-parameters, including
filter length (kernel size), activation choices, weight initializer, dropout,
and short cut patterns, etc.
Issues:
-------
1. even kernel size would case mismatch of shapes of main stream and short cut
"""
super().__init__()
self.__block_index = block_index
self.__in_channels = in_channels
self.__out_channels = num_filters
self.__kernel_size = filter_length
self.__down_scale = subsample_length
self.__stride = subsample_length
self.config = ED(deepcopy(config))
self.__num_convs = self.config.num_skip
self.__increase_channels = (self.__out_channels > self.__in_channels)
self.short_cut = self._make_short_cut_layer()
self.main_stream = nn.Sequential()
conv_in_channels = self.__in_channels
for i in range(self.__num_convs):
if not (block_index == 0 and i == 0):
self.main_stream.add_module(
f"ba_{self.__block_index}_{i}",
Bn_Activation(
num_features=self.__in_channels,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
dropout=self.config.dropout if i > 0 else 0,
),
)
self.main_stream.add_module(
f"conv_{self.__block_index}_{i}",
Conv_Bn_Activation(
in_channels=conv_in_channels,
out_channels=self.__out_channels,
kernel_size=self.__kernel_size,
stride=(self.__stride if i == 0 else 1),
batch_norm=False,
activation=None,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
))
conv_in_channels = self.__out_channels
def __init__(self,
in_channels: int,
out_channels: int,
filter_lengths: Union[Sequence[int], int],
subsample_lengths: Union[Sequence[int], int] = 1,
groups: int = 1,
mid_channels: Optional[int] = None,
**config) -> NoReturn:
""" finished, checked,
Parameters:
-----------
in_channels: int,
number of channels in the input
out_channels: int,
number of channels produced by the last convolutional layer
filter_lengths: int or sequence of int,
length(s) of the filters (kernel size)
subsample_lengths: int or sequence of int,
subsample length(s) (stride(s)) of the convolutions
groups: int, default 1,
connection pattern (of channels) of the inputs and outputs
mid_channels: int, optional,
number of channels produced by the first convolutional layer,
defaults to `out_channels`
config: dict,
other parameters, including
activation choices, weight initializer, batch normalization choices, etc.
for the convolutional layers
"""
super().__init__()
self.__num_convs = 2
self.__in_channels = in_channels
self.__mid_channels = mid_channels if mid_channels is not None else out_channels
self.__out_channels = out_channels
self.config = ED(deepcopy(config))
if self.__DEBUG__:
print(
f"configuration of {self.__name__} is as follows\n{dict_to_str(self.config)}"
)
if isinstance(filter_lengths, int):
kernel_sizes = list(repeat(filter_lengths, self.__num_convs))
else:
kernel_sizes = filter_lengths
assert len(kernel_sizes) == self.__num_convs
if isinstance(subsample_lengths, int):
strides = list(repeat(subsample_lengths, self.__num_convs))
else:
strides = subsample_lengths
assert len(strides) == self.__num_convs
self.add_module(
"cba_1",
Conv_Bn_Activation(
in_channels=self.__in_channels,
out_channels=self.__mid_channels,
kernel_size=kernel_sizes[0],
stride=strides[0],
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
),
)
self.add_module(
"cba_2",
Conv_Bn_Activation(
in_channels=self.__mid_channels,
out_channels=self.__out_channels,
kernel_size=kernel_sizes[1],
stride=strides[1],
batch_norm=self.config.batch_norm,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
))
def __init__(self,
in_channels: int,
scopes: Sequence[int],
num_filters: Union[int, Sequence[int]],
filter_lengths: Union[int, Sequence[int]],
subsample_length: int,
groups: int = 1,
**config) -> NoReturn:
""" finished, not checked,
Parameters:
-----------
in_channels: int,
number of channels in the input
scopes: sequence of int,
scopes of the convolutional layers, via `dilation`
num_filters: int or sequence of int,
filter_lengths: int or sequence of int,
subsample_length: int,
"""
super().__init__()
self.__in_channels = in_channels
self.__scopes = scopes
self.__num_convs = len(self.__scopes)
if isinstance(num_filters, int):
self.__out_channels = list(repeat(num_filters, self.__num_convs))
else:
self.__out_channels = num_filters
assert len(self.__out_channels) == self.__num_convs, \
f"`scopes` indicates {self.__num_convs} convolutional layers, while `num_filters` indicates {len(self.__out_channels)}"
if isinstance(filter_lengths, int):
self.__filter_lengths = list(
repeat(filter_lengths, self.__num_convs))
else:
self.__filter_lengths = filter_lengths
assert len(self.__filter_lengths) == self.__num_convs, \
f"`scopes` indicates {self.__num_convs} convolutional layers, while `filter_lengths` indicates {len(self.__filter_lengths)}"
self.__subsample_length = subsample_length
self.__groups = groups
self.config = ED(deepcopy(config))
conv_in_channels = self.__in_channels
for idx in range(self.__num_convs):
self.add_module(
f"ca_{idx}",
Conv_Bn_Activation(
in_channels=conv_in_channels,
out_channels=self.__out_channels[idx],
kernel_size=self.__filter_lengths[idx],
stride=1,
dilation=self.__scopes[idx],
groups=self.__groups,
batch_norm=False,
activation=self.config.activation,
kw_activation=self.config.kw_activation,
kernel_initializer=self.config.kernel_initializer,
kw_initializer=self.config.kw_initializer,
bias=self.config.bias,
))
conv_in_channels = self.__out_channels[idx]
self.add_module("bn", nn.BatchNorm1d(self.__out_channels[-1]))
self.add_module(
"down",
DownSample(
down_scale=self.__subsample_length,
in_channels=self.__out_channels[-1],
groups=self.__groups,
# padding=
batch_norm=False,
mode=self.config.subsample_mode,
))
if self.config.dropout > 0:
self.add_module("dropout",
nn.Dropout(self.config.dropout, inplace=False))