def custom_train_step(self, data):
"""
Custom training logic
:param data:
:return:
"""
data = data_adapter.expand_1d(data)
x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
with tf.GradientTape() as tape:
y_pred = self.keras_model(x, training=True)
loss = self.keras_model.compiled_loss(
y,
y_pred,
sample_weight,
regularization_losses=self.keras_model.losses)
if self.task == 'regression':
variance_loss = mse_var_wrapper(y_pred[0], x['labels_err'])
output_loss = mse_lin_wrapper(y_pred[1], x['labels_err'])
elif self.task == 'classification':
output_loss = bayesian_categorical_crossentropy_wrapper(
y_pred[1])
variance_loss = bayesian_categorical_crossentropy_var_wrapper(
y_pred[0])
elif self.task == 'binary_classification':
output_loss = bayesian_binary_crossentropy_wrapper(y_pred[1])
variance_loss = bayesian_binary_crossentropy_var_wrapper(
y_pred[0])
else:
raise RuntimeError(
'Only "regression", "classification" and "binary_classification" are supported'
)
loss = output_loss(y['output'], y_pred[0]) + variance_loss(
y['variance_output'], y_pred[1])
# apply gradient here
if version.parse(tf.__version__) >= version.parse("2.4.0"):
self.keras_model.optimizer.minimize(
loss, self.keras_model.trainable_variables, tape=tape)
else:
tf.python.keras.engine.training._minimize(
self.keras_model.distribute_strategy, tape,
self.keras_model.optimizer, loss,
self.keras_model.trainable_variables)
self.keras_model.compiled_metrics.update_state(y, y_pred,
sample_weight)
return {m.name: m.result() for m in self.keras_model.metrics}
def model(self):
input_tensor = Input(shape=self._input_shape, name='input')
labels_err_tensor = Input(shape=(self._labels_shape,), name='labels_err')
cnn_layer_1 = Conv1D(kernel_initializer=self.initializer, padding="same", filters=self.num_filters[0],
kernel_size=self.filter_len, kernel_regularizer=regularizers.l2(self.l2))(input_tensor)
activation_1 = Activation(activation=self.activation)(cnn_layer_1)
dropout_1 = MCDropout(self.dropout_rate, disable=self.disable_dropout)(activation_1)
cnn_layer_2 = Conv1D(kernel_initializer=self.initializer, padding="same", filters=self.num_filters[1],
kernel_size=self.filter_len, kernel_regularizer=regularizers.l2(self.l2))(dropout_1)
activation_2 = Activation(activation=self.activation)(cnn_layer_2)
maxpool_1 = MaxPooling1D(pool_size=self.pool_length)(activation_2)
flattener = Flatten()(maxpool_1)
dropout_2 = MCDropout(self.dropout_rate, disable=self.disable_dropout)(flattener)
layer_3 = Dense(units=self.num_hidden[0], kernel_regularizer=regularizers.l2(self.l2),
kernel_initializer=self.initializer,
activation=self.activation)(dropout_2)
activation_3 = Activation(activation=self.activation)(layer_3)
dropout_3 = MCDropout(self.dropout_rate, disable=self.disable_dropout)(activation_3)
layer_4 = Dense(units=self.num_hidden[1], kernel_regularizer=regularizers.l2(self.l2),
kernel_initializer=self.initializer,
activation=self.activation)(dropout_3)
activation_4 = Activation(activation=self.activation)(layer_4)
output = Dense(units=self._labels_shape, name='output')(activation_4)
output_activated = Activation(activation=self._last_layer_activation)(output)
variance_output = Dense(units=self._labels_shape, activation='linear', name='variance_output')(activation_4)
model = Model(inputs=[input_tensor, labels_err_tensor], outputs=[output, variance_output])
# new astroNN high performance dropout variational inference on GPU expects single output
model_prediction = Model(inputs=[input_tensor], outputs=concatenate([output, variance_output]))
if self.task == 'regression':
variance_loss = mse_var_wrapper(output, labels_err_tensor)
output_loss = mse_lin_wrapper(variance_output, labels_err_tensor)
elif self.task == 'classification':
output_loss = bayesian_categorical_crossentropy_wrapper(variance_output)
variance_loss = bayesian_categorical_crossentropy_var_wrapper(output)
elif self.task == 'binary_classification':
output_loss = bayesian_binary_crossentropy_wrapper(variance_output)
variance_loss = bayesian_binary_crossentropy_var_wrapper(output)
else:
raise RuntimeError('Only "regression", "classification" and "binary_classification" are supported')
return model, model_prediction, output_loss, variance_loss
def compile(self,
optimizer=None,
loss=None,
metrics=None,
weighted_metrics=None,
loss_weights=None,
sample_weight_mode=None):
if optimizer is not None:
self.optimizer = optimizer
elif self.optimizer is None or self.optimizer == 'adam':
self.optimizer = Adam(learning_rate=self.lr,
beta_1=self.beta_1,
beta_2=self.beta_2,
epsilon=self.optimizer_epsilon,
decay=0.0)
if metrics is not None:
self.metrics = metrics
if self.task == 'regression':
if self._last_layer_activation is None:
self._last_layer_activation = 'linear'
elif self.task == 'classification':
if self._last_layer_activation is None:
self._last_layer_activation = 'softmax'
elif self.task == 'binary_classification':
if self._last_layer_activation is None:
self._last_layer_activation = 'sigmoid'
else:
raise RuntimeError(
'Only "regression", "classification" and "binary_classification" are supported'
)
self.keras_model, self.keras_model_predict, self.output_loss, self.variance_loss = self.model(
)
if self.task == 'regression':
self._output_loss = lambda predictive, labelerr: mse_lin_wrapper(
predictive, labelerr)
elif self.task == 'classification':
self._output_loss = lambda predictive, labelerr: bayesian_categorical_crossentropy_wrapper(
predictive)
elif self.task == 'binary_classification':
self._output_loss = lambda predictive, labelerr: bayesian_binary_crossentropy_wrapper(
predictive)
else:
raise RuntimeError(
'Only "regression", "classification" and "binary_classification" are supported'
)
# all zero losss as dummy lose
if self.task == 'regression':
self.metrics = [mean_absolute_error, mean_error
] if not self.metrics else self.metrics
self.keras_model.compile(optimizer=self.optimizer,
loss=zeros_loss,
metrics=self.metrics,
weighted_metrics=weighted_metrics,
sample_weight_mode=sample_weight_mode)
elif self.task == 'classification':
self.metrics = [categorical_accuracy
] if not self.metrics else self.metrics
self.keras_model.compile(optimizer=self.optimizer,
loss=zeros_loss,
metrics={'output': self.metrics},
weighted_metrics=weighted_metrics,
sample_weight_mode=sample_weight_mode)
elif self.task == 'binary_classification':
self.metrics = [binary_accuracy
] if not self.metrics else self.metrics
self.keras_model.compile(optimizer=self.optimizer,
loss=zeros_loss,
metrics={'output': self.metrics},
weighted_metrics=weighted_metrics,
sample_weight_mode=sample_weight_mode)
# inject custom training step if needed
try:
self.custom_train_step()
except NotImplementedError:
pass
except TypeError:
self.keras_model.train_step = self.custom_train_step
# inject custom testing step if needed
try:
self.custom_test_step()
except NotImplementedError:
pass
except TypeError:
self.keras_model.test_step = self.custom_test_step
return None