Source code for nnsa.keras.layers

"""
This module contains helper functions to create a block of layers of a neural network.
"""

import numpy as np

from tensorflow import keras
from tensorflow.python.keras import Input
from tensorflow.python.keras.models import Model
from tensorflow.python.keras.layers import Conv1D, BatchNormalization, Activation, Dropout, \
    Layer, Conv2D, Conv2DTranspose, Dense, MaxPool1D, Concatenate, MaxPool2D, UpSampling2D, MaxPooling2D, Conv1DTranspose
from tensorflow.keras.regularizers import L2
import tensorflow as tf


[docs]class Sampling(Layer): """ Uses (z_mean, z_log_var) to sample z. Adapted from https://keras.io/examples/generative/vae/ Examples: >>> from artefact_detection.keras_utils.losses import kl_loss >>> x = Input(shape=[100, 1]) >>> z_mean = Dense(10, name="z_mean")(x) >>> z_log_var = Dense(10, name="z_log_var")(x) >>> z = Sampling(name='z')([z_mean, z_log_var]) >>> y = Dense(100)(z) >>> model = Model(x, y) >>> model.add_loss(kl_loss(z_mean, z_log_var, beta=1)) """ def __init__(self, *args, **kwargs): self.is_placeholder = True super(Sampling, self).__init__(*args, **kwargs)
[docs] def call(self, inputs): z_mean, z_log_var = inputs batch = tf.shape(z_mean)[0] dim = tf.shape(z_mean)[1] epsilon = tf.keras.backend.random_normal(shape=(batch, dim)) return z_mean + tf.exp(0.5 * z_log_var) * epsilon
[docs]def layer_block(layer, *args, batchnorm=True, batchnorm_kwargs=None, activation='relu', activation_kwargs=None, dropout=False, dropout_kwargs=None, **kwargs): """ Layer with optional batch normalization, activation and dropout regularization. Args: layer (keras.Layer): keras layer object. E.g., layer=Conv2D *args: optional positional arguments for the keras layer. batchnorm (bool): bool specifying whether to add a BatchNormalization layer. batchnorm_kwargs (dict): optional keyword arguments for BatchNormalization(). activation (str or None or Activation): specifies the activation used. See Activation(). activation_kwargs (dict): optional keyword arguments for Activation(). dropout (bool): bool specifying whether to add a Dropout layer. dropout_kwargs (dict): optional keyword arguments for Dropout(). **kwargs: optional keyword arguments for `layer`. Returns: block (function): function that takes in a keras layer and returns the output layer after passing the input thtough the specified `layer` and BatchNormalization, Activation and Dropout layers. """ def block(x): # Layer. x = layer(*args, **kwargs)(x) # Batch normalization layer. if batchnorm: x = BatchNormalization(**batchnorm_kwargs if batchnorm_kwargs is not None else dict())(x) # Activation layer. if activation is not None: x = Activation(activation, **activation_kwargs if activation_kwargs is not None else dict())(x) # Dropout layer. if dropout: x = Dropout(**dropout_kwargs if dropout_kwargs is not None else dict())(x) return x return block
[docs]def conv1d_block(filters, kernel_size, strides=1, padding='same', conv1d_kwargs=None, batchnorm=True, batchnorm_kwargs=None, activation='relu', activation_kwargs=None, dropout=True, dropout_kwargs=None): """ Helper function that implements layer_block with a Conv1D layer. """ return layer_block(Conv1D, filters, kernel_size, strides=strides, padding=padding, batchnorm=batchnorm, batchnorm_kwargs=batchnorm_kwargs, activation=activation, activation_kwargs=activation_kwargs, dropout=dropout, dropout_kwargs=dropout_kwargs, **(conv1d_kwargs if conv1d_kwargs is not None else dict()))
[docs]def conv2d_block(filters, kernel_size, strides=1, padding='same', conv2d_kwargs=None, batchnorm=True, batchnorm_kwargs=None, activation='relu', activation_kwargs=None, dropout=True, dropout_kwargs=None): """ Helper function that implements layer_block with a Conv2D layer. """ return layer_block(Conv2D, filters, kernel_size, strides=strides, padding=padding, batchnorm=batchnorm, batchnorm_kwargs=batchnorm_kwargs, activation=activation, activation_kwargs=activation_kwargs, dropout=dropout, dropout_kwargs=dropout_kwargs, **(conv2d_kwargs if conv2d_kwargs is not None else dict()))
[docs]def conv1dtranspose_block(filters, kernel_size, strides, padding='same', conv1dtranspose_kwargs=None, batchnorm=True, batchnorm_kwargs=None, activation='relu', activation_kwargs=None, dropout=True, dropout_kwargs=None): """ Helper function that implements layer_block with a Conv1DTranspose layer. """ return layer_block(Conv1DTranspose, filters, kernel_size, strides=strides, padding=padding, batchnorm=batchnorm, batchnorm_kwargs=batchnorm_kwargs, activation=activation, activation_kwargs=activation_kwargs, dropout=dropout, dropout_kwargs=dropout_kwargs, **(conv1dtranspose_kwargs if conv1dtranspose_kwargs is not None else dict()))
[docs]def conv2dtranspose_block(filters, kernel_size, strides=(1, 1), padding='same', conv2dtranspose_kwargs=None, batchnorm=True, batchnorm_kwargs=None, activation='relu', activation_kwargs=None, dropout=True, dropout_kwargs=None): """ Helper function that implements layer_block with a Conv2DTranspose layer. """ return layer_block(Conv2DTranspose, filters, kernel_size, batchnorm=batchnorm, batchnorm_kwargs=batchnorm_kwargs, activation=activation, activation_kwargs=activation_kwargs, dropout=dropout, dropout_kwargs=dropout_kwargs, strides=strides, padding=padding, **(conv2dtranspose_kwargs if conv2dtranspose_kwargs is not None else dict()))
[docs]def dense_block(units, dense_kwargs=None, batchnorm=True, batchnorm_kwargs=None, activation='relu', activation_kwargs=None, dropout=True, dropout_kwargs=None): """ Helper function that implements layer_block with a Dense layer. """ return layer_block(Dense, units, batchnorm=batchnorm, batchnorm_kwargs=batchnorm_kwargs, activation=activation, activation_kwargs=activation_kwargs, dropout=dropout, dropout_kwargs=dropout_kwargs, **(dense_kwargs if dense_kwargs is not None else dict()))
[docs]def inception1d(filters, kernel_sizes=(10, 20, 40), pool_size=None, strides=1, use_bottleneck=True, bottleneck_size=None, activation=None): """ Block of layers for inception module for 1D inputs. """ if bottleneck_size is None: bottleneck_size = filters activation = keras.activations.get(activation) kernel_sizes = np.round(np.asarray(kernel_sizes)) kernel_sizes = [int(k) for k in kernel_sizes] # Must be type int. if pool_size is None: # Take first kernel_size. pool_size = kernel_sizes[0] def _conv1d(filters, kernel_size, **kwargs): return Conv1D(filters=filters, kernel_size=kernel_size, padding='same', strides=strides, activation=None, use_bias=False, **kwargs) def module(x): if use_bottleneck and int(x.shape[-1]) > 1: bottleneck = _conv1d(filters=bottleneck_size, kernel_size=1)(x) else: bottleneck = x # Layers to concatenate. concat_list = [] # Loop over kernel sizes. for kernel_size in kernel_sizes: conv1d_layer = _conv1d( filters=filters, kernel_size=kernel_size)(bottleneck) concat_list.append(conv1d_layer) # Add pooling layer. maxpool = MaxPool1D(pool_size=pool_size, strides=strides, padding='same')(x) conv_pool = _conv1d(filters=filters, kernel_size=1)(maxpool) concat_list.append(conv_pool) # Concatenate. x = Concatenate(axis=-1)(concat_list) x = activation(x) return x return module
[docs]def inception2d(filters, kernel_sizes=((10, 1), (20, 1), (40, 1)), pool_size=None, strides=(1, 1), use_bottleneck=True, bottleneck_size=None, activation=None): """ Block of layers for inception module for 2D inputs. """ if bottleneck_size is None: bottleneck_size = filters activation = keras.activations.get(activation) kernel_sizes = np.round(np.asarray(kernel_sizes)) kernel_sizes = [tuple([int(i) for i in k]) for k in kernel_sizes] # Must be type int. if pool_size is None: # Take first kernel size. pool_size = kernel_sizes[0] def _conv2d(filters, kernel_size, **kwargs): return Conv2D(filters=filters, kernel_size=kernel_size, padding='same', strides=strides, activation=None, use_bias=False, **kwargs) def module(x): if use_bottleneck and int(x.shape[-1]) > 1: bottleneck = _conv2d(filters=bottleneck_size, kernel_size=(1, 1))(x) else: bottleneck = x # Layers to concatenate. concat_list = [] # Loop over kernel sizes. for kernel_size in kernel_sizes: conv1d_layer = _conv2d( filters=filters, kernel_size=kernel_size)(bottleneck) concat_list.append(conv1d_layer) # Add pooling layer. maxpool = MaxPool2D(pool_size=pool_size, strides=strides, padding='same')(x) conv_pool = _conv2d(filters=filters, kernel_size=(1, 1))(maxpool) concat_list.append(conv_pool) # Concatenate. x = Concatenate(axis=-1)(concat_list) x = activation(x) return x return module
[docs]def inception2dtranspose(filters, kernel_sizes=((10, 1), (20, 1), (40, 1)), strides=(2, 2), activation=None): """ Block of layers for transpose inception module for 2D inputs. """ activation = keras.activations.get(activation) kernel_sizes = np.round(np.asarray(kernel_sizes)) kernel_sizes = [tuple([int(i) for i in k]) for k in kernel_sizes] # Must be type int. def _conv2dtranpose(filters, kernel_size, **kwargs): return Conv2DTranspose(filters=filters, kernel_size=kernel_size, padding='same', strides=strides, activation=None, use_bias=False, **kwargs) def module(x): # Layers to concatenate. concat_list = [] # Loop over kernel sizes. for kernel_size in kernel_sizes: conv1dtranspose_layer = _conv2dtranpose( filters=filters, kernel_size=kernel_size)(x) concat_list.append(conv1dtranspose_layer) # Add upsampling layer (transpose of the pooling layer). up = Conv2D(filters=filters, kernel_size=(1, 1))(x) up = UpSampling2D(size=strides)(up) concat_list.append(up) # Concatenate. x = Concatenate(axis=-1)(concat_list) x = activation(x) return x return module
[docs]def sinc2d(filters=64, branches=5, l2_regularization=0.001, activation='elu', dropout = True, dropout_kwargs = None): """ Block of layers for shared inception (sinc) module for 2D inputs. """ # Default options. dropout_kwargs = dict({'rate': 0.2}, **dropout_kwargs if dropout_kwargs is not None else dict()) def block(x): # Branch 1 of sinc block. l1 = [] ldt = Conv2D(filters = filters, kernel_size = (1,1), padding = 'same', activation = activation, kernel_regularizer = L2(l2_regularization))(x) for i in range(int(branches)): ldt = Conv2D(filters = filters, kernel_size = (3,1), padding = 'same', activation = activation, kernel_regularizer = L2(l2_regularization))(ldt) l1.append(ldt) # Branch 2 of sinc block. l2 = MaxPooling2D(pool_size = (3,1), strides = (1,1), padding = 'same')(x) l2 = Conv2D(filters = filters, kernel_size = (1,1), padding = 'same', activation = activation, kernel_regularizer = L2(l2_regularization))(l2) # Branch 3 of sinc block. l3 = Conv2D(filters = filters, kernel_size = (1,1), padding = 'same', activation = activation, kernel_regularizer = L2(l2_regularization))(x) # Concatenate. x = Concatenate(axis=-1)([l3, l2] + l1) # Dropout. if dropout: x = Dropout(**dropout_kwargs)(x) return x return block