{"cells":[{"metadata":{},"cell_type":"markdown","source":"In this example, we reconstruct one part of MNIST digits given another part. This is a demonstration of the approach introduced in https://arxiv.org/abs/1904.10387"},{"metadata":{"trusted":true},"cell_type":"code","source":"num_feat = 200   # we are looking for that many features\nmax_epoch = 150  \n\nfast = False    # if false, we use the Moore-Penrose pseudo inverse\nfudge = 0.0     # otherwise, we add this much times the identity before inverting\n\nbatch_size = 5000  # should be at least 10 times greater than num_feat\n\nload = False","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# get rid of all the warnings\nimport os, warnings\nwarnings.simplefilter(\"ignore\", category=FutureWarning)\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"]=\"2\"\nimport tensorflow as tf\ntf.logging.set_verbosity(tf.logging.ERROR)\n\nimport keras;\nfrom keras.preprocessing.image import ImageDataGenerator;\nimport numpy as np\nimport tensorflow as tf\nimport numpy.linalg as nl\nimport tensorflow.linalg as tl\nfrom tensorflow.keras.models import Sequential, Model, load_model\nfrom tensorflow.keras.layers import Conv2D, Dense, Flatten, Input, concatenate, Lambda, ZeroPadding2D, BatchNormalization\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping\n\nimport pickle\n\nimport matplotlib.pyplot as plt\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Our data consists of two sub-areas of MNIST digits. For convenience, we actually use the whole images, but our two neural nets will be \"looking\" at different parts of it."},{"metadata":{"trusted":true},"cell_type":"code","source":"mnist = tf.keras.datasets.mnist;\n(train_images, train_labels), (test_images, test_labels) = mnist.load_data();\ntrain_images = train_images.reshape(-1,28,28,1).astype(np.float32)/255.0\ntest_images = test_images.reshape(-1,28,28,1).astype(np.float32)/255.0\n\n# keras really wants us to have targets, so let's make some dummy ones\ndummy = np.zeros(len(train_images))\ntest_dummy = np.zeros(len(test_images))\n    \n# we agument the data to make it a bit harder\ndatagen = ImageDataGenerator(\n    rotation_range=5,\n    width_shift_range=3.0,\n    height_shift_range=1.5)\n\nimport copy \ntest_datagen = copy.deepcopy(datagen)\ndatagen.fit(train_images)\ntest_datagen.fit(test_images)\n\n# defines which part of the image will be looked at by each net\nx_mask = np.zeros((28,28,1),dtype=np.bool)\ny_mask = np.ones((28,28,1),dtype=np.bool)\nx_mask[:,15:,0] = True\ny_mask[:,14:,0] = False","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's see what this looks like:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_images(pairs):\n    n = len(pairs)\n    width = 15\n    rows = (n-1) // width + 1\n    cols = min(n, width)\n    bigim = np.zeros((28*rows, 28*cols, 3), np.float32)\n    fig = plt.figure(figsize=(2*cols, 2*rows))\n    for (j,(x,y)) in enumerate(pairs):\n        im = np.zeros((28,28,3))\n        im[np.squeeze(x_mask),0:1] = x[np.squeeze(x_mask)]+0.2\n        im[np.squeeze(x_mask),2:3] = x[np.squeeze(x_mask)]*0.5+0.1\n        im[np.squeeze(y_mask),:] = y[np.squeeze(y_mask)]\n        a = (j % width)*28\n        b = (j // width)*28\n        bigim[b:b+28,a:a+28,:] = im\n        \n    plt.imshow(bigim)\n    plt.axis('off')\n    \nnum = 30\nexamples = datagen.flow(train_images[:num], dummy[:num], batch_size=num)[0][0]\n    \nshow_images([(examples[i], examples[i]) for i in range(num)])\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We build two convolutional neural nets, each acting on a different part of the image (defined by x_mask and y_mask). "},{"metadata":{"trusted":true},"cell_type":"code","source":"def build_model():\n    def cnn(image, mask):\n        x = Lambda(lambda y: y * mask.astype(np.float32))(image)\n        x = Conv2D(32, kernel_size=(3, 3), strides=(1,1), activation='relu')(x)\n        x = Conv2D(64, kernel_size=(4, 4), strides=(2,2), activation='relu')(x)\n        x = Conv2D(64, kernel_size=(3, 3), strides=(1,1), activation='relu')(x)\n        x = Conv2D(128, kernel_size=(4, 4), strides=(2,2), activation='relu')(x)\n        x = Flatten()(x)\n        x = Dense(1024, activation='relu')(x)\n        features = Dense(num_feat, activation='tanh')(x)\n        return features\n\n    image = Input(shape=(28,28,1))\n    out1 = cnn(image, x_mask)\n    out2 = cnn(image, y_mask)\n\n    # We will use each models separately for predictions...\n    feat1 = Model(inputs=image, outputs=out1)\n    feat2 = Model(inputs=image, outputs=out2)\n\n    # ...but together for training as they have a joint loss function\n    model = Model(inputs=image, outputs=concatenate([out1, out2], axis=1))\n    \n    return feat1, feat2, model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here are the functions implementing the RFA objective as well as the inference step:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def tf_prep(T):\n    z = tf.math.sqrt(tf.cast(tf.shape(T)[0], tf.float64))\n    return tf.cast(T, tf.float64)/z\n    \ndef prep(T):\n    z = np.sqrt(T.shape[0])\n    return T.astype(np.float64)/z\n\nif fast:\n    def pinv(T):\n        D = fudge * tf.eye(tf.shape(T)[0], dtype = tf.float64)\n        return tl.inv(T + D)\nelse:\n    import tensorflow_probability as tfp\n    pinv = tfp.math.pinv\n\n@tf.custom_gradient\ndef tf_relevance(F, G):  \n    F = tf_prep(F)\n    G = tf_prep(G)\n    D = fudge * tf.eye(tf.shape(F)[1], dtype = tf.float64)\n\n    iK = pinv(tl.transpose(F) @ F)\n    iL = pinv(tl.transpose(G) @ G)\n    \n    A = tl.transpose(F) @ G\n    U = iK @ A\n    V = iL @ tl.transpose(A)\n\n    dF = 2 * (G - F @ U) @ V @ iK\n    dG = 2 * (F - G @ V) @ U @ iL\n    grad = lambda delta: (tf.cast(delta * dF, tf.float32), tf.cast(delta * dG, tf.float32))\n        \n    return tl.trace(U @ V), grad\n\ndef cov(F, G):\n    F = prep(F)\n    G = prep(G)\n    return (np.transpose(F) @ F, \n            np.transpose(G) @ G, \n            np.transpose(F) @ G)\n    \n\n# produces a function that maps a vector of features on X to the inferred (expected) value of Y\n# More generally, Y can be the value of any function f of Y applied to each training example. \n# It's shape is (num training example, a,b,...) where (a,b,...) is the shape of the output of f.\ndef inferY(ker, G, Y):  \n    n, k = F.shape\n    data_shape = Y.shape[1:]\n    \n    K, L, A = ker\n    Y = prep(Y).reshape((n,-1))\n    G = prep(G)\n    \n    M = (np.transpose(Y) @ G) @ nl.solve(L, np.transpose(A)) \n    iK = nl.pinv(K)\n\n    def infer(tF):\n        X = np.transpose(M @ (iK @ np.transpose(tF))) \n        return X.reshape((-1,) + data_shape)\n    \n    return infer\n\n# in principle the above works for inferring the full pixel-pixel covariance matrix, \n# but it's not practical. So we have to do that differently:\ndef inferY2(ker, G, Y):  \n    n, k = F.shape\n    data_shape = Y.shape[1:]\n    \n    K, L, A = ker\n    Y = prep(Y).reshape((n,-1))\n    m = Y.shape[1]\n    G = prep(G)\n    \n    Y2G = np.zeros((m,m,k), np.float32)\n    for i in range(k): \n        print(i+1,\"/\",k,end=\"\\r\")\n        Y2G[:,:,i] = np.transpose(Y) @ (Y * G[:, i:i+1]) * np.sqrt(n)\n        \n    M = Y2G @ nl.solve(L, np.transpose(A)) \n    \n    iK = nl.pinv(K)\n    \n    def infer(tF):\n        X = np.transpose(M @ (iK @ np.transpose(tF))) \n        return X.reshape((-1,) + data_shape + data_shape)\n    \n    return infer\n\n# produces a function which maps a vector of features on Y to the inferred (expected) value of X\ndef inferX(ker, F, X): \n    K, L, A = ker\n    return inferY((L, K, np.transpose(A)), F, X)\n\ndef singvals(ker):\n    K, L, A = ker\n    M = nl.solve(K, A) @ nl.solve(L, np.transpose(A))\n    return sorted(nl.eigvals(M), reverse=True)\n    \ndef RFA_Loss():\n    def loss(dummy, features):\n        F, G = tf.split(features, 2, axis=1)\n        return -tf_relevance(F, G) \n    return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"It's time to learn the correlations"},{"metadata":{"trusted":true},"cell_type":"code","source":"feat1, feat2, model = build_model()\n\nmodel.compile(optimizer = Adam(lr=0.001), loss = RFA_Loss())\n\nif load:\n    model.load_weights('../input/rfa-mnist-weights/best_weights.h5')\nelse:\n    callbacks = [\n        ModelCheckpoint('best_weights.h5', monitor='val_loss', mode='min', verbose=0, save_best_only=True, save_weights_only=True),\n        EarlyStopping(monitor='val_loss', mode='min', verbose=0, patience=2)\n    ]\n    \n    train_data = datagen.flow(train_images, dummy, batch_size=batch_size, shuffle=True)\n    test_data = test_datagen.flow(test_images, test_dummy, batch_size=batch_size)\n    \n    model.fit_generator(train_data, steps_per_epoch = len(train_images) / batch_size,\n                        epochs = max_epoch, verbose=2,\n                        callbacks = callbacks,\n                        validation_data = test_data, validation_freq = 5)\n    \n    model.save_weights('last_weights.h5')\n    model.load_weights('best_weights.h5')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We compute the best features on the full training set"},{"metadata":{"trusted":true},"cell_type":"code","source":"#pred_data = datagen.flow(train_images, batch_size=5000)  # slowly fills memory and crashes\n#F = feat1.predict_generator(pred_data) #, steps=120000/batch_size, verbose=0)\n#G = feat2.predict_generator(pred_data) #, steps=120000/batch_size, verbose=0)\nF = feat1.predict(train_images, batch_size=2000, verbose=0)\nG = feat2.predict(train_images, batch_size=2000, verbose=0)\nker = cov(F, G)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"A little consistency check: we compute the singular values using different random batches. The result shouldn't depend too much on the sample size. This screens for a failure mode in which the objective is independent of the sample for a given sample size, but highly dependent on the sample size."},{"metadata":{"trusted":true},"cell_type":"code","source":"if not load:\n    for n in [5000, 10000, 20000]:\n        sample = np.random.choice(F.shape[0], n)\n        plt.plot(np.abs(singvals(cov(F[sample], G[sample]))), label=str(n))\n    \n    plt.legend();","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's graph the singular values evaluated on test data"},{"metadata":{"trusted":true},"cell_type":"code","source":"tF = feat1.predict(test_images * x_mask, batch_size=2000, verbose=0)\ntG = feat2.predict(test_images * y_mask, batch_size=2000, verbose=0)\n\ntker = cov(tF, tG)\nplt.plot(np.abs(singvals(tker)));\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Finally, we can try inferring one part of test images given their other part (in red). These are meant to be averages over all possibilities, not samples from the conditional distributions."},{"metadata":{"trusted":true},"cell_type":"code","source":"# make function that predicts conditional mean of pixel value\ninfer = inferY(ker, G, train_images * y_mask)\n\n# try it on a sample of test images\nnp.random.seed(1337)\nsamples = np.random.choice(test_images.shape[0], size=15*10, replace=False)\n\ny_pred = infer(tF[samples])\nshow_images(list(zip(test_images[samples], y_pred))) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Okay, but maybe a traditional feed-forward encoder-decoder might have done just as well. However, we have a full probabilist model. For instance, let us check the uncertainty on each pixel values:"},{"metadata":{"trusted":true},"cell_type":"code","source":"# predict the mean of the square of the pixel values\ninfer2_local = inferY(ker, G, (train_images * y_mask)**2)\n\ny2_pred = infer2_local(tF[samples])\n\n# that allows us to compute the standard deviations\ndev = y2_pred - y_pred**2\nshow_images(list(zip(test_images[samples], dev*3)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Even better, we can predict all pixel-pixel correlations!"},{"metadata":{"trusted":true},"cell_type":"code","source":"infer2 = inferY2(ker, G, train_images * y_mask)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And uncover the main predictions ambiguities by exploring variations in the direction of largest covariance"},{"metadata":{"trusted":true},"cell_type":"code","source":"s = 28**2\ny2_pred = infer2(tF[samples])\nleft_images = []\nright_images = []\nfor i in range(10*15):\n    y = y_pred[i]\n    y2 = y2_pred[i]\n    a, ef = nl.eigh(y2.reshape((s,s)) - y.reshape((s,1)) * y.reshape((1,s)))\n    idx = a.argsort()[::-1]\n    alpha = np.sqrt(a[idx[0]])                     # largest eigenvalue\n    vec = ef[:,idx[0]].reshape((28,28,1))          # corresponding eigenvector \n    vec = vec/nl.norm(vec)\n    left_images.append(y - alpha*vec)\n    left_images.append(y)\n    left_images.append(y + alpha*vec)\n    for j in range(3): right_images.append(test_images[samples[i]])\n\nshow_images(list(zip(right_images, left_images)))","execution_count":null,"outputs":[]}],"metadata":{"language_info":{"name":"python","version":"3.6.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"}},"nbformat":4,"nbformat_minor":1}