Barlow Twins 用于对比自监督学习

(pred_1.T @ pred_2) / batch_size

c[i][i] = 1
c[i][j] = 0

其中：
  c 是交叉相关矩阵
  i 是一个表示的神经元的索引
  j 是第二个表示的神经元的索引

!pip install tensorflow-addons

import os

# 略微加快的改进，在第一个周期减少了30秒，在周期时间上减少了1-2秒
# 总体节省约5分钟的训练时间

# 为gpu私有分配两个线程，这样可以更快地完成更多操作
os.environ["TF_GPU_THREAD_MODE"] = "gpu_private"

import tensorflow as tf  # 框架
from tensorflow import keras  # 用于tf.keras
import tensorflow_addons as tfa  # LAMB优化器和gaussian_blur_2d函数
import numpy as np  # np.random.random
import matplotlib.pyplot as plt  # 图表
import datetime  # tensorboard日志命名

# XLA优化以获得更快的性能（节省总时间高达10-15分钟）
tf.config.optimizer.set_jit(True)

[
    (train_features, train_labels),
    (test_features, test_labels),
] = keras.datasets.cifar10.load_data()

train_features = train_features / 255.0
test_features = test_features / 255.0

# 数据集的批次大小
BATCH_SIZE = 512
# 图像的宽度和高度
IMAGE_SIZE = 32

class Augmentation(keras.layers.Layer):
    """Base augmentation class.

    Base augmentation class. Contains the random_execute method.

    Methods:
        random_execute: method that returns true or false based
          on a probability. Used to determine whether an augmentation
          will be run.
    """

    def __init__(self):
        super().__init__()

    @tf.function
    def random_execute(self, prob: float) -> bool:
        """random_execute function.

        Arguments:
            prob: a float value from 0-1 that determines the
              probability.

        Returns:
            returns true or false based on the probability.
        """

        return tf.random.uniform([], minval=0, maxval=1) < prob


class RandomToGrayscale(Augmentation):
    """RandomToGrayscale class.

    RandomToGrayscale class. Randomly makes an image
    grayscaled based on the random_execute method. There
    is a 20% chance that an image will be grayscaled.

    Methods:
        call: method that grayscales an image 20% of
          the time.
    """

    @tf.function
    def call(self, x: tf.Tensor) -> tf.Tensor:
        """call function.

        Arguments:
            x: a tf.Tensor representing the image.

        Returns:
            returns a grayscaled version of the image 20% of the time
              and the original image 80% of the time.
        """

        if self.random_execute(0.2):
            x = tf.image.rgb_to_grayscale(x)
            x = tf.tile(x, [1, 1, 3])
        return x


class RandomColorJitter(Augmentation):
    """RandomColorJitter class.

    RandomColorJitter class. Randomly adds color jitter to an image.
    Color jitter means to add random brightness, contrast,
    saturation, and hue to an image. There is a 80% chance that an
    image will be randomly color-jittered.

    Methods:
        call: method that color-jitters an image 80% of
          the time.
    """

    @tf.function
    def call(self, x: tf.Tensor) -> tf.Tensor:
        """call function.

        Adds color jitter to image, including:
          Brightness change by a max-delta of 0.8
          Contrast change by a max-delta of 0.8
          Saturation change by a max-delta of 0.8
          Hue change by a max-delta of 0.2
        Originally, the same deltas of the original paper
        were used, but a performance boost of almost 2% was found
        when doubling them.

        Arguments:
            x: a tf.Tensor representing the image.

        Returns:
            returns a color-jittered version of the image 80% of the time
              and the original image 20% of the time.
        """

        if self.random_execute(0.8):
            x = tf.image.random_brightness(x, 0.8)
            x = tf.image.random_contrast(x, 0.4, 1.6)
            x = tf.image.random_saturation(x, 0.4, 1.6)
            x = tf.image.random_hue(x, 0.2)
        return x


class RandomFlip(Augmentation):
    """RandomFlip class.

    RandomFlip class. Randomly flips image horizontally. There is a 50%
    chance that an image will be randomly flipped.

    Methods:
        call: method that flips an image 50% of
          the time.
    """

    @tf.function
    def call(self, x: tf.Tensor) -> tf.Tensor:
        """call function.

        Randomly flips the image.

        Arguments:
            x: a tf.Tensor representing the image.

        Returns:
            returns a flipped version of the image 50% of the time
              and the original image 50% of the time.
        """

        if self.random_execute(0.5):
            x = tf.image.random_flip_left_right(x)
        return x


class RandomResizedCrop(Augmentation):
    """RandomResizedCrop class.

    RandomResizedCrop class. Randomly crop an image to a random size,
    then resize the image back to the original size.

    Attributes:
        image_size: The dimension of the image

    Methods:
        __call__: method that does random resize crop to the image.
    """

    def __init__(self, image_size):
        super().__init__()
        self.image_size = image_size

    def call(self, x: tf.Tensor) -> tf.Tensor:
        """call function.

        Does random resize crop by randomly cropping an image to a random
        size 75% - 100% the size of the image. Then resizes it.

        Arguments:
            x: a tf.Tensor representing the image.

        Returns:
            returns a randomly cropped image.
        """

        rand_size = tf.random.uniform(
            shape=[],
            minval=int(0.75 * self.image_size),
            maxval=1 * self.image_size,
            dtype=tf.int32,
        )

        crop = tf.image.random_crop(x, (rand_size, rand_size, 3))
        crop_resize = tf.image.resize(crop, (self.image_size, self.image_size))
        return crop_resize


class RandomSolarize(Augmentation):
    """RandomSolarize class.

    RandomSolarize class. Randomly solarizes an image.
    Solarization is when pixels accidentally flip to an inverted state.

    Methods:
        call: method that does random solarization 20% of the time.
    """

    @tf.function
    def call(self, x: tf.Tensor) -> tf.Tensor:
        """call function.

        Randomly solarizes the image.

        Arguments:
            x: a tf.Tensor representing the image.

        Returns:
            returns a solarized version of the image 20% of the time
              and the original image 80% of the time.
        """

        if self.random_execute(0.2):
            # flips abnormally low pixels to abnormally high pixels
            x = tf.where(x < 10, x, 255 - x)
        return x


class RandomBlur(Augmentation):
    """RandomBlur class.

    RandomBlur class. Randomly blurs an image.

    Methods:
        call: method that does random blur 20% of the time.
    """

    @tf.function
    def call(self, x: tf.Tensor) -> tf.Tensor:
        """call function.

        Randomly solarizes the image.

        Arguments:
            x: a tf.Tensor representing the image.

        Returns:
            returns a blurred version of the image 20% of the time
              and the original image 80% of the time.
        """

        if self.random_execute(0.2):
            s = np.random.random()
            return tfa.image.gaussian_filter2d(image=x, sigma=s)
        return x


class RandomAugmentor(keras.Model):
    """RandomAugmentor class.

    RandomAugmentor class. Chains all the augmentations into
    one pipeline.

    Attributes:
        image_size: An integer represing the width and height
          of the image. Designed to be used for square images.
        random_resized_crop: Instance variable representing the
          RandomResizedCrop layer.
        random_flip: Instance variable representing the
          RandomFlip layer.
        random_color_jitter: Instance variable representing the
          RandomColorJitter layer.
        random_blur: Instance variable representing the
          RandomBlur layer
        random_to_grayscale: Instance variable representing the
          RandomToGrayscale layer
        random_solarize: Instance variable representing the
          RandomSolarize layer

    Methods:
        call: chains layers in pipeline together
    """

    def __init__(self, image_size: int):
        super().__init__()

        self.image_size = image_size
        self.random_resized_crop = RandomResizedCrop(image_size)
        self.random_flip = RandomFlip()
        self.random_color_jitter = RandomColorJitter()
        self.random_blur = RandomBlur()
        self.random_to_grayscale = RandomToGrayscale()
        self.random_solarize = RandomSolarize()

    def call(self, x: tf.Tensor) -> tf.Tensor:
        x = self.random_resized_crop(x)
        x = self.random_flip(x)
        x = self.random_color_jitter(x)
        x = self.random_blur(x)
        x = self.random_to_grayscale(x)
        x = self.random_solarize(x)

        x = tf.clip_by_value(x, 0, 1)
        return x


bt_augmentor = RandomAugmentor(IMAGE_SIZE)

class BTDatasetCreator:
    """Barlow双胞胎数据集创建类。

    BTDatasetCreator类。负责创建
    Barlow双胞胎的数据集。

    属性：
        options: tf.data.Options，用于配置可能提高性能的设置。
        seed: 随机种子，用于洗牌。用于同步两个
          增强版本。
        augmentor: 用于增强的增强器。

    方法：
        __call__: 创建Barlow数据集。
        augmented_version: 创建数据集的一半。
    """

    def __init__(self, augmentor: RandomAugmentor, seed: int = 1024):
        self.options = tf.data.Options()
        self.options.threading.max_intra_op_parallelism = 1
        self.seed = seed
        self.augmentor = augmentor

    def augmented_version(self, ds: list) -> tf.data.Dataset:
        return (
            tf.data.Dataset.from_tensor_slices(ds)
            .shuffle(1000, seed=self.seed)
            .map(self.augmentor, num_parallel_calls=tf.data.AUTOTUNE)
            .batch(BATCH_SIZE, drop_remainder=True)
            .prefetch(tf.data.AUTOTUNE)
            .with_options(self.options)
        )

    def __call__(self, ds: list) -> tf.data.Dataset:
        a1 = self.augmented_version(ds)
        a2 = self.augmented_version(ds)

        return tf.data.Dataset.zip((a1, a2)).with_options(self.options)


augment_versions = BTDatasetCreator(bt_augmentor)(train_features)

sample_augment_versions = iter(augment_versions)


def plot_values(batch: tuple):
    fig, axs = plt.subplots(3, 3)
    fig1, axs1 = plt.subplots(3, 3)

    fig.suptitle("增强 1")
    fig1.suptitle("增强 2")

    a1, a2 = batch

    # 在两个表上绘制图像
    for i in range(3):
        for j in range(3):
            # 改变（添加 / 255）
            axs[i][j].imshow(a1[3 * i + j])
            axs[i][j].axis("off")
            axs1[i][j].imshow(a2[3 * i + j])
            axs1[i][j].axis("off")

    plt.show()


plot_values(next(sample_augment_versions))

y_a: 原始图像的第一个增强版本。
y_b: 原始图像的第二个增强版本。
z_a: y_a的模型表示（嵌入）。
z_b: y_b的模型表示（嵌入）。
z_a_norm: z_a的归一化。
z_b_norm: z_b的归一化。
c: 交叉相关矩阵。
c_diff: 损失的对角部分（不变性项）。
off_diag: 损失的非对角部分（冗余减少项）。

class BarlowLoss(keras.losses.Loss):
    """BarlowLoss 类。

    BarlowLoss 类。基于交叉相关矩阵创建损失函数。

    属性：
        batch_size: 数据集的批次大小
        lambda_amt: lambda 的值（用于 cross_corr_matrix_loss）

    方法：
        __init__: 获取实例变量
        call: 根据交叉相关矩阵获取损失
          make_diag_zeros: 用于计算损失函数的非对角部分；将对角线设为零
        cross_corr_matrix_loss: 基于交叉相关矩阵创建损失。
    """

    def __init__(self, batch_size: int):
        """__init__ 方法。

        获取实例变量

        参数：
            batch_size: 一个整数值，表示数据集的批次大小。用于交叉相关矩阵计算。
        """

        super().__init__()
        self.lambda_amt = 5e-3
        self.batch_size = batch_size

    def get_off_diag(self, c: tf.Tensor) -> tf.Tensor:
        """get_off_diag 方法。

        将交叉相关矩阵的对角线设为零。
        这用于损失函数的非对角部分，我们对非对角值取平方并求和。

        参数：
            c: 表示交叉相关矩阵的 tf.tensor

        返回：
            返回一个 tf.tensor，表示对角线为零的交叉相关矩阵。
        """

        zero_diag = tf.zeros(c.shape[-1])
        return tf.linalg.set_diag(c, zero_diag)

    def cross_corr_matrix_loss(self, c: tf.Tensor) -> tf.Tensor:
        """cross_corr_matrix_loss 方法。

        根据交叉相关矩阵获取损失。
        我们希望对角线为 1，其他所有值为零，以显示这两幅增强图像相似。

        损失函数过程：
        取交叉相关矩阵的对角线，减去 1，然后平方该值以避免负值。

        取 cc 矩阵的非对角部分（见 get_off_diag()），
        将这些值平方以消除负值并增加值，
        并乘以一个 lambda 值，使其与对角线的值相等（非对角线的值比对角线值多）

        将第一部分和第二部分相加，然后再相加。

        参数：
            c: 表示交叉相关矩阵的 tf.tensor

        返回：
            返回一个 tf.tensor，表示对角线为零的交叉相关矩阵。
        """

        # 将对角线减去一并平方（第一部分）
        c_diff = tf.pow(tf.linalg.diag_part(c) - 1, 2)

        # 取非对角线，平方并乘以 lambda（第二部分）
        off_diag = tf.pow(self.get_off_diag(c), 2) * self.lambda_amt

        # 将第一部分和第二部分相加
        loss = tf.reduce_sum(c_diff) + tf.reduce_sum(off_diag)

        return loss

    def normalize(self, output: tf.Tensor) -> tf.Tensor:
        """normalize 方法。

        归一化模型预测。

        参数：
            output: 模型预测。

        返回：
            返回归一化后的模型预测。
        """

        return (output - tf.reduce_mean(output, axis=0)) / tf.math.reduce_std(
            output, axis=0
        )

    def cross_corr_matrix(self, z_a_norm: tf.Tensor, z_b_norm: tf.Tensor) -> tf.Tensor:
        """cross_corr_matrix 方法。

        从预测中创建交叉相关矩阵。
        它转置第一个预测并与第二个预测相乘，创建形状为 (n_dense_units, n_dense_units) 的矩阵。
        有关更多信息，请参见 build_twin()。然后将其除以批次大小。

        参数：
            z_a_norm: 第一个预测的归一化版本。
            z_b_norm: 第二个预测的归一化版本。

        返回：
            返回一个交叉相关矩阵。
        """
        return (tf.transpose(z_a_norm) @ z_b_norm) / self.batch_size

    def call(self, z_a: tf.Tensor, z_b: tf.Tensor) -> tf.Tensor:
        """call 方法。

        计算交叉相关损失。使用 CreateCrossCorr 类生成交叉相关矩阵，然后找到损失并返回（见 cross_corr_matrix_loss()）。

        参数：
            z_a: 第一组增强数据的预测。
            z_b: 第二组增强数据的预测。

        返回：
            返回一个（秩为 0 的）tf.Tensor，表示损失。
        """

        z_a_norm, z_b_norm = self.normalize(z_a), self.normalize(z_b)
        c = self.cross_corr_matrix(z_a_norm, z_b_norm)
        loss = self.cross_corr_matrix_loss(c)
        return loss

class ResNet34:
    """Resnet34类。

        负责Resnet 34架构。
    修改自
    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2。
    https://www.analyticsvidhya.com/blog/2021/08/how-to-code-your-resnet-from-scratch-in-tensorflow/#h2_2。
        更多信息请查看他们的网站。
    """

    def identity_block(self, x, filter):
        # 将张量复制到名为x_skip的变量
        x_skip = x
        # 层1
        x = tf.keras.layers.Conv2D(filter, (3, 3), padding="same")(x)
        x = tf.keras.layers.BatchNormalization(axis=3)(x)
        x = tf.keras.layers.Activation("relu")(x)
        # 层2
        x = tf.keras.layers.Conv2D(filter, (3, 3), padding="same")(x)
        x = tf.keras.layers.BatchNormalization(axis=3)(x)
        # 添加残差
        x = tf.keras.layers.Add()([x, x_skip])
        x = tf.keras.layers.Activation("relu")(x)
        return x

    def convolutional_block(self, x, filter):
        # 将张量复制到名为x_skip的变量
        x_skip = x
        # 层1
        x = tf.keras.layers.Conv2D(filter, (3, 3), padding="same", strides=(2, 2))(x)
        x = tf.keras.layers.BatchNormalization(axis=3)(x)
        x = tf.keras.layers.Activation("relu")(x)
        # 层2
        x = tf.keras.layers.Conv2D(filter, (3, 3), padding="same")(x)
        x = tf.keras.layers.BatchNormalization(axis=3)(x)
        # 使用conv(1,1)处理残差
        x_skip = tf.keras.layers.Conv2D(filter, (1, 1), strides=(2, 2))(x_skip)
        # 添加残差
        x = tf.keras.layers.Add()([x, x_skip])
        x = tf.keras.layers.Activation("relu")(x)
        return x

    def __call__(self, shape=(32, 32, 3)):
        # 步骤1（设置输入层）
        x_input = tf.keras.layers.Input(shape)
        x = tf.keras.layers.ZeroPadding2D((3, 3))(x_input)
        # 步骤2（初始卷积层和最大池化层）
        x = tf.keras.layers.Conv2D(64, kernel_size=7, strides=2, padding="same")(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation("relu")(x)
        x = tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding="same")(x)
        # 定义子块的大小和初始滤波器大小
        block_layers = [3, 4, 6, 3]
        filter_size = 64
        # 步骤3 添加Resnet块
        for i in range(4):
            if i == 0:
                # 对于子块1不需要残差/卷积块
                for j in range(block_layers[i]):
                    x = self.identity_block(x, filter_size)
            else:
                # 一个残差/卷积块后跟身份块
                # 滤波器大小将以2的倍数增加
                filter_size = filter_size * 2
                x = self.convolutional_block(x, filter_size)
                for j in range(block_layers[i] - 1):
                    x = self.identity_block(x, filter_size)
        # 步骤4 结束密集网络
        x = tf.keras.layers.AveragePooling2D((2, 2), padding="same")(x)
        x = tf.keras.layers.Flatten()(x)
        model = tf.keras.models.Model(inputs=x_input, outputs=x, name="ResNet34")
        return model

def build_twin() -> keras.Model:
    """build_twin方法。

    构建一个包含编码器（resnet-34）和投影器的barlow twins模型，
    用于为图像生成嵌入

    返回：
        返回一个barlow twins模型
    """

    # 投影器中的稠密神经元数量
    n_dense_neurons = 5000

    # 编码器网络
    resnet = ResNet34()()
    last_layer = resnet.layers[-1].output

    # 投影网络的中间层
    n_layers = 2
    for i in range(n_layers):
        dense = tf.keras.layers.Dense(n_dense_neurons, name=f"projector_dense_{i}")
        if i == 0:
            x = dense(last_layer)
        else:
            x = dense(x)
        x = tf.keras.layers.BatchNormalization(name=f"projector_bn_{i}")(x)
        x = tf.keras.layers.ReLU(name=f"projector_relu_{i}")(x)

    x = tf.keras.layers.Dense(n_dense_neurons, name=f"projector_dense_{n_layers}")(x)

    model = keras.Model(resnet.input, x)
    return model

class BarlowModel(keras.Model):
    """BarlowModel 类。

    BarlowModel 类。负责进行预测和处理
    优化器的梯度下降。

    属性：
        model: barlow 模型架构。
        loss_tracker: 损失指标。

    方法：
        train_step: 一个训练步骤；进行模型预测、计算损失和
            优化器步骤。
        metrics: 返回指标。
    """

    def __init__(self):
        super().__init__()
        self.model = build_twin()
        self.loss_tracker = keras.metrics.Mean(name="loss")

    @property
    def metrics(self):
        return [self.loss_tracker]

    def train_step(self, batch: tf.Tensor) -> tf.Tensor:
        """train_step 方法。

        进行一个训练步骤。进行模型预测，计算损失，将损失传递给
        优化器，并使优化器应用梯度。

        参数：
            batch: 一批数据，用于损失函数。

        返回：
            返回一个包含损失指标的字典。
        """

        # 从批次中获取两个增强版本
        y_a, y_b = batch

        with tf.GradientTape() as tape:
            # 获取两个版本的预测
            z_a, z_b = self.model(y_a, training=True), self.model(y_b, training=True)
            loss = self.loss(z_a, z_b)

        grads_model = tape.gradient(loss, self.model.trainable_variables)

        self.optimizer.apply_gradients(zip(grads_model, self.model.trainable_variables))
        self.loss_tracker.update_state(loss)

        return {"loss": self.loss_tracker.result()}

# 设置模型、优化器、损失

bm = BarlowModel()
# 选择 LAMB 优化器是因为批大小较大。收敛得比 ADAM 或 SGD 快得多
optimizer = tfa.optimizers.LAMB()
loss = BarlowLoss(BATCH_SIZE)

bm.compile(optimizer=optimizer, loss=loss)

# 预计训练时间：1 小时 30 分钟

history = bm.fit(augment_versions, epochs=160)
plt.plot(history.history["loss"])
plt.show()

# 近似：这个barlow twins模型的准确率为64%。

xy_ds = (
    tf.data.Dataset.from_tensor_slices((train_features, train_labels))
    .shuffle(1000)
    .batch(BATCH_SIZE, drop_remainder=True)
    .prefetch(tf.data.AUTOTUNE)
)

test_ds = (
    tf.data.Dataset.from_tensor_slices((test_features, test_labels))
    .shuffle(1000)
    .batch(BATCH_SIZE, drop_remainder=True)
    .prefetch(tf.data.AUTOTUNE)
)

model = keras.models.Sequential(
    [
        bm.model,
        keras.layers.Dense(
            10, activation="softmax", kernel_regularizer=keras.regularizers.l2(0.02)
        ),
    ]
)

model.layers[0].trainable = False

linear_optimizer = tfa.optimizers.LAMB()
model.compile(
    optimizer=linear_optimizer,
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

model.fit(xy_ds, epochs=35, validation_data=test_ds)