diff --git a/mlair/model_modules/convolutional_networks.py b/mlair/model_modules/convolutional_networks.py index 0a16be7c8079641dc9752fe5b572e9e2c12a7017..5146fe529ef01aabc2440581198d37d0d2ab6a4b 100644 --- a/mlair/model_modules/convolutional_networks.py +++ b/mlair/model_modules/convolutional_networks.py @@ -22,7 +22,7 @@ class CNN(AbstractModelClass): _requirements = ["lr", "beta_1", "beta_2", "epsilon", "decay", "amsgrad"] def __init__(self, input_shape: list, output_shape: list, activation="relu", activation_output="linear", - optimizer="adam", regularizer=None, **kwargs): + optimizer="adam", regularizer=None, kernel_size=1, **kwargs): assert len(input_shape) == 1 assert len(output_shape) == 1 @@ -35,6 +35,7 @@ class CNN(AbstractModelClass): self.activation_output_name = activation_output self.kernel_initializer = self._initializer.get(activation, "glorot_uniform") self.kernel_regularizer = self._set_regularizer(regularizer, **kwargs) + self.kernel_size = kernel_size self.optimizer = self._set_optimizer(optimizer, **kwargs) # apply to model @@ -81,7 +82,7 @@ class CNN(AbstractModelClass): Build the model. """ x_input = keras.layers.Input(shape=self._input_shape) - kernel = (5, 1) + kernel = (self.kernel_size, 1) pad_size = PadUtils.get_padding_for_same(kernel) x_in = Padding2D("SymPad2D")(padding=pad_size, name="SymPad1")(x_input) x_in = keras.layers.Conv2D(filters=16, kernel_size=kernel,