► 代码示例 / 计算机视觉 / RandAugment 用于提高图像分类模型的鲁棒性

RandAugment 用于提高图像分类模型的鲁棒性

作者： Sayak Paul Sachin Prasad
创建日期 2021/03/13
最后修改 2023/12/12
描述： 使用 RandAugment 训练图像分类模型以提高鲁棒性。

ⓘ 此示例使用 Keras 3

数据增强是一种非常有用的技术，可以帮助提高卷积神经网络（CNN）的平移不变性。RandAugment 是一种用于视觉数据的随机数据增强例程，最早在RandAugment: 减少搜索空间的实用自动化数据增强中提出。它由颜色抖动、高斯模糊、饱和度等强力增强变换以及随机裁剪等更传统的增强变换组成。

这些参数针对给定的数据集和网络架构进行调整。RandAugment 的作者还在原始论文（图 2）中提供了 RandAugment 的伪代码。

最近，它已成为诸如Noisy Student Training和用于一致性训练的无监督数据增强等工作的关键组成部分。它也对EfficientNets的成功至关重要。

pip install keras-cv

导入和设置

import os

os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import keras_cv
from keras import ops
from keras import layers
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import tensorflow_datasets as tfds

tfds.disable_progress_bar()
keras.utils.set_random_seed(42)

加载 CIFAR10 数据集

在本例中，我们将使用CIFAR10 数据集。

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
print(f"Total training examples: {len(x_train)}")
print(f"Total test examples: {len(x_test)}")

Total training examples: 50000
Total test examples: 10000

定义超参数

AUTO = tf.data.AUTOTUNE
BATCH_SIZE = 128
EPOCHS = 1
IMAGE_SIZE = 72

初始化 `RandAugment` 对象

现在，我们将使用 RandAugment 作者建议的参数，从 imgaug.augmenters 模块中初始化一个 RandAugment 对象。

rand_augment = keras_cv.layers.RandAugment(
    value_range=(0, 255), augmentations_per_image=3, magnitude=0.8
)

创建 TensorFlow `Dataset` 对象

train_ds_rand = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(BATCH_SIZE * 100)
    .batch(BATCH_SIZE)
    .map(
        lambda x, y: (tf.image.resize(x, (IMAGE_SIZE, IMAGE_SIZE)), y),
        num_parallel_calls=AUTO,
    )
    .map(
        lambda x, y: (rand_augment(tf.cast(x, tf.uint8)), y),
        num_parallel_calls=AUTO,
    )
    .prefetch(AUTO)
)

test_ds = (
    tf.data.Dataset.from_tensor_slices((x_test, y_test))
    .batch(BATCH_SIZE)
    .map(
        lambda x, y: (tf.image.resize(x, (IMAGE_SIZE, IMAGE_SIZE)), y),
        num_parallel_calls=AUTO,
    )
    .prefetch(AUTO)
)

为了比较，我们还将定义一个简单的增强流程，包括随机翻转、随机旋转和随机缩放。

simple_aug = keras.Sequential(
    [
        layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(factor=0.02),
        layers.RandomZoom(height_factor=0.2, width_factor=0.2),
    ]
)

# Now, map the augmentation pipeline to our training dataset
train_ds_simple = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(BATCH_SIZE * 100)
    .batch(BATCH_SIZE)
    .map(lambda x, y: (simple_aug(x), y), num_parallel_calls=AUTO)
    .prefetch(AUTO)
)

可视化使用 RandAugment 增强后的数据集

sample_images, _ = next(iter(train_ds_rand))
plt.figure(figsize=(10, 10))
for i, image in enumerate(sample_images[:9]):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(image.numpy().astype("int"))
    plt.axis("off")

png

建议你多次运行上面的代码块，以查看不同的变化。

可视化使用 `simple_aug` 增强后的数据集

sample_images, _ = next(iter(train_ds_simple))
plt.figure(figsize=(10, 10))
for i, image in enumerate(sample_images[:9]):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(image.numpy().astype("int"))
    plt.axis("off")

png

定义模型构建辅助函数

现在，我们定义一个基于ResNet50V2 架构的 CNN 模型。另外，请注意，该网络内部已经包含一个缩放层。这使得无需对数据集进行任何单独的预处理，这对于部署目的特别有用。

def get_training_model():
    resnet50_v2 = keras.applications.ResNet50V2(
        weights=None,
        include_top=True,
        input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3),
        classes=10,
    )
    model = keras.Sequential(
        [
            layers.Input((IMAGE_SIZE, IMAGE_SIZE, 3)),
            layers.Rescaling(scale=1.0 / 127.5, offset=-1),
            resnet50_v2,
        ]
    )
    return model


get_training_model().summary()

Model: "sequential_1"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape              ┃    Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ rescaling (Rescaling)           │ (None, 72, 72, 3)         │          0 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ resnet50v2 (Functional)         │ (None, 10)                │ 23,585,290 │
└─────────────────────────────────┴───────────────────────────┴────────────┘

 Total params: 23,585,290 (89.97 MB)

 Trainable params: 23,539,850 (89.80 MB)

 Non-trainable params: 45,440 (177.50 KB)

我们将使用两个不同版本的数据集来训练该网络

一个使用 RandAugment 进行增强。
另一个使用 simple_aug 进行增强。

由于 RandAugment 以增强模型对常见干扰和损坏的鲁棒性而闻名，我们还将根据 Hendrycks 等人在评估神经网络对常见损坏和干扰的鲁棒性中提出的 CIFAR-10-C 数据集来评估我们的模型。CIFAR-10-C 数据集包含 19 种不同的图像损坏和干扰（例如散斑噪声、雾、高斯模糊等），并且严重程度各不相同。在本例中，我们将使用以下配置：cifar10_corrupted/saturate_5。此配置中的图像如下所示

为了确保可复现性，我们序列化了浅层网络的初始随机权重。

initial_model = get_training_model()
initial_model.save_weights("initial.weights.h5")

使用 RandAugment 训练模型

rand_aug_model = get_training_model()
rand_aug_model.load_weights("initial.weights.h5")
rand_aug_model.compile(
    loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
)
rand_aug_model.fit(train_ds_rand, validation_data=test_ds, epochs=EPOCHS)
_, test_acc = rand_aug_model.evaluate(test_ds)
print("Test accuracy: {:.2f}%".format(test_acc * 100))

 391/391 ━━━━━━━━━━━━━━━━━━━━ 1146s 3s/step - accuracy: 0.1677 - loss: 2.3232 - val_accuracy: 0.2818 - val_loss: 1.9966
 79/79 ━━━━━━━━━━━━━━━━━━━━ 39s 489ms/step - accuracy: 0.2803 - loss: 2.0073
Test accuracy: 28.18%

使用 `simple_aug` 训练模型

simple_aug_model = get_training_model()
simple_aug_model.load_weights("initial.weights.h5")
simple_aug_model.compile(
    loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
)
simple_aug_model.fit(train_ds_simple, validation_data=test_ds, epochs=EPOCHS)
_, test_acc = simple_aug_model.evaluate(test_ds)
print("Test accuracy: {:.2f}%".format(test_acc * 100))

 391/391 ━━━━━━━━━━━━━━━━━━━━ 1132s 3s/step - accuracy: 0.3673 - loss: 1.7929 - val_accuracy: 0.4789 - val_loss: 1.4296
 79/79 ━━━━━━━━━━━━━━━━━━━━ 39s 494ms/step - accuracy: 0.4762 - loss: 1.4368
Test accuracy: 47.89%

加载 CIFAR-10-C 数据集并评估性能

# Load and prepare the CIFAR-10-C dataset
# (If it's not already downloaded, it takes ~10 minutes of time to download)
cifar_10_c = tfds.load("cifar10_corrupted/saturate_5", split="test", as_supervised=True)
cifar_10_c = cifar_10_c.batch(BATCH_SIZE).map(
    lambda x, y: (tf.image.resize(x, (IMAGE_SIZE, IMAGE_SIZE)), y),
    num_parallel_calls=AUTO,
)

# Evaluate `rand_aug_model`
_, test_acc = rand_aug_model.evaluate(cifar_10_c, verbose=0)
print(
    "Accuracy with RandAugment on CIFAR-10-C (saturate_5): {:.2f}%".format(
        test_acc * 100
    )
)

# Evaluate `simple_aug_model`
_, test_acc = simple_aug_model.evaluate(cifar_10_c, verbose=0)
print(
    "Accuracy with simple_aug on CIFAR-10-C (saturate_5): {:.2f}%".format(
        test_acc * 100
    )
)

 Downloading and preparing dataset 2.72 GiB (download: 2.72 GiB, generated: Unknown size, total: 2.72 GiB) to /home/sachinprasad/tensorflow_datasets/cifar10_corrupted/saturate_5/1.0.0...
 Dataset cifar10_corrupted downloaded and prepared to /home/sachinprasad/tensorflow_datasets/cifar10_corrupted/saturate_5/1.0.0. Subsequent calls will reuse this data.
Accuracy with RandAugment on CIFAR-10-C (saturate_5): 30.36%
Accuracy with simple_aug on CIFAR-10-C (saturate_5): 37.18%

为了本例的目的，我们只训练了模型一个 epoch。在 CIFAR-10-C 数据集上，使用 RandAugment 的模型可以获得更好的性能和更高的准确率（例如，在一项实验中达到 76.64%），而使用 simple_aug 训练的模型准确率较低（例如，64.80%）。RandAugment 还可以帮助稳定训练。

在 notebook 中，你可能会注意到，虽然使用 RandAugment 增加了训练时间，但我们在 CIFAR-10-C 数据集上获得了更好的性能。你可以在同一 CIFAR-10-C 数据集提供的其他损坏和干扰设置下进行实验，看看 RandAugment 是否有帮助。

你还可以尝试使用 RandAugment 对象中 n 和 m 的不同值。在原始论文中，作者展示了针对特定任务的单个增强变换的影响以及一系列消融研究。欢迎查阅。

RandAugment 在提高计算机视觉深度模型的鲁棒性方面取得了显著进展，如Noisy Student Training和FixMatch等工作所示。这使得 RandAugment 成为训练不同视觉模型的非常有用的方法。

你可以使用托管在Hugging Face Hub上的训练好的模型，并在Hugging Face Spaces上试用 demo。

RandAugment 用于提高图像分类模型的鲁棒性

导入和设置

加载 CIFAR10 数据集

定义超参数

初始化 RandAugment 对象

创建 TensorFlow Dataset 对象

可视化使用 RandAugment 增强后的数据集

可视化使用 simple_aug 增强后的数据集