deep_mil.py

#!/usr/bin/env python

import argparse
import os

import keras
import numpy as np
import tensorflow as tf
from keras import backend as K
from keras.layers import Dense, Conv2D, MaxPooling2D
from keras.layers import Layer
from keras.models import Sequential, load_model

from datasets import Mnist
from tf_cnnvis import deconv_visualization


class NoisyAnd(Layer):
    """Custom NoisyAND layer from the Deep MIL paper"""

    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super(NoisyAnd, self).__init__(**kwargs)

    def build(self, input_shape):
        self.a = 10  # fixed, controls the slope of the activation
        self.b = self.add_weight(name='b',
                                 shape=(1, input_shape[3]),
                                 initializer='uniform',
                                 trainable=True)
        super(NoisyAnd, self).build(input_shape)  # Be sure to call this at the end

    def call(self, x):
        mean = tf.reduce_mean(x, axis=[1, 2])
        res = (tf.nn.sigmoid(self.a * (mean - self.b)) - tf.nn.sigmoid(-self.a * self.b)) / (
                tf.nn.sigmoid(self.a * (1 - self.b)) - tf.nn.sigmoid(-self.a * self.b))
        return res

    def compute_output_shape(self, input_shape):
        return input_shape[0], input_shape[3]


def define_model(input_shape, num_classes):
    """Define Deep FCN for MIL, layer-by-layer from original paper"""
    model = Sequential()
    model.add(Conv2D(32, kernel_size=(3, 3),
                     activation='relu',
                     input_shape=input_shape))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(128, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(128, (3, 3), activation='relu'))
    model.add(Conv2D(1000, (3, 3), activation='relu'))
    model.add(Conv2D(num_classes, (1, 1), activation='relu'))
    model.add(NoisyAnd(num_classes))
    model.add(Dense(num_classes, activation='softmax'))
    return model


def train(epochs, seed, batch_size, dataset):
    """Train FCN"""
    np.random.seed(seed)

    model = define_model(dataset.input_shape, dataset.num_classes)
    model.compile(loss=keras.losses.categorical_crossentropy,
                  optimizer=keras.optimizers.Adadelta(),
                  metrics=['accuracy'])
    model.fit(dataset.x_train, dataset.y_train,
              batch_size=batch_size,
              epochs=epochs,
              verbose=1,
              validation_data=(dataset.x_test, dataset.y_test))
    return model


def visualize(sess, model, dataset):
    """Save png of deconvolution image from first image in test set"""
    deconv_visualization(sess, {model.input: dataset.x_test[0:1, :, :, :]})


def main():
    # Parse Arguments
    parser = argparse.ArgumentParser(description='Deep MIL Arguments')
    parser.add_argument('-e', '--NUM_EPOCHS', default=1, type=int)  # Number of epochs for which to train the model
    parser.add_argument('-r', '--SEED', default=123)  # specify the seed
    parser.add_argument('-b', '--BATCH_SIZE', default=128, type=int)  # batch size for training
    parser.add_argument('-s', '--SAVE_DIRECTORY', default='conv_mil/', type=str)  # Where to save model
    parser.add_argument('-m', '--MODEL_NAME', default='model.h5', type=str)  # To save individual run
    parser.add_argument('-t', '--TRAIN', default=1, type=int)  # whether to train (1) or load model (0)
    parser.add_argument('-v', '--VISUALIZE', default=1, type=int)  # whether to visualize the output
    flags = vars(parser.parse_args())

    # Build MNIST dataset
    dataset = Mnist()

    # Make save directory if it doesn't exist
    if not os.path.exists(flags['SAVE_DIRECTORY']):
        os.makedirs(flags['SAVE_DIRECTORY'])

    filepath = os.path.join(flags['SAVE_DIRECTORY'], flags['MODEL_NAME'])
    with tf.Graph().as_default():
        with tf.Session() as sess:
            K.set_session(sess)

            # Train or load model
            if flags['TRAIN']:
                model = train(epochs=flags['NUM_EPOCHS'],
                              seed=flags['SEED'],
                              batch_size=flags['BATCH_SIZE'],
                              dataset=dataset)
                model.save(filepath)
            else:
                model = load_model(filepath)

            # Visualize with tf_cnnvis
            visualize(sess, model, dataset)


if __name__ == "__main__":
    main()