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
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