blackbox.py

"""
Collection of functions related to analysing Deep Neural Networks (DNNs) from a blackbox perspective.
"""

import numpy as np


def getOrthogonalBasis(V):
    '''
    Compute an orthogonal basis for the space generated by the vectors given
    by the rows of matrix V.

    Parameters
    ----------
    V : array
        2-dimensional array with the row vectors.

    Returns
    -------
    array
        Orthogonal basis given by the rows of the returned matrix.
    '''
    rk = np.linalg.matrix_rank(V)
    q, _ = np.linalg.qr(V.T)
    return q[:,0:rk].T

def getReconstructedModel(Ws, Bs, w, b, inputShape=(32,32,1), endWithReLU=False):
    """
    Reconstruct a tensorflow model with hidden dense ReLU layer weights and biases 'Ws', 'Bs'
    and a linear output neuron with weights and biases 'w', 'b'.
    """

    from tensorflow.keras.models import Sequential
    from tensorflow.keras.layers import Input, Dense, Flatten

    # --- adapt dimension of current neuron weight and bias ---
    # The current neurons weight 'w' and bias 'b' are simple arrays of shape (neurons in previous layer,), e.g. (10,).
    # To be used as weight for the construction of a neural network model, we have to expand the dimension to
    # (neurons in previous layer, neurons in the current layer (=1)):
    if len(w.shape) == 1:
        w = np.expand_dims(w, -1)
        b = np.expand_dims(b, -1)

    # --- create model ---
    myModel = Sequential()

    myModel.add(Input(shape=inputShape))
    myModel.add(Flatten())

    for weights, biases in zip(Ws, Bs):
        nNeurons = weights.shape[-1]
        myModel.add(Dense(nNeurons, activation="relu",))
        myModel.layers[-1].set_weights([weights, biases])

    # add output neurons
    if endWithReLU: activation="relu"
    else: activation="linear"
    nNeurons = w.shape[-1]
    myModel.add(Dense(nNeurons, activation=activation))
    myModel.layers[-1].set_weights([w, b])

    return myModel

def getCIFAR10testImage(): 
    """Pick a random flattened input image from CIFAR10."""
    
    import tensorflow as tf
    from keras.datasets import cifar10

    def normalize_resize(image, label):
        image = tf.cast(image, tf.float32)
        image = tf.divide(image, 255)
        image = tf.image.resize(image, (32,32))
        return image, label

    (trainX, trainy), (testX, testy) = cifar10.load_data()

    #trainX, trainy = normalize_resize(trainX, trainy)
    testX, testy = normalize_resize(testX, testy)
    # Flatten the test dataset
    testX = tf.keras.layers.Flatten()(testX)
    # pick a random input image
    pickID = np.random.choice(np.arange(testX.shape[0]))
    return testX[pickID].numpy()
    

def findCorner(weights, biases, shape, targetNeurons, targetValue=1, tol=1e-6,
               timeout=100, dataset=None):
    """
    Given the weights and biases for a sequence of layers, find an input such
    that targetNeurons in the last layer are close to targetValue, with an error
    vector of norm less than tol.

    Parameters
    ----------
    weights :
        A list of weights for the known layers (each of them a 2D array).
    biases :
        A list of biases for the known layers (each of them a 1D array).
    shape : tuple
        The shape of the input that is returned.
    targetNeurons : list
        The indices of the neurons that need to be fixed.
    targetValue : float
        The value that they should be fixed to.
    tol :
        The allowed norm of the error vector.
    timeout:
        Max number of steps before abandoning the current walk and starting from
        scratch with double timeout
    """
    if dataset is None:
        x0 = np.random.uniform(low=-1, high=1, size=shape).flatten()
    elif dataset=='CIFAR10':
        x0 = getCIFAR10testImage()

    for i in range(timeout):

        # Local matrix at the starting point
        M0, b0 = getLocalMatrixAndBias(weights, biases, x0)
        y0 = np.matmul(x0, M0) + b0

        # Compute the point where the target neurons would be zero assuming linearity never breaks
        dx = np.matmul(targetValue - y0[targetNeurons], np.linalg.pinv(M0[:,targetNeurons]))
        x1 = x0 + dx

        # If linearity breaks we can end up with unexpected values, so we halve the shift until
        # we get strictly smaller outputs than before.
        M1, b1 = getLocalMatrixAndBias(weights, biases, x1)
        y1 = np.matmul(x1, M1) + b1
        while ((targetValue-y0)**2)[targetNeurons].sum() < ((targetValue-y1)**2)[targetNeurons].sum():
            dx /= 2
            x1 = x0 + dx
            M1, b1 = getLocalMatrixAndBias(weights, biases, x1)
            y1 = np.matmul(x1, M1) + b1

        # If the outputs are small enough, we exit the loop
        if ((targetValue-y1)**2)[targetNeurons].sum() < tol**2:
            # Ensure the outputs are all slightly larger than the target
            dy = y1[targetNeurons] - targetValue
            dy[dy > 0] = 0
            dy *= -2
            dx = np.matmul(dy, np.linalg.pinv(M1[:,targetNeurons]))

            return (x1 + dx).reshape(shape)

        # Otherwise we restart from the new x
        x0 = x1

    # If timeout was reached, we restart from a fresh x0 and double the timout
    return findCorner(weights, biases, shape, targetNeurons, targetValue, tol, 2 * timeout)

def getLocalMatrixAndBias(weights, biases, x0):
    """
    Given the weights and biases up to a certain layer, find the equivalent matrix and bias
    around the vicinity of an input x.

    Parameters
    ----------
    weights:
        A list of weights for the known layers (each of them a 2D array).
    biases:
        A list of biases for the known layers (each of them a 1D array).
    x0:
        A 1D array of inputs (of dimension equal to the second dimension of weights[0])
        OR a 2D array which consists of a vector of inputs (along the first dimension)

    Returns
    -------
    M
        A 2D array representing the local matrix
        OR a 3D array representing one matrix per input (along the first dimension)
    b
        A 1D array representing the local bias vector
        OR a 2D array representing one bias per input (along the first dimension)
    """

    # Special case if x0 is not vectorized
    if len(x0.shape) < 2:
        M, b = getLocalMatrixAndBias(weights, biases, np.array([x0]))
        return M[0], b[0]

    MM = []
    bb = []
    for x0i in x0:
        M = weights[0].copy()
        b = biases[0].copy()
        x = np.matmul(x0i, M) + b
        for layer_id in range(1, len(weights)):
            M_hat = weights[layer_id].copy()
            M_hat[x < 0] = 0
            x = np.matmul(x, M_hat) + biases[layer_id]
            b = np.matmul(b, M_hat) + biases[layer_id]
            M = np.matmul(M, M_hat)

        MM.append(M)
        bb.append(b)

    return np.array(MM), np.array(bb)

def getHiddenVector(weights, biases, l, x, relu=False):
    """
    Computes the hidden vector resulting from applying the first l hidden layers
    to the given input x.

    Parameters
    ----------
    weights : array
        List of weights corresponding to the hidden layers. The i-th element in
        the list is a 2-dimensional array with the weights of the incoming
        connections to the neurons in the i-th hidden layer.
    biases : array
        List of biases corresponding to the hidden layers. The i-th element in
        the list is a 1-dimensional array with the biases of the neurons in the
        i-th hidden layer.
    l : int
        Number of hidden layers to consider.
    x : array
        1-dimensional array representing an input to the DNN.
    relu : bool, optional
        Specifies whether to compute the hidden vector before (relu=False) or
        after (relu=True) the ReLU in layer l.

    Returns
    -------
    array
        The hidden vector corresponding to x after applying the first l hidden
        layers.
    """
    y = x
    for i in range(l):
        y = np.matmul(y, weights[i]) + biases[i]
        if (i < (l - 1)) or relu:
            y *= (y > 0)
    return y

def getOrthogonalBasisForInnerLayerSpace(x, weights, biases, l, eps):
    '''
    Given the input x to the DNN, compute an orthogonal basis for the vector
    space generated by the linear neighbourhood of x after the l first layers,
    inclusive.

    Parameters
    ----------
    x : array
        1-dimensional array representing an input to the DNN.
    weights : array
        List of weights corresponding to the hidden layers. The i-th element in
        the list is a 2-dimensional array with the weights of the incoming
        connections to the neurons in the i-th hidden layer.
    biases : array
        List of biases corresponding to the hidden layers. The i-th element in
        the list is a 1-dimensional array with the biases of the neurons in the
        i-th hidden layer.
    l : int
        Number of hidden layers to consider.
    eps : float
        "Infinitesimal" variation used for computing the vector space after the
        first l layers. This value should be small enough so that varying each
        coordinate at the input does not cross the boundary of the linear
        neighbourhood of x.

    Returns
    -------
    B : array
        Orthogonal basis of the vector space after the first l layers. The
        basis is given by the rows of the returned matrix.
    diffs : array
        Let f_{1..l} be the function that evaluates the first l layers of the
        DNN. Also, let y = f_{1..l}(x) and h_i = f_{1..l}(x + eps * e_i), where
        e_i is the i-th canonical vector.
        diffs is a 2-dimensional array representing a matrix where the i-th row
        is the vector h_i - y.
    '''
    n = np.prod(x.shape)
    if l > 0:
        M, b = getLocalMatrixAndBias(weights[0:l], biases[0:l], x)
    else:
        M = np.identity(n, dtype=np.double)
        b = np.zeros_like(x)
    # y = f_{1..l}(x)
    y = np.matmul(x, M) + b
    y *= (y > 0)
    # h_i = f_{1..l}(x + eps * e_i)
    xdx = (eps * np.identity(n, dtype=np.double)) + x
    hidVecs = np.matmul(xdx, M) + b
    hidVecs *= (hidVecs > 0)
    # Find orthogonal basis for the space generated by all h_i - y
    diffs = hidVecs - y
    B = getOrthogonalBasis(diffs)

    return B, diffs

def getLastLayerOutputMatrixBlackbox(func, weights, biases, inputShape, layerId, dataset, eps, tol):
    out = []
    timeout = 3
    for i in range(weights[-1].shape[1]):
        print("Computing output coefficients for neuron "+str(i+1)+"/"+str(weights[-1].shape[1]))
        for k in range(timeout):
            x = findCorner(weights, biases, (1,)+inputShape, [i], targetValue=-eps, dataset=dataset, tol=tol)
            Ml,_ = getLocalMatrixAndBias(weights, biases, x.flatten())
            dx = Ml[:,i].copy().flatten()
            dx = (dx * eps / (dx@Ml)[i]).reshape((1,)+inputShape)
            Mr,_ = getLocalMatrixAndBias(weights, biases, (x+2*dx).flatten())
            if (np.abs(Ml-Mr) > 1e-10).any():
                if k == timeout-1:
                    print("LINEARITY ERROR: try increasing --eps")
                    exit(-1)
                else:
                    continue
            c = func(x+2*dx) - 2 * func(x+dx) + func(x)
            if (c == 0).any():
                print("PRECISION ERROR(1): Try decreasing --eps")
                exit(-1)
            out.append(c.flatten() / eps)
            break
    return np.array(out)