# # Copyright (c) 2019 Jonathan Weyn # # See the file LICENSE for your rights. # """ Custom Keras and PyTorch classes. """ import tensorflow as tf from tensorflow.compat.v1 import keras as tfk from tensorflow.keras.callbacks import Callback, EarlyStopping from tensorflow.keras.layers import ZeroPadding2D, ZeroPadding3D, LocallyConnected2D, Lambda, Layer from tensorflow.keras.losses import mean_absolute_error, mean_squared_error from tensorflow.python.keras.utils import conv_utils from tensorflow.python.keras.engine.base_layer import InputSpec from tensorflow.keras import activations, initializers, regularizers, constraints import numpy as np try: from s2cnn import S2Convolution, SO3Convolution except ImportError: pass K = tfk.backend # ==================================================================================================================== # # Keras utility classes # ==================================================================================================================== # class AdamLearningRateTracker(Callback): """ Log the current learning rate used by and Adam optimizer. """ def on_epoch_end(self, epoch, logs=None, beta_1=0.9, beta_2=0.999, ): optimizer = self.model.optimizer it = K.cast(optimizer.iterations, K.floatx()) lr = K.cast(optimizer.lr, K.floatx()) decay = K.cast(optimizer.decay, K.floatx()) t = K.eval(it + 1.) new_lr = K.eval(lr * (1. / (1. + decay * it))) lr_t = K.eval(new_lr * (K.sqrt(1. - K.pow(beta_2, t)) / (1. - K.pow(beta_1, t)))) print(' - LR: {:.6f}'.format(lr_t)) class SGDLearningRateTracker(Callback): """ Log the current learning rate used by an SGD optimizer. """ def on_epoch_end(self, epoch, logs=None): optimizer = self.model.optimizer it = K.cast(optimizer.iterations, K.floatx()) lr = K.cast(optimizer.lr, K.floatx()) decay = K.cast(optimizer.decay, K.floatx()) new_lr = K.eval(lr * (1. / (1. + decay * it))) print(' - LR: {:.6f}'.format(new_lr)) class BatchHistory(Callback): """ Log training metrics for each batch of training data. """ def on_train_begin(self, logs=None): self.history = [] self.epoch = 0 def on_epoch_begin(self, epoch, logs=None): self.history.append({}) def on_epoch_end(self, epoch, logs=None): self.epoch += 1 def on_batch_end(self, batch, logs=None): logs = logs or {} for k, v in logs.items(): self.history[self.epoch].setdefault(k, []).append(v) class RunHistory(Callback): """ Callback that records events into a `History` object. Adapted from keras.callbacks.History to include logging to Azure experiment runs. """ def __init__(self, run): super(RunHistory, self).__init__() self.epoch = [] self.history = {} self.run = run def on_train_begin(self, logs=None): self.epoch = [] self.history = {} def on_epoch_end(self, epoch, logs=None): logs = logs or {} self.epoch.append(epoch) for k, v in logs.items(): self.history.setdefault(k, []).append(v) self.run.log(k, v) class RNNResetStates(Callback): def on_epoch_begin(self, epoch, logs=None): self.model.reset_states() class EarlyStoppingMin(EarlyStopping): """ Extends the keras.callbacks.EarlyStopping class to provide the option to force training for a minimum number of epochs or restore the best weights after the maximum epochs have been reached. """ def __init__(self, min_epochs=0, max_epochs=None, **kwargs): """ :param min_epochs: int: train the network for at least this number of epochs before early stopping :param max_epochs: int: train the network for at most this number of epochs before early stopping :param kwargs: passed to EarlyStopping.__init__() """ super(EarlyStoppingMin, self).__init__(**kwargs) if not isinstance(min_epochs, int) or min_epochs < 0: raise ValueError('min_epochs must be an integer >= 0') self.min_epochs = int(min_epochs) self.max_epochs = int(max_epochs) if max_epochs is not None else None def on_epoch_end(self, epoch, logs=None): if epoch < self.min_epochs: return current = self.get_monitor_value(logs) if current is None: return if self.monitor_op(current - self.min_delta, self.best): self.best = current self.wait = 0 if self.restore_best_weights: self.best_weights = self.model.get_weights() else: self.wait += 1 if self.wait >= self.patience: self.stopped_epoch = epoch self.model.stop_training = True if self.restore_best_weights: if self.verbose > 0: print('Restoring model weights from the end of ' 'the best epoch') self.model.set_weights(self.best_weights) if self.max_epochs is not None and epoch >= self.max_epochs: self.stopped_epoch = epoch self.model.stop_training = True if self.restore_best_weights: if self.verbose > 0: print('Maximum epochs reached; restoring model weights from the end of ' 'the best epoch') self.model.set_weights(self.best_weights) if self.verbose > 1: print('EarlyStoppingMin: %d epochs since last minimum' % self.wait) class SaveWeightsOnEpoch(Callback): """ Saves the model weights to a temporary file at the end of each epoch. This is useful for avoiding complete loss of a run that fails for any reason. """ def __init__(self, weights_file, interval=None): """ :param weights_file: str: file name to save weights """ super(SaveWeightsOnEpoch, self).__init__() self.weights_file = str(weights_file) if interval is not None: assert isinstance(interval, int) and interval > 0, "'interval' must be an integer > 0" self.interval = interval def on_epoch_end(self, epoch, logs=None): if self.interval is not None and epoch % self.interval == 0: self.model.save_weights('%s.%s' % (self.weights_file, epoch), save_format='h5') else: try: self.model.save_weights(self.weights_file, save_format='h5') except OSError: # If the file is temporarily system locked for any reason, avoid crashing pass class GeneratorEpochEnd(Callback): """ Used in conjunction with a tensorflow.data.Dataset generator to manually execute the on_epoch_end() method of a DLWP.model.generators Generator instance after each epoch. """ def __init__(self, generator): """ :param generator: DLWP.model.generator instance """ super(GeneratorEpochEnd, self).__init__() self.generator = generator def on_epoch_end(self, epoch, logs=None): self.generator.on_epoch_end() # ==================================================================================================================== # # Keras padding layers # ==================================================================================================================== # class PeriodicPadding2D(ZeroPadding2D): """Periodic-padding layer for 2D input (e.g. image). This layer can add periodic rows and columns at the top, bottom, left and right side of an image tensor. Adapted from keras.layers.ZeroPadding2D by @jweyn # Arguments padding: int, or tuple of 2 ints, or tuple of 2 tuples of 2 ints. - If int: the same symmetric padding is applied to height and width. - If tuple of 2 ints: interpreted as two different symmetric padding values for height and width: `(symmetric_height_pad, symmetric_width_pad)`. - If tuple of 2 tuples of 2 ints: interpreted as `((top_pad, bottom_pad), (left_pad, right_pad))` data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, height, width, channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, height, width)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". # Input shape 4D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, rows, cols, channels)` - If `data_format` is `"channels_first"`: `(batch, channels, rows, cols)` # Output shape 4D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, padded_rows, padded_cols, channels)` - If `data_format` is `"channels_first"`: `(batch, channels, padded_rows, padded_cols)` """ def __init__(self, padding=(1, 1), data_format=None, **kwargs): super(PeriodicPadding2D, self).__init__(padding=padding, data_format=data_format, **kwargs) def call(self, inputs): if K.backend() == 'plaidml.keras.backend': shape = inputs.shape.dims else: shape = inputs.shape if self.data_format == 'channels_first': top_slice = slice(shape[2] - self.padding[0][0], shape[2]) bottom_slice = slice(0, self.padding[0][1]) left_slice = slice(shape[3] - self.padding[1][0], shape[3]) right_slice = slice(0, self.padding[1][1]) # Pad the horizontal outputs = K.concatenate([inputs[:, :, :, left_slice], inputs, inputs[:, :, :, right_slice]], axis=3) # Pad the vertical outputs = K.concatenate([outputs[:, :, top_slice], outputs, outputs[:, :, bottom_slice]], axis=2) else: top_slice = slice(shape[1] - self.padding[0][0], shape[1]) bottom_slice = slice(0, self.padding[0][1]) left_slice = slice(shape[2] - self.padding[1][0], shape[2]) right_slice = slice(0, self.padding[1][1]) # Pad the horizontal outputs = K.concatenate([inputs[:, :, left_slice], inputs, inputs[:, :, right_slice]], axis=2) # Pad the vertical outputs = K.concatenate([outputs[:, top_slice], outputs, outputs[:, bottom_slice]], axis=1) return outputs class PeriodicPadding3D(ZeroPadding3D): """Periodic-padding layer for 3D input (e.g. image). This layer can add periodic rows, columns, and depth to an image tensor. Adapted from keras.layers.ZeroPadding3D by @jweyn # Arguments padding: int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints. - If int: the same symmetric padding is applied to height and width. - If tuple of 3 ints: interpreted as two different symmetric padding values for height and width: `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`. - If tuple of 3 tuples of 2 ints: interpreted as `((left_dim1_pad, right_dim1_pad), (left_dim2_pad, right_dim2_pad), (left_dim3_pad, right_dim3_pad))` data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". # Input shape 5D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad, depth)` - If `data_format` is `"channels_first"`: `(batch, depth, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad)` # Output shape 5D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, first_padded_axis, second_padded_axis, third_axis_to_pad, depth)` - If `data_format` is `"channels_first"`: `(batch, depth, first_padded_axis, second_padded_axis, third_axis_to_pad)` """ def __init__(self, padding=(1, 1, 1), data_format=None, **kwargs): super(PeriodicPadding3D, self).__init__(padding=padding, data_format=data_format, **kwargs) def call(self, inputs): if K.backend() == 'plaidml.keras.backend': shape = inputs.shape.dims else: shape = inputs.shape if self.data_format == 'channels_first': low_slice = slice(shape[2] - self.padding[0][0], shape[2]) high_slice = slice(0, self.padding[0][1]) top_slice = slice(shape[3] - self.padding[1][0], shape[3]) bottom_slice = slice(0, self.padding[1][1]) left_slice = slice(shape[4] - self.padding[2][0], shape[4]) right_slice = slice(0, self.padding[2][1]) # Pad the horizontal outputs = K.concatenate([inputs[:, :, :, :, left_slice], inputs, inputs[:, :, :, :, right_slice]], axis=4) # Pad the vertical outputs = K.concatenate([outputs[:, :, :, top_slice], outputs, outputs[:, :, :, bottom_slice]], axis=3) # Pad the depth outputs = K.concatenate([outputs[:, :, low_slice], outputs, outputs[:, :, high_slice]], axis=2) else: low_slice = slice(shape[1] - self.padding[0][0], shape[1]) high_slice = slice(0, self.padding[0][1]) top_slice = slice(shape[2] - self.padding[1][0], shape[2]) bottom_slice = slice(0, self.padding[1][1]) left_slice = slice(shape[3] - self.padding[2][0], shape[3]) right_slice = slice(0, self.padding[2][1]) # Pad the horizontal outputs = K.concatenate([inputs[:, :, :, left_slice], inputs, inputs[:, :, :, right_slice]], axis=3) # Pad the vertical outputs = K.concatenate([outputs[:, :, top_slice], outputs, outputs[:, :, bottom_slice]], axis=2) # Pad the depth outputs = K.concatenate([outputs[:, low_slice], outputs, outputs[:, high_slice]], axis=1) return outputs class FillPadding2D(ZeroPadding2D): """Fill-padding layer for 2D input (e.g. image). This layer can add rows or columns that duplicate the edge values. Adapted from keras.layers.ZeroPadding2D by @jweyn # Arguments padding: int, or tuple of 2 ints, or tuple of 2 tuples of 2 ints. - If int: the same symmetric padding is applied to height and width. - If tuple of 2 ints: interpreted as two different symmetric padding values for height and width: `(symmetric_height_pad, symmetric_width_pad)`. - If tuple of 2 tuples of 2 ints: interpreted as `((top_pad, bottom_pad), (left_pad, right_pad))` data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, height, width, channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, height, width)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". # Input shape 4D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, rows, cols, channels)` - If `data_format` is `"channels_first"`: `(batch, channels, rows, cols)` # Output shape 4D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, padded_rows, padded_cols, channels)` - If `data_format` is `"channels_first"`: `(batch, channels, padded_rows, padded_cols)` """ def __init__(self, padding=(1, 1), data_format=None, **kwargs): super(FillPadding2D, self).__init__(padding=padding, data_format=data_format, **kwargs) def call(self, inputs): if self.data_format == 'channels_first': # Pad the vertical if self.padding[0][0] > 0: top_slice = K.stack([inputs[:, :, 0]] * self.padding[0][0], axis=2) else: top_slice = inputs[:, :, slice(0, 0)] if self.padding[0][1] > 0: bottom_slice = K.stack([inputs[:, :, -1]] * self.padding[0][1], axis=2) else: bottom_slice = inputs[:, :, slice(0, 0)] outputs = K.concatenate([top_slice, inputs, bottom_slice], axis=2) # Pad the horizontal if self.padding[1][0] > 0: left_slice = K.stack([outputs[:, :, :, 0]] * self.padding[1][0], axis=3) else: left_slice = outputs[:, :, :, slice(0, 0)] if self.padding[1][1] > 0: right_slice = K.stack([outputs[:, :, :, -1]] * self.padding[1][1], axis=3) else: right_slice = outputs[:, :, :, slice(0, 0)] outputs = K.concatenate([left_slice, outputs, right_slice], axis=3) else: # Pad the vertical if self.padding[0][0] > 0: top_slice = K.stack([inputs[:, 0]] * self.padding[0][0], axis=1) else: top_slice = inputs[:, slice(0, 0)] if self.padding[0][1] > 0: bottom_slice = K.stack([inputs[:, -1]] * self.padding[0][1], axis=1) else: bottom_slice = inputs[:, slice(0, 0)] outputs = K.concatenate([top_slice, inputs, bottom_slice], axis=1) # Pad the horizontal if self.padding[1][0] > 0: left_slice = K.stack([outputs[:, :, 0]] * self.padding[1][0], axis=2) else: left_slice = outputs[:, :, slice(0, 0)] if self.padding[1][1] > 0: right_slice = K.stack([outputs[:, :, -1]] * self.padding[1][1], axis=2) else: right_slice = outputs[:, :, slice(0, 0)] outputs = K.concatenate([left_slice, outputs, right_slice], axis=2) return outputs class FillPadding3D(ZeroPadding3D): """Fill-padding layer for 3D input (e.g. image). This layer can add rows or columns that duplicate the edge values. Adapted from keras.layers.ZeroPadding3D by @jweyn # Arguments padding: int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints. - If int: the same symmetric padding is applied to height and width. - If tuple of 3 ints: interpreted as two different symmetric padding values for height and width: `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`. - If tuple of 3 tuples of 2 ints: interpreted as `((left_dim1_pad, right_dim1_pad), (left_dim2_pad, right_dim2_pad), (left_dim3_pad, right_dim3_pad))` data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". # Input shape 5D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad, depth)` - If `data_format` is `"channels_first"`: `(batch, depth, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad)` # Output shape 5D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, first_padded_axis, second_padded_axis, third_axis_to_pad, depth)` - If `data_format` is `"channels_first"`: `(batch, depth, first_padded_axis, second_padded_axis, third_axis_to_pad)` """ def __init__(self, padding=(1, 1, 1), data_format=None, **kwargs): super(FillPadding3D, self).__init__(padding=padding, data_format=data_format, **kwargs) def call(self, inputs): if self.data_format == 'channels_first': # Pad the depth if self.padding[0][0] > 0: low_slice = K.stack([inputs[:, :, 0]] * self.padding[0][0], axis=2) else: low_slice = inputs[:, :, slice(0, 0)] if self.padding[0][1] > 0: high_slice = K.stack([inputs[:, :, -1]] * self.padding[0][1], axis=2) else: high_slice = inputs[:, :, slice(0, 0)] outputs = K.concatenate([low_slice, inputs, high_slice], axis=2) # Pad the vertical if self.padding[1][0] > 0: top_slice = K.stack([outputs[:, :, :, 0]] * self.padding[1][0], axis=3) else: top_slice = outputs[:, :, :, slice(0, 0)] if self.padding[1][1] > 0: bottom_slice = K.stack([outputs[:, :, :, -1]] * self.padding[1][1], axis=3) else: bottom_slice = outputs[:, :, :, slice(0, 0)] outputs = K.concatenate([top_slice, outputs, bottom_slice], axis=3) # Pad the horizontal if self.padding[2][0] > 0: left_slice = K.stack([outputs[:, :, :, :, 0]] * self.padding[2][0], axis=4) else: left_slice = outputs[:, :, :, :, slice(0, 0)] if self.padding[2][1] > 0: right_slice = K.stack([outputs[:, :, :, :, -1]] * self.padding[2][1], axis=4) else: right_slice = outputs[:, :, :, :, slice(0, 0)] outputs = K.concatenate([left_slice, outputs, right_slice], axis=4) else: # Pad the depth if self.padding[0][0] > 0: low_slice = K.stack([inputs[:, 0]] * self.padding[0][0], axis=1) else: low_slice = inputs[:, slice(0, 0)] if self.padding[0][1] > 0: high_slice = K.stack([inputs[:, -1]] * self.padding[0][1], axis=1) else: high_slice = inputs[:, slice(0, 0)] outputs = K.concatenate([low_slice, inputs, high_slice], axis=1) # Pad the vertical if self.padding[1][0] > 0: top_slice = K.stack([outputs[:, :, 0]] * self.padding[1][0], axis=2) else: top_slice = outputs[:, :, slice(0, 0)] if self.padding[1][1] > 0: bottom_slice = K.stack([outputs[:, :, -1]] * self.padding[1][1], axis=2) else: bottom_slice = outputs[:, :, slice(0, 0)] outputs = K.concatenate([top_slice, outputs, bottom_slice], axis=2) # Pad the horizontal if self.padding[2][0] > 0: left_slice = K.stack([outputs[:, :, :, 0]] * self.padding[2][0], axis=3) else: left_slice = outputs[:, :, :, slice(0, 0)] if self.padding[2][1] > 0: right_slice = K.stack([outputs[:, :, :, -1]] * self.padding[2][1], axis=3) else: right_slice = outputs[:, :, :, slice(0, 0)] outputs = K.concatenate([left_slice, outputs, right_slice], axis=3) return outputs class TFPadding2D(ZeroPadding2D): """Padding layer for 2D input (e.g. image) using TensorFlow's padding function. Adapted from keras.layers.ZeroPadding2D by @jweyn # Arguments padding: int, or tuple of 2 ints, or tuple of 2 tuples of 2 ints. - If int: the same symmetric padding is applied to height and width. - If tuple of 2 ints: interpreted as two different symmetric padding values for height and width: `(symmetric_height_pad, symmetric_width_pad)`. - If tuple of 2 tuples of 2 ints: interpreted as `((top_pad, bottom_pad), (left_pad, right_pad))` data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, height, width, channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, height, width)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". mode: A string, one of `"CONSTANT"`, `"SYMMETRIC"`, or `"REFLECT"`. constant_values: A float. The value to pad if mode=='CONSTANT'. # Input shape 5D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad, depth)` - If `data_format` is `"channels_first"`: `(batch, depth, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad)` # Output shape 5D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, first_padded_axis, second_padded_axis, third_axis_to_pad, depth)` - If `data_format` is `"channels_first"`: `(batch, depth, first_padded_axis, second_padded_axis, third_axis_to_pad)` """ def __init__(self, padding=(1, 1), data_format=None, mode='CONSTANT', constant_values=0., **kwargs): super(TFPadding2D, self).__init__(padding=padding, data_format=data_format, **kwargs) self.mode = mode self.constant_values = constant_values def call(self, inputs): if self.data_format == 'channels_first': padding = ((0, 0), (0, 0)) + self.padding else: padding = ((0, 0),) + self.padding + ((0, 0),) return tf.pad(inputs, padding, mode=self.mode, constant_values=self.constant_values) def get_config(self): config = {'padding': self.padding, 'data_format': self.data_format, 'mode': self.mode, 'constant_values': self.constant_values} base_config = super(TFPadding2D, self).get_config() return dict(list(base_config.items()) + list(config.items())) class TFPadding3D(ZeroPadding3D): """Padding layer for 3D input (e.g. image) using TensorFlow's padding function. Adapted from keras.layers.ZeroPadding3D by @jweyn # Arguments padding: int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints. - If int: the same symmetric padding is applied to height and width. - If tuple of 3 ints: interpreted as two different symmetric padding values for height and width: `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`. - If tuple of 3 tuples of 2 ints: interpreted as `((left_dim1_pad, right_dim1_pad), (left_dim2_pad, right_dim2_pad), (left_dim3_pad, right_dim3_pad))` data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". mode: A string, one of `"CONSTANT"`, `"SYMMETRIC"`, or `"REFLECT"`. constant_values: A float. The value to pad if mode=='CONSTANT'. # Input shape 4D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, rows, cols, channels)` - If `data_format` is `"channels_first"`: `(batch, channels, rows, cols)` # Output shape 4D tensor with shape: - If `data_format` is `"channels_last"`: `(batch, padded_rows, padded_cols, channels)` - If `data_format` is `"channels_first"`: `(batch, channels, padded_rows, padded_cols)` """ def __init__(self, padding=(1, 1, 1), data_format=None, mode='CONSTANT', constant_values=0., **kwargs): super(TFPadding3D, self).__init__(padding=padding, data_format=data_format, **kwargs) self.mode = mode self.constant_values = constant_values def call(self, inputs): if self.data_format == 'channels_first': padding = ((0, 0), (0, 0)) + self.padding else: padding = ((0, 0),) + self.padding + ((0, 0),) return tf.pad(inputs, padding, mode=self.mode, constant_values=self.constant_values) def get_config(self): config = {'padding': self.padding, 'data_format': self.data_format, 'mode': self.mode, 'constant_values': self.constant_values} base_config = super(TFPadding3D, self).get_config() return dict(list(base_config.items()) + list(config.items())) # ==================================================================================================================== # # Keras classes for cube sphere convolution # ==================================================================================================================== # class CubeSphereConv2D(Layer): """ 2D convolutional layer for data that is assumed on a cubed sphere. The requirements for using this layer are as follows: - The input data is 5-dimensional: (batch, channels, 6, height, width) for channels_first (batch, 6, height, width, channels) for channels_last - The first spatial dimension must have a length of 6 for the 6 faces of the cubed sphere - The last two faces (indices 4 and 5) are the polar faces This layer learns two separate convolutional kernels and biases, one for the equatorial faces of the cube, and one for the polar faces. Optionally, it can learn separate kernels and biases for each polar face. Note that this layer should be preceded by CubeSpherePadding2D otherwise there is no connection between faces of the cube. Adapted from keras.layers._Conv by @jweyn # Arguments filters: Integer, the dimensionality of the output space (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of n integers, specifying the dimensions of the convolution window. strides: An integer or tuple/list of n integers, specifying the strides of the convolution. Specifying any stride value != 1 is incompatible with specifying any `dilation_rate` value != 1. padding: One of `"valid"` or `"same"` (case-insensitive). data_format: A string, one of `"channels_last"` or `"channels_first"`. The ordering of the dimensions in the inputs. `"channels_last"` corresponds to inputs with shape `(batch, ..., channels)` while `"channels_first"` corresponds to inputs with shape `(batch, channels, ...)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". dilation_rate: An integer or tuple/list of n integers, specifying the dilation rate to use for dilated convolution. Currently, specifying any `dilation_rate` value != 1 is incompatible with specifying any `strides` value != 1. activation: Activation function to use (see [activations](../activations.md)). If you don't specify anything, no activation is applied (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. flip_north_pole: Boolean, whether to reverse the direction of the north pole should that be necessary to match the rotation direction of the south pole in the data independent_north_pole: Boolean, if true, learn separate filters for the north and south poles (3 total sets) kernel_initializer: Initializer for the `kernel` weights matrix (see [initializers](../initializers.md)). bias_initializer: Initializer for the bias vector (see [initializers](../initializers.md)). kernel_regularizer: Regularizer function applied to the `kernel` weights matrix (see [regularizer](../regularizers.md)). bias_regularizer: Regularizer function applied to the bias vector (see [regularizer](../regularizers.md)). activity_regularizer: Regularizer function applied to the output of the layer (its "activation"). (see [regularizer](../regularizers.md)). kernel_constraint: Constraint function applied to the kernel matrix (see [constraints](../constraints.md)). bias_constraint: Constraint function applied to the bias vector (see [constraints](../constraints.md)). """ def __init__(self, filters, kernel_size, strides=1, padding='valid', data_format='channels_first', dilation_rate=1, activation=None, use_bias=True, flip_north_pole=True, independent_north_pole=False, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, **kwargs): super(CubeSphereConv2D, self).__init__(**kwargs) self.filters = filters self.kernel_size = conv_utils.normalize_tuple(kernel_size, 2, 'kernel_size') self.strides = conv_utils.normalize_tuple(strides, 2, 'strides') self.padding = conv_utils.normalize_padding(padding) self.data_format = conv_utils.normalize_data_format(data_format) self.dilation_rate = conv_utils.normalize_tuple(dilation_rate, 2, 'dilation_rate') self.activation = activations.get(activation) self.use_bias = use_bias self.flip_north_pole = flip_north_pole self.independent_north_pole = independent_north_pole self.kernel_initializer = initializers.get(kernel_initializer) self.bias_initializer = initializers.get(bias_initializer) self.kernel_regularizer = regularizers.get(kernel_regularizer) self.bias_regularizer = regularizers.get(bias_regularizer) self.activity_regularizer = regularizers.get(activity_regularizer) self.kernel_constraint = constraints.get(kernel_constraint) self.bias_constraint = constraints.get(bias_constraint) self.rank = 3 self.input_spec = InputSpec(ndim=self.rank + 2) self.equatorial_kernel = None self.equatorial_bias = None self.polar_kernel = None self.polar_bias = None self.north_pole_kernel = None self.north_pole_bias = None def build(self, input_shape): if self.data_format == 'channels_first': channel_axis = 1 else: channel_axis = -1 if input_shape[channel_axis] is None: raise ValueError('The channel dimension of the inputs ' 'should be defined. Found `None`.') input_dim = input_shape[channel_axis] kernel_shape = self.kernel_size + (input_dim, self.filters) self.equatorial_kernel = self.add_weight(shape=kernel_shape, initializer=self.kernel_initializer, name='equatorial_kernel', regularizer=self.kernel_regularizer, constraint=self.kernel_constraint) self.polar_kernel = self.add_weight(shape=kernel_shape, initializer=self.kernel_initializer, name='polar_kernel', regularizer=self.kernel_regularizer, constraint=self.kernel_constraint) if self.independent_north_pole: self.north_pole_kernel = self.add_weight(shape=kernel_shape, initializer=self.kernel_initializer, name='north_pole_kernel', regularizer=self.kernel_regularizer, constraint=self.kernel_constraint) if self.use_bias: self.equatorial_bias = self.add_weight(shape=(self.filters,), initializer=self.bias_initializer, name='equatorial_bias', regularizer=self.bias_regularizer, constraint=self.bias_constraint) self.polar_bias = self.add_weight(shape=(self.filters,), initializer=self.bias_initializer, name='polar_bias', regularizer=self.bias_regularizer, constraint=self.bias_constraint) if self.independent_north_pole: self.north_pole_bias = self.add_weight(shape=(self.filters,), initializer=self.bias_initializer, name='north_pole_bias', regularizer=self.bias_regularizer, constraint=self.bias_constraint) # Set input spec. self.input_spec = InputSpec(ndim=self.rank + 2, axes={channel_axis: input_dim}) self.built = True def call(self, inputs, **kwargs): outputs = [] channels_first = self.data_format == 'channels_first' # Equatorial faces for f in range(4): outputs.append( K.conv2d( inputs[:, :, f, :, :] if channels_first else inputs[:, f], self.equatorial_kernel, strides=self.strides, padding=self.padding, data_format=self.data_format, dilation_rate=self.dilation_rate ) ) if self.use_bias: outputs[f] = K.bias_add( outputs[f], self.equatorial_bias, data_format=self.data_format ) outputs[f] = K.expand_dims(outputs[f], 2 if channels_first else 1) # South pole face outputs.append( K.conv2d( inputs[:, :, 4, :, :] if channels_first else inputs[:, 4], self.polar_kernel, strides=self.strides, padding=self.padding, data_format=self.data_format, dilation_rate=self.dilation_rate ) ) if self.use_bias: outputs[4] = K.bias_add( outputs[4], self.polar_bias, data_format=self.data_format ) outputs[4] = K.expand_dims(outputs[4], 2 if channels_first else 1) # North pole face if self.flip_north_pole: # Reverse the height dimension, 2 if channels_first otherwise 1 outputs.append( K.conv2d( K.reverse(inputs[:, :, 5, :, :], 2) if channels_first else K.reverse(inputs[:, 5], 1), self.north_pole_kernel if self.independent_north_pole else self.polar_kernel, strides=self.strides, padding=self.padding, data_format=self.data_format, dilation_rate=self.dilation_rate ) ) else: outputs.append( K.conv2d( inputs[:, :, 5, :, :] if channels_first else inputs[:, 5], self.north_pole_kernel if self.independent_north_pole else self.polar_kernel, strides=self.strides, padding=self.padding, data_format=self.data_format, dilation_rate=self.dilation_rate ) ) if self.use_bias: outputs[5] = K.bias_add( outputs[5], self.north_pole_bias if self.independent_north_pole else self.polar_bias, data_format=self.data_format ) if self.flip_north_pole: outputs[5] = K.reverse(outputs[5], 2 if channels_first else 1) outputs[5] = K.expand_dims(outputs[5], 2 if channels_first else 1) outputs = K.concatenate(outputs, axis=2 if channels_first else 1) if self.activation is not None: return self.activation(outputs) return outputs def compute_output_shape(self, input_shape): if self.data_format == 'channels_last': # batch, face, height, width, ... space = input_shape[2:4] new_space = [] for i in range(len(space)): new_dim = conv_utils.conv_output_length( space[i], self.kernel_size[i], padding=self.padding, stride=self.strides[i], dilation=self.dilation_rate[i]) new_space.append(new_dim) return (input_shape[0], 6) + tuple(new_space) + (self.filters,) if self.data_format == 'channels_first': # batch, channels, ..., face, height, width space = input_shape[-2:] new_space = [] for i in range(len(space)): new_dim = conv_utils.conv_output_length( space[i], self.kernel_size[i], padding=self.padding, stride=self.strides[i], dilation=self.dilation_rate[i]) new_space.append(new_dim) return (input_shape[0], self.filters, 6) + tuple(new_space) def get_config(self): config = { 'filters': self.filters, 'kernel_size': self.kernel_size, 'strides': self.strides, 'padding': self.padding, 'data_format': self.data_format, 'dilation_rate': self.dilation_rate, 'activation': activations.serialize(self.activation), 'use_bias': self.use_bias, 'flip_north_pole': self.flip_north_pole, 'independent_north_pole': self.independent_north_pole, 'kernel_initializer': initializers.serialize(self.kernel_initializer), 'bias_initializer': initializers.serialize(self.bias_initializer), 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), 'bias_regularizer': regularizers.serialize(self.bias_regularizer), 'activity_regularizer': regularizers.serialize(self.activity_regularizer), 'kernel_constraint': constraints.serialize(self.kernel_constraint), 'bias_constraint': constraints.serialize(self.bias_constraint) } base_config = super(CubeSphereConv2D, self).get_config() return dict(list(base_config.items()) + list(config.items())) class CubeSpherePadding2D(ZeroPadding3D): """ Padding layer for 2D data on a cubed sphere. The requirements for using this layer are as follows: - The input data is 5-dimensional (batch, channels, height, width, 6) - Must follow "channels_first" order - The last dimension must have a length of 6 for the 6 faces of the cubed sphere - The last two faces (indices 4 and 5) are the polar faces Adapted from keras.layers.ZeroPadding3D by @jweyn # Arguments padding: int data_format: 'channels_first' or 'channels_last' """ def __init__(self, padding=(1, 1), data_format='channels_first', **kwargs): data_format = conv_utils.normalize_data_format(data_format) super(CubeSpherePadding2D, self).__init__(padding=padding, data_format=data_format, **kwargs) self.padding = ((0, 0),) + self.padding[1:] def call(self, inputs): p = self.padding[1][0] if self.data_format == 'channels_first': tr = (0, 1, 3, 2) # Pad the equatorial upper/lower boundaries and the polar upper/lower boundaries out = list() # Face 0 out.append(K.expand_dims( K.concatenate([ inputs[:, :, 4, -p:, :], inputs[:, :, 0], inputs[:, :, 5, :p, :] ], axis=2), 2 )) # Face 1 out.append(K.expand_dims( K.concatenate([ tf.transpose(inputs[:, :, 4, ::-1, -p:], tr), inputs[:, :, 1], tf.transpose(K.reverse(inputs[:, :, 5, :, -p:], 3), tr) ], axis=2), 2 )) # Face 2 out.append(K.expand_dims( K.concatenate([ K.reverse(inputs[:, :, 4, :p, ::-1], 2), inputs[:, :, 2], K.reverse(inputs[:, :, 5, -p:, ::-1], 2) ], axis=2), 2 )) # Face 3 out.append(K.expand_dims( K.concatenate([ tf.transpose(K.reverse(inputs[:, :, 4, :, :p], 3), tr), inputs[:, :, 3], tf.transpose(inputs[:, :, 5, ::-1, :p], tr) ], axis=2), 2 )) # Face 4 (south pole) out.append(K.expand_dims( K.concatenate([ K.reverse(inputs[:, :, 2, :p, ::-1], 2), inputs[:, :, 4], inputs[:, :, 0, :p, :] ], axis=2), 2 )) # Face 5 (north pole) out.append(K.expand_dims( K.concatenate([ inputs[:, :, 0, -p:, :], inputs[:, :, 5], K.reverse(inputs[:, :, 2, -p:, ::-1], 2) ], axis=2), 2 )) out1 = K.concatenate(out, axis=2) del out # Pad the equatorial periodic lateral boundaries and the polar left/right boundaries out = list() # Face 0 out.append(K.expand_dims( K.concatenate([ out1[:, :, 3, :, -p:], out1[:, :, 0], out1[:, :, 1, :, :p] ], axis=3), 2 )) # Face 1 out.append(K.expand_dims( K.concatenate([ out1[:, :, 0, :, -p:], out1[:, :, 1], out1[:, :, 2, :, :p] ], axis=3), 2 )) # Face 2 out.append(K.expand_dims( K.concatenate([ out1[:, :, 1, :, -p:], out1[:, :, 2], out1[:, :, 3, :, :p] ], axis=3), 2 )) # Face 3 out.append(K.expand_dims( K.concatenate([ out1[:, :, 2, :, -p:], out1[:, :, 3], out1[:, :, 0, :, :p] ], axis=3), 2 )) # Face 4 out.append(K.expand_dims( K.concatenate([ tf.transpose(K.reverse(out[3][:, :, 0, p:2 * p, :], 2), tr), out1[:, :, 4], tf.transpose(out[1][:, :, 0, p:2 * p, ::-1], tr) ], axis=3), 2 )) # Face 5 out.append(K.expand_dims( K.concatenate([ tf.transpose(out[3][:, :, 0, -2 * p:-p, ::-1], tr), out1[:, :, 5], tf.transpose(K.reverse(out[1][:, :, 0, -2 * p:-p, :], 2), tr) ], axis=3), 2 )) del out1 outputs = K.concatenate(out, axis=2) del out return outputs else: # channels_last tr = (0, 2, 1, 3) # Pad the equatorial upper/lower boundaries and the polar upper/lower boundaries out = list() # Face 0 out.append(K.expand_dims( K.concatenate([ inputs[:, 4, -p:, :], inputs[:, 0], inputs[:, 5, :p, :] ], axis=1), 1 )) # Face 1 out.append(K.expand_dims( K.concatenate([ tf.transpose(inputs[:, 4, ::-1, -p:], tr), inputs[:, 1], tf.transpose(K.reverse(inputs[:, 5, :, -p:], 2), tr) ], axis=1), 1 )) # Face 2 out.append(K.expand_dims( K.concatenate([ K.reverse(inputs[:, 4, :p, ::-1], 1), inputs[:, 2], K.reverse(inputs[:, 5, -p:, ::-1], 1) ], axis=1), 1 )) # Face 3 out.append(K.expand_dims( K.concatenate([ tf.transpose(K.reverse(inputs[:, 4, :, :p], 2), tr), inputs[:, 3], tf.transpose(inputs[:, 5, ::-1, :p], tr) ], axis=1), 1 )) # Face 4 (south pole) out.append(K.expand_dims( K.concatenate([ K.reverse(inputs[:, 2, :p, ::-1], 1), inputs[:, 4], inputs[:, 0, :p, :] ], axis=1), 1 )) # Face 5 (north pole) out.append(K.expand_dims( K.concatenate([ inputs[:, 0, -p:, :], inputs[:, 5], K.reverse(inputs[:, 2, -p:, ::-1], 1) ], axis=1), 1 )) out1 = K.concatenate(out, axis=1) del out # Pad the equatorial periodic lateral boundaries and the polar left/right boundaries out = list() # Face 0 out.append(K.expand_dims( K.concatenate([ out1[:, 3, :, -p:], out1[:, 0], out1[:, 1, :, :p] ], axis=2), 1 )) # Face 1 out.append(K.expand_dims( K.concatenate([ out1[:, 0, :, -p:], out1[:, 1], out1[:, 2, :, :p] ], axis=2), 1 )) # Face 2 out.append(K.expand_dims( K.concatenate([ out1[:, 1, :, -p:], out1[:, 2], out1[:, 3, :, :p] ], axis=2), 1 )) # Face 3 out.append(K.expand_dims( K.concatenate([ out1[:, 2, :, -p:], out1[:, 3], out1[:, 0, :, :p] ], axis=2), 1 )) # Face 4 out.append(K.expand_dims( K.concatenate([ tf.transpose(K.reverse(out[3][:, 0, p:2 * p, :], 1), tr), out1[:, 4], tf.transpose(out[1][:, 0, p:2 * p, ::-1], tr) ], axis=2), 1 )) # Face 5 out.append(K.expand_dims( K.concatenate([ tf.transpose(out[3][:, 0, -2 * p:-p, ::-1], tr), out1[:, 5], tf.transpose(K.reverse(out[1][:, 0, -2 * p:-p, :], 1), tr) ], axis=2), 1 )) del out1 outputs = K.concatenate(out, axis=1) del out return outputs # ==================================================================================================================== # # Keras other layer classes and methods # ==================================================================================================================== # def slice_layer(start, end, step=None, axis=1): """ Return a Lambda layer that performs slicing on a tensor. :param start: int: start index :param end: int: end index :param step: int: stepping parameter :param axis: int: axis along which to slice """ if axis < 1: raise ValueError("'slice_layer' can only work on a specified axis > 0") def slice_func(x): slices = [slice(None)] * axis slices.append(slice(start, end, step)) return x[tuple(slices)] return Lambda(slice_func) class RowConnected2D(LocallyConnected2D): """Row-connected layer for 2D inputs. The `RowConnected2D` layer works similarly to the `Conv2D` layer, except that weights are shared only along rows, that is, a different set of filters is applied at each different row of the input. Adapted from keras.layers.local.LocallyConnected2D by @jweyn # Examples ```python # apply a 3x3 unshared weights convolution with 64 output filters # on a 32x32 image with `data_format="channels_last"`: model = Sequential() model.add(LocallyConnected2D(64, (3, 3), input_shape=(32, 32, 3))) # now model.output_shape == (None, 30, 30, 64) # notice that this layer will consume (30*30)*(3*3*3*64) # + (30*30)*64 parameters # add a 3x3 unshared weights convolution on top, with 32 output filters: model.add(LocallyConnected2D(32, (3, 3))) # now model.output_shape == (None, 28, 28, 32) ``` # Arguments filters: Integer, the dimensionality of the output space (i.e. the number of output filters in the convolution). kernel_size: An integer or tuple/list of 2 integers, specifying the width and height of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. strides: An integer or tuple/list of 2 integers, specifying the strides of the convolution along the width and height. Can be a single integer to specify the same value for all spatial dimensions. padding: Currently only support `"valid"` (case-insensitive). `"same"` will be supported in future. data_format: A string, one of `channels_last` (default) or `channels_first`. The ordering of the dimensions in the inputs. `channels_last` corresponds to inputs with shape `(batch, height, width, channels)` while `channels_first` corresponds to inputs with shape `(batch, channels, height, width)`. It defaults to the `image_data_format` value found in your Keras config file at `~/.keras/keras.json`. If you never set it, then it will be "channels_last". activation: Activation function to use (see [activations](../activations.md)). If you don't specify anything, no activation is applied (ie. "linear" activation: `a(x) = x`). use_bias: Boolean, whether the layer uses a bias vector. kernel_initializer: Initializer for the `kernel` weights matrix (see [initializers](../initializers.md)). bias_initializer: Initializer for the bias vector (see [initializers](../initializers.md)). kernel_regularizer: Regularizer function applied to the `kernel` weights matrix (see [regularizer](../regularizers.md)). bias_regularizer: Regularizer function applied to the bias vector (see [regularizer](../regularizers.md)). activity_regularizer: Regularizer function applied to the output of the layer (its "activation"). (see [regularizer](../regularizers.md)). kernel_constraint: Constraint function applied to the kernel matrix (see [constraints](../constraints.md)). bias_constraint: Constraint function applied to the bias vector (see [constraints](../constraints.md)). # Input shape 4D tensor with shape: `(samples, channels, rows, cols)` if data_format='channels_first' or 4D tensor with shape: `(samples, rows, cols, channels)` if data_format='channels_last'. # Output shape 4D tensor with shape: `(samples, filters, new_rows, new_cols)` if data_format='channels_first' or 4D tensor with shape: `(samples, new_rows, new_cols, filters)` if data_format='channels_last'. `rows` and `cols` values might have changed due to padding. """ def __init__(self, *args, **kwargs): super(RowConnected2D, self).__init__(*args, **kwargs) def build(self, input_shape): if self.data_format == 'channels_last': input_row, input_col = input_shape[1:-1] input_filter = input_shape[3] else: input_row, input_col = input_shape[2:] input_filter = input_shape[1] if input_row is None or input_col is None: raise ValueError('The spatial dimensions of the inputs to ' ' a LocallyConnected2D layer ' 'should be fully-defined, but layer received ' 'the inputs shape ' + str(input_shape)) output_row = conv_utils.conv_output_length(input_row, self.kernel_size[0], self.padding, self.strides[0]) output_col = conv_utils.conv_output_length(input_col, self.kernel_size[1], self.padding, self.strides[1]) self.output_row = output_row self.output_col = output_col self.kernel_shape = ( output_row, self.kernel_size[0], self.kernel_size[1], input_filter, self.filters) self.kernel = self.add_weight(shape=self.kernel_shape, initializer=self.kernel_initializer, name='kernel', regularizer=self.kernel_regularizer, constraint=self.kernel_constraint) if self.use_bias: self.bias = self.add_weight(shape=(output_row, 1, self.filters), initializer=self.bias_initializer, name='bias', regularizer=self.bias_regularizer, constraint=self.bias_constraint) else: self.bias = None if self.data_format == 'channels_first': self.input_spec = InputSpec(ndim=4, axes={1: input_filter}) else: self.input_spec = InputSpec(ndim=4, axes={-1: input_filter}) self.built = True def call(self, inputs): output = row_conv2d(inputs, self.kernel, self.kernel_size, self.strides, (self.output_row, self.output_col), self.data_format) if self.use_bias: output = K.bias_add(output, self.bias, data_format=self.data_format) output = self.activation(output) return output def row_conv2d(inputs, kernel, kernel_size, strides, output_shape, data_format=None): """Apply 2D conv with weights shared only along rows. Adapted from K.local_conv2d by @jweyn # Arguments inputs: 4D tensor with shape: (batch_size, filters, new_rows, new_cols) if data_format='channels_first' or 4D tensor with shape: (batch_size, new_rows, new_cols, filters) if data_format='channels_last'. kernel: the row-shared weights for convolution, with shape (output_rows, kernel_size, input_channels, filters) kernel_size: a tuple of 2 integers, specifying the width and height of the 2D convolution window. strides: a tuple of 2 integers, specifying the strides of the convolution along the width and height. output_shape: a tuple with (output_row, output_col) data_format: the data format, channels_first or channels_last # Returns A 4d tensor with shape: (batch_size, filters, new_rows, new_cols) if data_format='channels_first' or 4D tensor with shape: (batch_size, new_rows, new_cols, filters) if data_format='channels_last'. # Raises ValueError: if `data_format` is neither `channels_last` or `channels_first`. """ data_format = K.normalize_data_format(data_format) stride_row, stride_col = strides output_row, output_col = output_shape out = [] for i in range(output_row): # Slice the rows with the neighbors they need slice_row = slice(i * stride_row, i * stride_col + kernel_size[0]) if data_format == 'channels_first': x = inputs[:, :, slice_row, :] # batch, 16, 5, 144 else: x = inputs[:, slice_row, :, :] # batch, 5, 144, 16 # Convolve, resulting in an array with only one row: batch, 1, 140, 6 or batch, 6, 1, 140 x = K.conv2d(x, kernel[i], strides=strides, padding='valid', data_format=data_format) out.append(x) if data_format == 'channels_first': output = K.concatenate(out, axis=2) else: output = K.concatenate(out, axis=1) del x del out return output # ==================================================================================================================== # # Keras loss functions # ==================================================================================================================== # def latitude_weighted_loss(loss_function=mean_squared_error, lats=None, output_shape=(), axis=-2, weighting='cosine'): """ Create a loss function that weights inputs by a function of latitude before calculating the loss. :param loss_function: method: Keras loss function to apply after the weighting :param lats: ndarray: 1-dimensional array of latitude coordinates :param output_shape: tuple: shape of expected model output :param axis: int: latitude axis in model output shape :param weighting: str: type of weighting to apply. Options are: cosine: weight by the cosine of the latitude (default) midlatitude: weight by the cosine of the latitude but also apply a 25% reduction to the equator and boost to the mid-latitudes :return: callable loss function """ if weighting not in ['cosine', 'midlatitude']: raise ValueError("'weighting' must be one of 'cosine' or 'midlatitude'") if lats is not None: lat_tensor = K.zeros(lats.shape) lat_tensor.assign(K.cast_to_floatx(lats[:])) weights = K.cos(lat_tensor * np.pi / 180.) if weighting == 'midlatitude': weights = weights + 0.5 * K.pow(K.sin(lat_tensor * 2 * np.pi / 180.), 2.) weight_shape = output_shape[axis:] for d in weight_shape[1:]: weights = K.expand_dims(weights, axis=-1) weights = K.repeat_elements(weights, d, axis=-1) else: weights = K.ones(output_shape) def lat_loss(y_true, y_pred): return loss_function(y_true * weights, y_pred * weights) return lat_loss def anomaly_correlation(y_true, y_pred, mean=0., regularize_mean='mse', reverse=True): """ Calculate the anomaly correlation. FOR NOW, ASSUMES THAT THE CLIMATOLOGICAL MEAN IS 0, AND THEREFORE REQUIRES DATA TO BE SCALED TO REMOVE SPATIALLY-DEPENDENT MEAN. :param y_true: Tensor: target values :param y_pred: Tensor: model-predicted values :param mean: float: subtract this global mean from all predicted and target array values. IGNORED FOR NOW. :param regularize_mean: str or None: if not None, also penalizes a form of mean squared error: global: penalize differences in the global mean spatial: penalize differences in spatially-averaged mean (last two dimensions) mse: penalize the mean squared error mae: penalize the mean absolute error :param reverse: bool: if True, inverts the loss so that -1 is the target score :return: float: anomaly correlation loss """ if regularize_mean is not None: assert regularize_mean in ['global', 'spatial', 'mse', 'mae'] a = (K.mean(y_pred * y_true) / K.sqrt(K.mean(K.square(y_pred)) * K.mean(K.square(y_true)))) if regularize_mean is not None: if regularize_mean == 'global': m = K.abs((K.mean(y_true) - K.mean(y_pred)) / K.mean(y_true)) elif regularize_mean == 'spatial': m = K.mean(K.abs((K.mean(y_true, axis=[-2, -1]) - K.mean(y_pred, axis=[-2, -1])) / K.mean(y_true, axis=[-2, -1]))) elif regularize_mean == 'mse': m = mean_squared_error(y_true, y_pred) elif regularize_mean == 'mae': m = mean_absolute_error(y_true, y_pred) if reverse: if regularize_mean is not None: return m - a else: return -a else: if regularize_mean: return a - m else: return a def anomaly_correlation_loss(mean=None, regularize_mean='mse', reverse=True): """ Create a Keras loss function for anomaly correlation. :param mean: ndarray or None: if not None, must be an array with the same shape as the expected prediction, except that the first (batch) axis should have a dimension of 1. :param regularize_mean: str or None: if not None, also penalizes a form of mean squared error: global: penalize differences in the global mean spatial: penalize differences in spatially-averaged mean (last two dimensions) mse: penalize the mean squared error mae: penalize the mean absolute error :param reverse: bool: if True, inverts the loss so that -1 is the (minimized) target score. Must be True if regularize_mean is not None. :return: method: anomaly correlation loss function """ if mean is not None: assert len(mean.shape) > 1 assert mean.shape[0] == 1 mean_tensor = K.variable(mean, name='anomaly_correlation_mean') if regularize_mean is not None: assert regularize_mean in ['global', 'spatial', 'mse', 'mae'] reverse = True def acc_loss(y_true, y_pred): if mean is not None: a = (K.mean((y_pred - mean_tensor) * (y_true - mean_tensor)) / K.sqrt(K.mean(K.square((y_pred - mean_tensor))) * K.mean(K.square((y_true - mean_tensor))))) else: a = (K.mean(y_pred * y_true) / K.sqrt(K.mean(K.square(y_pred)) * K.mean(K.square(y_true)))) if regularize_mean is not None: if regularize_mean == 'global': m = K.abs((K.mean(y_true) - K.mean(y_pred)) / K.mean(y_true)) elif regularize_mean == 'spatial': m = K.mean(K.abs((K.mean(y_true, axis=[-2, -1]) - K.mean(y_pred, axis=[-2, -1])) / K.mean(y_true, axis=[-2, -1]))) elif regularize_mean == 'mse': m = mean_squared_error(y_true, y_pred) elif regularize_mean == 'mae': m = mean_absolute_error(y_true, y_pred) if reverse: if regularize_mean is not None: return m - a else: return -a else: if regularize_mean: return a - m else: return a return acc_loss # Compatibility names # lat_loss = latitude_weighted_loss() # acc_loss = anomaly_correlation_loss() class AsymptoteReLU(Layer): """ Rectified Linear Unit activation function. Returns: `f(x) = x / (beta * x + 1)` for `x >= 0`, `f(x) = alpha * x` otherwise. Input shape: Arbitrary. Use the keyword argument `input_shape` (tuple of integers, does not include the samples axis) when using this layer as the first layer in a model. Output shape: Same shape as the input. Arguments: alpha: Float >= 0. Negative slope coefficient. beta: Float >= 0. Positive asymptote coefficient. """ def __init__(self, alpha=0., beta=1., **kwargs): super(AsymptoteReLU, self).__init__(**kwargs) if alpha < 0.: raise ValueError('alpha of AsymptoteReLU layer ' 'cannot be negative value: ' + str(alpha)) if beta <= 0.: raise ValueError('beta of AsymptoteReLU layer ' 'cannot be zero or negative value: ' + str(beta)) self.support_masking = True self.alpha = K.cast_to_floatx(alpha) self.beta = K.cast_to_floatx(beta) self.one = K.cast_to_floatx(1.) def call(self, inputs, **kwargs): return K.relu(inputs / (K.abs(self.beta * inputs) + self.one), alpha=self.alpha) def get_config(self): config = { 'alpha': self.alpha, 'beta': self.beta } base_config = super(AsymptoteReLU, self).get_config() return dict(list(base_config.items()) + list(config.items())) def compute_output_shape(self, input_shape): return input_shape # ==================================================================================================================== # # PyTorch classes # ==================================================================================================================== # class TorchReshape(object): def __init__(self, shape): if not isinstance(shape, tuple): raise ValueError("'shape' must be a tuple of integers") self.shape = shape def __call__(self, x): return x.view(*self.shape)