MNISTSeq.py

__author__ = "Yinchong Yang"
__copyright__ = "Siemens AG, 2018"
__licencse__ = "MIT"
__version__ = "0.1"

"""
MIT License
Copyright (c) 2018 Siemens AG
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""



"""
We first sample MNIST digits to form sequences of random lengths. 
The sequence is labeled as one if it contains a zero, and is labeled zero otherwise.
This simulates a high dimensional sequence classification task, such as predicting therapy decision
and survival of patients based on their historical clinical event information.    
We train plain LSTM and Tensor-Train LSTM for this task. 
After the training, we apply Layer-wise Relevance Propagation to identify the digit(s) that 
have influenced the classification. 
Apparently, we would expect the LRP algorithm would assign high relevance value to the zero(s)
in the sequence. 
These experiments turn out to be successful, which demonstrates that 
i) the LSTM and TT-LSTM can indeed learn the mapping from a zero to the sequence class, and that 
ii) both LSTMs have no problem in storing the zero pattern over a period of time, because the 
classifier is deployed only at the last hidden state, and that   
iii) the implementation of the LRP algorithm, complex as it is, is also correct, in that 
the zeros are assigned high relevance scores. 

Especially the experiments with the plain LSTM serve as simulation study supporting our submission of 
“Yinchong Yang, Volker Tresp, Marius Wunderle, Peter A. Fasching, 
Explaining Therapy Predictions with Layer-wise Relevance Propagation in Neural Networks, at IEEE ICHI 2018”. 

The original LRP for LSTM from the repository: 
                        https://github.com/ArrasL/LRP_for_LSTM
which we modified and adjusted for keras models.  

Feel free to experiment with the hyper parameters and suggest other sequence classification tasks.
Have fun ;)  
"""


import pickle
import sys
import numpy as np
from numpy import newaxis as na
import keras
from keras.layers.recurrent import Recurrent
from keras import backend as K
from keras.engine import InputSpec
from keras import activations
from keras import initializers
from keras import regularizers
from keras import constraints
from keras.engine.topology import Layer

from TTLayer import *
from TTRNN import TT_LSTM


def make_seq(n, x, y, maxlen=32, seed=123):
    np.random.seed(seed)
    lens = np.random.choice(range(2, maxlen), n)
    seqs = np.zeros((n, maxlen, 28**2))
    labels = np.zeros(n)
    digits_label = np.zeros((n, maxlen), dtype='int32')-1
    ids = np.zeros((n, maxlen), dtype='int64')-1
    for i in range(n):
        digits_inds = np.random.choice(range(x.shape[0]), lens[i])
        ids[i, -lens[i]::] = digits_inds
        seqs[i, -lens[i]::, :] = x[digits_inds]
        digits_label[i, -lens[i]::] = y[digits_inds]
        class_inds = y[digits_inds]

        if True:
            # option 1: is there any 0 in the sequence?
            labels[i] = (0 in class_inds)
        else:
            # option 2: even number of 0 -> label=0, odd number of 0 -> label=1
            labels[i] = len(np.where(class_inds == 0)[0]) % 2 == 1
    return [seqs, labels, digits_label, ids]


# From: https://github.com/ArrasL/LRP_for_LSTM
def lrp_linear(hin, w, b, hout, Rout, bias_nb_units, eps, bias_factor, debug=False):
    """
    LRP for a linear layer with input dim D and output dim M.
    Args:
    - hin:            forward pass input, of shape (D,)
    - w:              connection weights, of shape (D, M)
    - b:              biases, of shape (M,)
    - hout:           forward pass output, of shape (M,) (unequal to np.dot(w.T,hin)+b if more than one incoming layer!)
    - Rout:           relevance at layer output, of shape (M,)
    - bias_nb_units:  number of lower-layer units onto which the bias/stabilizer contribution is redistributed
    - eps:            stabilizer (small positive number)
    - bias_factor:    for global relevance conservation set to 1.0, otherwise 0.0 to ignore bias redistribution
    Returns:
    - Rin:            relevance at layer input, of shape (D,)
    """
    sign_out = np.where(hout[na, :] >= 0, 1., -1.)  # shape (1, M)

    numer = (w * hin[:, na]) + \
            ((bias_factor * b[na, :] * 1. + eps * sign_out * 1.) * 1. / bias_nb_units)  # shape (D, M)

    denom = hout[na, :] + (eps * sign_out * 1.)  # shape (1, M)

    message = (numer / denom) * Rout[na, :]  # shape (D, M)

    Rin = message.sum(axis=1)  # shape (D,)

    # Note: local  layer   relevance conservation if bias_factor==1.0 and bias_nb_units==D
    #       global network relevance conservation if bias_factor==1.0 (can be used for sanity check)
    if debug:
        print("local diff: ", Rout.sum() - Rin.sum())

    return Rin


def sigmoid(x):
    x = x.astype('float128')
    return 1. / (1. + np.exp(-x))

# Modified from https://github.com/ArrasL/LRP_for_LSTM 
def lstm_lrp(l, d, train_data = True):
    if train_data:
        x_l = X_tr[l]
        y_l = Y_tr[l]
        z_l = Z_tr[l]
        # d_l = d_tr[l]
    else:
        x_l = X_te[l]
        y_l = Y_te[l]
        z_l = Z_te[l]
        # d_l = d_te[l]


    # calculate the FF pass in LSTM for every time step
    pre_gates = np.zeros((MAXLEN, d*4))
    gates = np.zeros((MAXLEN, d * 4))
    h = np.zeros((MAXLEN, d))
    c = np.zeros((MAXLEN, d))

    for t in range(MAXLEN):
        z = np.dot(x_l[t], Ws)
        if t > 0:
            z += np.dot(h[t-1], Us)
        z += b
        pre_gates[t] = z
        z0 = z[0:d]
        z1 = z[d:2*d]
        z2 = z[2*d:3*d]
        z3 = z[3 * d::]
        i = sigmoid(z0)
        f = sigmoid(z1)
        c[t] = f * c[t-1] + i * np.tanh(z2)
        o = sigmoid(z3)
        h[t] = o * np.tanh(c[t])
        gates[t] = np.concatenate([i, f, np.tanh(z2), o])

    # check: z_l[12] / h[-1][12]

    Rh = np.zeros((MAXLEN, d))
    Rc = np.zeros((MAXLEN, d))
    Rg = np.zeros((MAXLEN, d))
    Rx = np.zeros((MAXLEN, 28**2))

    bias_factor = 0

    Rh[MAXLEN-1] = lrp_linear(hin=z_l,
                              w=Dense_w,
                              b=np.array(Dense_b),
                              hout=np.dot(z_l, Dense_w)+Dense_b,
                              Rout=np.array([y_l]),
                              bias_nb_units=len(z_l),
                              eps=eps,
                              bias_factor=bias_factor)


    for t in reversed(range(MAXLEN)):
        # t = MAXLEN-1
        # print t

        Rc[t] += Rh[t]
        # Rc[t] = Rh[t]
        if t > 0:
            Rc[t-1] = lrp_linear(gates[t, d: 2 * d] * c[t - 1],  # gates[t , 2 *d: 3 *d ] *c[ t -1],
                                  np.identity(d),
                                  np.zeros((d)),
                                  c[t],
                                  Rc[t],
                                  2*d,
                                  eps,
                                  bias_factor,
                                  debug=False)

        Rg[t] = lrp_linear(gates[t, 0:d] * gates[t, 2*d:3*d],  # h_input: i + g
                           np.identity(d),                     # W
                           np.zeros((d)),                      # b
                           c[t],                               # h_output
                           Rc[t],                              # R_output
                           2 * d,
                           eps,
                           bias_factor,
                           debug=False)

        # foo = np.dot(x_l[t], Ws[:,2*d:3*d]) + np.dot(h[t-1], Us[:, 2*d:3*d]) + b[2*d:3*d]

        Rx[t] = lrp_linear(x_l[t],
                           Ws[:,2*d:3*d],
                           b[2*d:3*d],
                           pre_gates[t, 2*d:3*d],
                           Rg[t],
                           d + 28 ** 2,
                           eps,
                           bias_factor,
                           debug=False)

        if t > 0:
            Rh[t-1] = lrp_linear(h[t-1],
                                 Us[:,2*d:3*d],
                                 b[2*d:3*d],
                                 pre_gates[t, 2 * d:3 * d],
                                 Rg[t],
                                 d + 28**2,
                                 eps,
                                 bias_factor,
                                 debug=False)

    # hin, w, b, hout, Rout, bias_nb_units, eps, bias_factor, debug=False

    # Rx[np.where(d_l==-1.)[0]] *= 0
    return Rx


from keras.datasets import mnist
from keras.utils import to_categorical
from keras.models import Model, Input
from keras.layers import Dense, GRU, LSTM, Dropout, Masking
from keras.optimizers import *
from keras.regularizers import l2

from sklearn.metrics import *


# Script configurations ###################################################################

seed=111111
use_TT = True  # whether use Tensor-Train or plain RNNs


# Prepare the data ########################################################################
# Load the MNIST data and build sequences:
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(x_train.shape[0], -1)
x_test = x_test.reshape(x_test.shape[0], -1)

MAXLEN = 32  # max length of the sequences

X_tr, Y_tr, d_tr, idx_tr = make_seq(n=10000, x=x_train, y=y_train, maxlen=MAXLEN, seed=seed)
X_te, Y_te, d_te, idx_te = make_seq(n=1000, x=x_test, y=y_test, maxlen=MAXLEN, seed=seed+1)

# Define the model ######################################################################

if use_TT:
    # TT settings
    tt_input_shape = [7, 7, 16]
    tt_output_shape = [4, 4, 4]
    tt_ranks = [1, 4, 4, 1]

rnn_size = 64

X = Input(shape=X_tr.shape[1::])
X_mask = Masking(mask_value=0.0, input_shape=X_tr.shape[1::])(X)

if use_TT:
    Z = TT_LSTM(tt_input_shape=tt_input_shape, tt_output_shape=tt_output_shape, tt_ranks=tt_ranks,
                return_sequences=False, recurrent_dropout=.5)(X_mask)
    Out = Dense(units=1, activation='sigmoid', kernel_regularizer=l2(1e-2))(Z)
else:
    Z = LSTM(units=rnn_size, return_sequences=False, recurrent_dropout=.5)(X_mask) # dropout=.5,
    Out = Dense(units=1, activation='sigmoid', kernel_regularizer=l2(1e-2))(Z)

rnn_model = Model(X, Out)
rnn_model.compile(optimizer=Adam(1e-3), loss='binary_crossentropy',
                  metrics=['accuracy'])

# Train the model and save the results ######################################################
rnn_model.fit(X_tr, Y_tr, epochs=50, batch_size=32, validation_split=.2, verbose=2)


Y_hat = rnn_model.predict(X_tr, verbose=2).reshape(-1)
train_acc = (np.round(Y_hat) == Y_tr).mean()
Y_pred = rnn_model.predict(X_te, verbose=2).reshape(-1)
(np.round(Y_pred) == Y_te).mean()
pred_acc = (np.round(Y_pred) == Y_te).mean()


# Collect all hidden layers ################################################################
if use_TT:
    # Reconstruct the fully connected input-to-hidden weights:
    from keras.initializers import constant
    _tt_output_shape = np.copy(tt_output_shape)
    _tt_output_shape[0] *= 4

    fc_w = rnn_model.get_weights()[0]
    fc_layer = TT_Layer(tt_input_shape=tt_input_shape, tt_output_shape=_tt_output_shape, tt_ranks=tt_ranks,
                            kernel_initializer=constant(value=fc_w), use_bias=False)
    fc_input = Input(shape=(X_tr.shape[2],))
    fc_output = fc_layer(fc_input)
    fc_model = Model(fc_input, fc_output)
    fc_model.compile('sgd', 'mse')

    fc_recon_mat = fc_model.predict(np.identity(X_tr.shape[2]))

    # Reconstruct the entire LSTM:
    fc_Z = LSTM(units=np.prod(tt_output_shape), return_sequences=False, dropout=.5, recurrent_dropout=.5,
                weights=[fc_recon_mat, rnn_model.get_weights()[2], rnn_model.get_weights()[1]])(X_mask)

else:
    fc_Z = LSTM(units=rnn_size, return_sequences=False, dropout=.5, recurrent_dropout=.5,
                weights=rnn_model.get_weights()[0:3])(X_mask)

fc_Out = Dense(units=1, activation='sigmoid', kernel_regularizer=l2(1e-3),
               weights=rnn_model.get_weights()[3::])(fc_Z)
fc_rnn_model = Model(X, fc_Out)
fc_rnn_model.compile(optimizer=Adam(1e-3), loss='binary_crossentropy',
                  metrics=['accuracy'])

fc_rnn_model.evaluate(X_te, Y_te, verbose=2)



# Calculate the LRP: #########################################################################
fc_Z_model = Model(X, fc_Z)
fc_Z_model.compile('sgd', 'mse')

Y_hat_fc = fc_rnn_model.predict(X_tr)
Y_pred_fc = fc_rnn_model.predict(X_te)

Ws = fc_rnn_model.get_weights()[0]
Us = fc_rnn_model.get_weights()[1]
b = fc_rnn_model.get_weights()[2]
Dense_w = fc_rnn_model.get_weights()[3]
Dense_b = fc_rnn_model.get_weights()[4]

Z_tr = fc_Z_model.predict(X_tr)
Z_te = fc_Z_model.predict(X_te)

eps = 1e-4

is_number_flag = np.where(d_te != -1)

# All relevance scores of the test sequences
lrp_te = np.vstack([lstm_lrp(i, rnn_size, False).sum(1) for i in range(X_te.shape[0])])

lrp_auroc = roc_auc_score((d_te == 0).astype('int')[is_number_flag].reshape(-1),
                           lrp_te[is_number_flag].reshape(-1))
lrp_auprc = average_precision_score((d_te == 0).astype('int')[is_number_flag].reshape(-1),
                           lrp_te[is_number_flag].reshape(-1))


# The reported results:
print pred_acc
print lrp_auroc
print lrp_auprc