基于注意力的深度多实例学习（MIL）分类

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]

        # 默认情况下，所有袋子标记为 0。
        bag_label = 0

        # 检查袋中是否至少有一个正类样本。
        if positive_class in instances_labels:
            # 正袋将标记为 1。
            bag_label = 1
            count += 1

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

    print(f"正袋子: {count}")
    print(f"负袋子: {bag_count - count}")

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


# 加载 MNIST 数据集。
(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):
    """实现基于注意力的深度MIL层。

    Args:
      weight_params_dim: 正整数。权重矩阵的维度。
      kernel_initializer: `kernel`矩阵的初始化器。
      kernel_regularizer: 应用于`kernel`矩阵的正则化函数。
      use_gated: 布尔值，是否使用门控机制。

    Returns:
      长度为 BAG_SIZE 的 2D 张量列表。
      这些张量是经过softmax处理后的注意力分数，形状为 `(batch_size, 1)`。
    """

    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):
        # 输入形状。
        # 形状为: (batch_size, input_dim) 的 2D 张量列表。
        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)

        # 在实例上应用softmax，使输出总和等于1。
        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

        # tanh(v*h_k^T)
        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)
            )

        # w^T*(tanh(v*h_k^T)) / w^T*(tanh(v*h_k^T)*sigmoid(u*h_k^T))
        return ops.tensordot(instance, self.w_weight_params, axes=1)

def plot(data, labels, bag_class, predictions=None, attention_weights=None):
    """ "用于绘制包和注意力权重的实用工具。

    参数：
      data: 包含实例袋的输入数据。
      labels: 输入数据的关联包标签。
      bag_class: 所需包类的字符串名称。
        选项包括：“positive”或“negative”。
      predictions: 类别标签模型预测。
      如果您不指定任何内容，将使用真实标签。
      attention_weights: 输入数据中每个实例的注意力权重。
      如果您不指定任何内容，则不会显示这些值。
    """
    return  ## TODO
    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"没有类 {bag_class}")
        return

    print(f"包类标签是 {bag_class}")
    for i in range(PLOT_SIZE):
        figure = plt.figure(figsize=(8, 8))
        print(f"包编号: {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)

        # 创建中间模型以获取 MIL 注意力层权重。
        intermediate_model = keras.Model(model.input, model.get_layer("alpha").output)

        # 预测 MIL 注意力层权重。
        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"平均损失和准确率为 {np.sum(models_losses, axis=0) / ENSEMBLE_AVG_COUNT:.2f}"
        f" 和 {100 * np.sum(models_accuracies, axis=0) / ENSEMBLE_AVG_COUNT:.2f} % "
        "响应。"
    )

    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,
)