title classification using attentionbased deep multiple instance learning mil mohamad jaberhttps www linkedin cominmohamadjaber1 date created 20210816 last modified 20211125 description mil approach to classify bags of instances and get their individual instance score accelerator gpu introduction what is multiple instance learning mil usually with supervised learning algorithms the learner receives labels for a set of instances in the case of mil the learner receives labels for a set of bags each of which contains a set of instances the bag is labeled positive if it contains at least one positive instance and negative if it does not contain any motivation it is often assumed in image classification tasks that each image clearly represents a class label in medical imaging e g computational pathology etc an entire image is represented by a single class label cancerousnoncancerous or a region of interest could be given however one will be interested in knowing which patterns in the image is actually causing it to belong to that class in this context the images will be divided and the subimages will form the bag of instances therefore the goals are to 1 learn a model to predict a class label for a bag of instances 2 find out which instances within the bag caused a position class label prediction implementation the following steps describe how the model works 1 the feature extractor layers extract feature embeddings 2 the embeddings are fed into the mil attention layer to get the attention scores the layer is designed as permutationinvariant 3 input features and their corresponding attention scores are multiplied together 4 the resulting output is passed to a softmax function for classification references attentionbased deep multiple instance learninghttps arxiv orgabs1802 04712 some of the attention operator code implementation was inspired from https github comutayaoattendeepmil imbalanced data tutorialhttps www tensorflow orgtutorialsstructureddataimbalanceddata by tensorflow setup create dataset we will create a set of bags and assign their labels according to their contents if at least one positive instance is available in a bag the bag is considered as a positive bag if it does not contain any positive instance the bag will be considered as negative configuration parameters positiveclass the desired class to be kept in the positive bag bagcount the number of training bags valbagcount the number of validation bags bagsize the number of instances in a bag plotsize the number of bags to plot ensembleavgcount the number of models to create and average together optional often results in better performance set to 1 for single model prepare bags since the attention operator is a permutationinvariant operator an instance with a positive class label is randomly placed among the instances in the positive bag set up bags normalize input data count positive samples pick a fixed size random subset of samples by default all bags are labeled as 0 check if there is at least a positive class in the bag positive bag will be labeled as 1 load the mnist dataset create training data create validation data create the model we will now build the attention layer prepare some utilities then build and train the entire model attention operator implementation the output size of this layer is decided by the size of a single bag the attention mechanism uses a weighted average of instances in a bag in which the sum of the weights must equal to 1 invariant of the bag size the weight matrices parameters are w and v to include positive and negative values hyperbolic tangent elementwise nonlinearity is utilized a gated attention mechanism can be used to deal with complex relations another weight matrix u is added to the computation a sigmoid nonlinearity is used to overcome approximately linear behavior for x 1 1 by hyperbolic tangent nonlinearity implementation of the attentionbased deep mil layer args weightparamsdim positive integer dimension of the weight matrix kernelinitializer initializer for the kernel matrix kernelregularizer regularizer function applied to the kernel matrix usegated boolean whether or not to use the gated mechanism returns list of 2d tensors with bagsize length the tensors are the attention scores after softmax with shape batchsize 1 input shape list of 2d tensors with shape batchsize inputdim assigning variables from the number of inputs stack instances into a single tensor apply softmax over instances such that the output summation is equal to 1 split to recreate the same array of tensors we had as inputs reserve incase gated mechanism used tanhvhkt for learning nonlinear relations efficiently wttanhvhkt wttanhvhktsigmoiduhkt visualizer tool plot the number of bags given by plotsize with respect to the class moreover if activated the class label prediction with its associated instance score for each bag after the model has been trained can be seen utility for plotting bags and attention weights args data input data that contains the bags of instances labels the associated bag labels of the input data bagclass string name of the desired bag class the options are positive or negative predictions class labels model predictions if you don t specify anything ground truth labels will be used attentionweights attention weights for each instance within the input data if you don t specify anything the values won t be displayed plot some of validation data bags per class create model first we will create some embeddings per instance invoke the attention operator and then use the softmax function to output the class probabilities extract features from inputs invoke the attention layer multiply attention weights with the input layers concatenate layers classification output node class weights since this kind of problem could simply turn into imbalanced data classification problem class weighting should be considered let s say there are 1000 bags there often could be cases were 90 of the bags do not contain any positive label and 10 do such data can be referred to as imbalanced data using class weights the model will tend to give a higher weight to the rare class count number of postive and negative bags build class weight dictionary build and train model the model is built and trained in this section train model prepare callbacks path where to save best weights take the file name from the wrapper initialize model checkpoint callback initialize early stopping callback the model performance is monitored across the validation data and stops training when the generalization error cease to decrease compile model fit model load best weights building models show single model architecture training models model evaluation the models are now ready for evaluation with each model we also create an associated intermediate model to get the weights from the attention layer we will compute a prediction for each of our ensembleavgcount models and average them together for our final prediction collect info per model predict output classes on data create intermediate model to get mil attention layer weights predict mil attention layer weights evaluate and predict classes and attention scores on validation data plot some results from our validation data conclusion from the above plot you can notice that the weights always sum to 1 in a positively predict bag the instance which resulted in the positive labeling will have a substantially higher attention score than the rest of the bag however in a negatively predicted bag there are two cases all instances will have approximately similar scores an instance will have relatively higher score but not as high as of a positive instance this is because the feature space of this instance is close to that of the positive instance remarks if the model is overfit the weights will be equally distributed for all bags hence the regularization techniques are necessary in the paper the bag sizes can differ from one bag to another for simplicity the bag sizes are fixed here in order not to rely on the random initial weights of a single model averaging ensemble methods should be considered introduction what is multiple instance learning mil usually with supervised learning algorithms the learner receives labels for a set of instances in the case of mil the learner receives labels for a set of bags each of which contains a set of instances the bag is labeled positive if it contains at least one positive instance and negative if it does not contain any motivation it is often assumed in image classification tasks that each image clearly represents a class label in medical imaging e g computational pathology etc an entire image is represented by a single class label cancerous non cancerous or a region of interest could be given however one will be interested in knowing which patterns in the image is actually causing it to belong to that class in this context the image s will be divided and the subimages will form the bag of instances therefore the goals are to 1 learn a model to predict a class label for a bag of instances 2 find out which instances within the bag caused a position class label prediction implementation the following steps describe how the model works 1 the feature extractor layers extract feature embeddings 2 the embeddings are fed into the mil attention layer to get the attention scores the layer is designed as permutation invariant 3 input features and their corresponding attention scores are multiplied together 4 the resulting output is passed to a softmax function for classification references attention based deep multiple instance learning https arxiv org abs 1802 04712 some of the attention operator code implementation was inspired from https github com utayao atten_deep_mil imbalanced data tutorial https www tensorflow org tutorials structured_data imbalanced_data by tensorflow setup create dataset we will create a set of bags and assign their labels according to their contents if at least one positive instance is available in a bag the bag is considered as a positive bag if it does not contain any positive instance the bag will be considered as negative configuration parameters positive_class the desired class to be kept in the positive bag bag_count the number of training bags val_bag_count the number of validation bags bag_size the number of instances in a bag plot_size the number of bags to plot ensemble_avg_count the number of models to create and average together optional often results in better performance set to 1 for single model prepare bags since the attention operator is a permutation invariant operator an instance with a positive class label is randomly placed among the instances in the positive bag set up bags normalize input data count positive samples pick a fixed size random subset of samples by default all bags are labeled as 0 check if there is at least a positive class in the bag positive bag will be labeled as 1 load the mnist dataset create training data create validation data create the model we will now build the attention layer prepare some utilities then build and train the entire model attention operator implementation the output size of this layer is decided by the size of a single bag the attention mechanism uses a weighted average of instances in a bag in which the sum of the weights must equal to 1 invariant of the bag size the weight matrices parameters are w and v to include positive and negative values hyperbolic tangent element wise non linearity is utilized a gated attention mechanism can be used to deal with complex relations another weight matrix u is added to the computation a sigmoid non linearity is used to overcome approximately linear behavior for x 1 1 by hyperbolic tangent non linearity implementation of the attention based deep mil layer args weight_params_dim positive integer dimension of the weight matrix kernel_initializer initializer for the kernel matrix kernel_regularizer regularizer function applied to the kernel matrix use_gated boolean whether or not to use the gated mechanism returns list of 2d tensors with bag_size length the tensors are the attention scores after softmax with shape batch_size 1 input shape list of 2d tensors with shape batch_size input_dim assigning variables from the number of inputs stack instances into a single tensor apply softmax over instances such that the output summation is equal to 1 split to recreate the same array of tensors we had as inputs reserve in case gated mechanism used tanh v h_k t for learning non linear relations efficiently w t tanh v h_k t w t tanh v h_k t sigmoid u h_k t visualizer tool plot the number of bags given by plot_size with respect to the class moreover if activated the class label prediction with its associated instance score for each bag after the model has been trained can be seen utility for plotting bags and attention weights args data input data that contains the bags of instances labels the associated bag labels of the input data bag_class string name of the desired bag class the options are positive or negative predictions class labels model predictions if you don t specify anything ground truth labels will be used attention_weights attention weights for each instance within the input data if you don t specify anything the values won t be displayed todo plot some of validation data bags per class create model first we will create some embeddings per instance invoke the attention operator and then use the softmax function to output the class probabilities extract features from inputs invoke the attention layer multiply attention weights with the input layers concatenate layers classification output node class weights since this kind of problem could simply turn into imbalanced data classification problem class weighting should be considered let s say there are 1000 bags there often could be cases were 90 of the bags do not contain any positive label and 10 do such data can be referred to as imbalanced data using class weights the model will tend to give a higher weight to the rare class count number of postive and negative bags build class weight dictionary build and train model the model is built and trained in this section train model prepare callbacks path where to save best weights take the file name from the wrapper initialize model checkpoint callback initialize early stopping callback the model performance is monitored across the validation data and stops training when the generalization error cease to decrease compile model fit model load best weights building model s show single model architecture training model s model evaluation the models are now ready for evaluation with each model we also create an associated intermediate model to get the weights from the attention layer we will compute a prediction for each of our ensemble_avg_count models and average them together for our final prediction collect info per model predict output classes on data create intermediate model to get mil attention layer weights predict mil attention layer weights evaluate and predict classes and attention scores on validation data plot some results from our validation data conclusion from the above plot you can notice that the weights always sum to 1 in a positively predict bag the instance which resulted in the positive labeling will have a substantially higher attention score than the rest of the bag however in a negatively predicted bag there are two cases all instances will have approximately similar scores an instance will have relatively higher score but not as high as of a positive instance this is because the feature space of this instance is close to that of the positive instance remarks if the model is overfit the weights will be equally distributed for all bags hence the regularization techniques are necessary in the paper the bag sizes can differ from one bag to another for simplicity the bag sizes are fixed here in order not to rely on the random initial weights of a single model averaging ensemble methods should be considered
import numpy as np
import keras
from keras import layers
from keras import ops
from tqdm import tqdm
from matplotlib import pyplot as plt

plt.style.use("ggplot")



POSITIVE_CLASS = 1
BAG_COUNT = 1000
VAL_BAG_COUNT = 300
BAG_SIZE = 3
PLOT_SIZE = 3
ENSEMBLE_AVG_COUNT = 1




def create_bags(
    input_data, input_labels, positive_class, bag_count, instance_count
):
    
    bags = []
    bag_labels = []

    
    input_data = np.divide(input_data, 255.0)

    
    count = 0

    for _ in range(bag_count):
        
        index = np.random.choice(
            input_data.shape[0], instance_count, replace=False
        )
        instances_data = input_data[index]
        instances_labels = input_labels[index]

        
        bag_label = 0

        
        if positive_class in instances_labels:
            
            bag_label = 1
            count += 1

        bags.append(instances_data)
        bag_labels.append(np.array([bag_label]))

    print(f"Positive bags: {count}")
    print(f"Negative bags: {bag_count - count}")

    return (list(np.swapaxes(bags, 0, 1)), np.array(bag_labels))



(x_train, y_train), (x_val, y_val) = keras.datasets.mnist.load_data()


train_data, train_labels = create_bags(
    x_train, y_train, POSITIVE_CLASS, BAG_COUNT, BAG_SIZE
)


val_data, val_labels = create_bags(
    x_val, y_val, POSITIVE_CLASS, VAL_BAG_COUNT, BAG_SIZE
)




class MILAttentionLayer(layers.Layer):
    

    def __init__(
        self,
        weight_params_dim,
        kernel_initializer="glorot_uniform",
        kernel_regularizer=None,
        use_gated=False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.weight_params_dim = weight_params_dim
        self.use_gated = use_gated

        self.kernel_initializer = keras.initializers.get(kernel_initializer)
        self.kernel_regularizer = keras.regularizers.get(kernel_regularizer)

        self.v_init = self.kernel_initializer
        self.w_init = self.kernel_initializer
        self.u_init = self.kernel_initializer

        self.v_regularizer = self.kernel_regularizer
        self.w_regularizer = self.kernel_regularizer
        self.u_regularizer = self.kernel_regularizer

    def build(self, input_shape):
        
        
        input_dim = input_shape[0][1]

        self.v_weight_params = self.add_weight(
            shape=(input_dim, self.weight_params_dim),
            initializer=self.v_init,
            name="v",
            regularizer=self.v_regularizer,
            trainable=True,
        )

        self.w_weight_params = self.add_weight(
            shape=(self.weight_params_dim, 1),
            initializer=self.w_init,
            name="w",
            regularizer=self.w_regularizer,
            trainable=True,
        )

        if self.use_gated:
            self.u_weight_params = self.add_weight(
                shape=(input_dim, self.weight_params_dim),
                initializer=self.u_init,
                name="u",
                regularizer=self.u_regularizer,
                trainable=True,
            )
        else:
            self.u_weight_params = None

        self.input_built = True

    def call(self, inputs):
        
        instances = [
            self.compute_attention_scores(instance) for instance in inputs
        ]

        
        instances = ops.stack(instances)

        
        alpha = ops.softmax(instances, axis=0)

        
        return [alpha[i] for i in range(alpha.shape[0])]

    def compute_attention_scores(self, instance):
        
        original_instance = instance

        
        instance = ops.tanh(
            ops.tensordot(instance, self.v_weight_params, axes=1)
        )

        
        if self.use_gated:
            instance = instance * ops.sigmoid(
                ops.tensordot(original_instance, self.u_weight_params, axes=1)
            )

        
        return ops.tensordot(instance, self.w_weight_params, axes=1)





def plot(data, labels, bag_class, predictions=None, attention_weights=None):
    
    return  
    labels = np.array(labels).reshape(-1)

    if bag_class == "positive":
        if predictions is not None:
            labels = np.where(predictions.argmax(1) == 1)[0]
            bags = np.array(data)[:, labels[0:PLOT_SIZE]]

        else:
            labels = np.where(labels == 1)[0]
            bags = np.array(data)[:, labels[0:PLOT_SIZE]]

    elif bag_class == "negative":
        if predictions is not None:
            labels = np.where(predictions.argmax(1) == 0)[0]
            bags = np.array(data)[:, labels[0:PLOT_SIZE]]
        else:
            labels = np.where(labels == 0)[0]
            bags = np.array(data)[:, labels[0:PLOT_SIZE]]

    else:
        print(f"There is no class {bag_class}")
        return

    print(f"The bag class label is {bag_class}")
    for i in range(PLOT_SIZE):
        figure = plt.figure(figsize=(8, 8))
        print(f"Bag number: {labels[i]}")
        for j in range(BAG_SIZE):
            image = bags[j][i]
            figure.add_subplot(1, BAG_SIZE, j + 1)
            plt.grid(False)
            if attention_weights is not None:
                plt.title(np.around(attention_weights[labels[i]][j], 2))
            plt.imshow(image)
        plt.show()



plot(val_data, val_labels, "positive")
plot(val_data, val_labels, "negative")




def create_model(instance_shape):
    
    inputs, embeddings = [], []
    shared_dense_layer_1 = layers.Dense(128, activation="relu")
    shared_dense_layer_2 = layers.Dense(64, activation="relu")
    for _ in range(BAG_SIZE):
        inp = layers.Input(instance_shape)
        flatten = layers.Flatten()(inp)
        dense_1 = shared_dense_layer_1(flatten)
        dense_2 = shared_dense_layer_2(dense_1)
        inputs.append(inp)
        embeddings.append(dense_2)

    
    alpha = MILAttentionLayer(
        weight_params_dim=256,
        kernel_regularizer=keras.regularizers.L2(0.01),
        use_gated=True,
        name="alpha",
    )(embeddings)

    
    multiply_layers = [
        layers.multiply([alpha[i], embeddings[i]]) for i in range(len(alpha))
    ]

    
    concat = layers.concatenate(multiply_layers, axis=1)

    
    output = layers.Dense(2, activation="softmax")(concat)

    return keras.Model(inputs, output)





def compute_class_weights(labels):
    
    negative_count = len(np.where(labels == 0)[0])
    positive_count = len(np.where(labels == 1)[0])
    total_count = negative_count + positive_count

    
    return {
        0: (1 / negative_count) * (total_count / 2),
        1: (1 / positive_count) * (total_count / 2),
    }





def train(train_data, train_labels, val_data, val_labels, model):
    
    
    

    
    file_path = "/tmp/best_model.weights.h5"

    
    model_checkpoint = keras.callbacks.ModelCheckpoint(
        file_path,
        monitor="val_loss",
        verbose=0,
        mode="min",
        save_best_only=True,
        save_weights_only=True,
    )

    
    
    
    early_stopping = keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=10, mode="min"
    )

    
    model.compile(
        optimizer="adam",
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"],
    )

    
    model.fit(
        train_data,
        train_labels,
        validation_data=(val_data, val_labels),
        epochs=20,
        class_weight=compute_class_weights(train_labels),
        batch_size=1,
        callbacks=[early_stopping, model_checkpoint],
        verbose=0,
    )

    
    model.load_weights(file_path)

    return model



instance_shape = train_data[0][0].shape
models = [create_model(instance_shape) for _ in range(ENSEMBLE_AVG_COUNT)]


print(models[0].summary())


trained_models = [
    train(train_data, train_labels, val_data, val_labels, model)
    for model in tqdm(models)
]




def predict(data, labels, trained_models):
    
    models_predictions = []
    models_attention_weights = []
    models_losses = []
    models_accuracies = []

    for model in trained_models:
        
        predictions = model.predict(data)
        models_predictions.append(predictions)

        
        intermediate_model = keras.Model(
            model.input, model.get_layer("alpha").output
        )

        
        intermediate_predictions = intermediate_model.predict(data)

        attention_weights = np.squeeze(
            np.swapaxes(intermediate_predictions, 1, 0)
        )
        models_attention_weights.append(attention_weights)

        loss, accuracy = model.evaluate(data, labels, verbose=0)
        models_losses.append(loss)
        models_accuracies.append(accuracy)

    print(
        f"The average loss and accuracy are {np.sum(models_losses, axis=0) / ENSEMBLE_AVG_COUNT:.2f}"
        f" and {100 * np.sum(models_accuracies, axis=0) / ENSEMBLE_AVG_COUNT:.2f} % resp."
    )

    return (
        np.sum(models_predictions, axis=0) / ENSEMBLE_AVG_COUNT,
        np.sum(models_attention_weights, axis=0) / ENSEMBLE_AVG_COUNT,
    )



class_predictions, attention_params = predict(
    val_data, val_labels, trained_models
)


plot(
    val_data,
    val_labels,
    "positive",
    predictions=class_predictions,
    attention_weights=attention_params,
)
plot(
    val_data,
    val_labels,
    "negative",
    predictions=class_predictions,
    attention_weights=attention_params,
)