imgaug数据增强神器：第三章 调用多核CPU

官网：https://imgaug.readthedocs.io/en/latest/
教程：https://nbviewer.jupyter.org/github/aleju/imgaug-doc/blob/master/notebooks/A03 - Multicore Augmentation.ipynb

一、在多个CPU核心上进行增强

增强可能是一个缓慢的过程，尤其是在处理大型图像和组合许多不同的增强技术时。 使用过时的硬件，查看性能文档，了解预期的单核性能的下限。

提高性能的一种方法是在多个CPU内核上同时增强。 imgaug提供了一个本地系统来做到这一点。 它基于大致以下步骤：

1. 将数据集拆分为批次。每批包含一个或多个图像以及与它们相关的附加数据，例如，边界框或分割图。 （拆分数据也可以 用生成器动态完成。）
2. 启动一个或多个子进程。它们中的每一个都在自己的CPU核心上运行。
3. 将批次发送到子进程。尝试在子进程上平均分配它们，以便它们中的每一个都有相似的工作量。
4. 让子进程增强数据。
5. 从子进程接收增强的批处理。

从这些步骤可以得出一些重要的观点。首先，数据必须分成批次。其次，将所有数据合并为一个批处理并使用多核增强是没有意义的，因为每个单独的批处理只增加了一个核心。第三，对少量数据使用多核增强也没有意义，因为启动子进程可能比简单地在单个CPU核心上增加数据集花费更多时间。 （因此您可以在多个迭代中重复使用子进程，以此获得回报。）

重要提示：imgaug提供多核功能，建议将它们用于多核增强。 建议不要在定制的多核例程中运行imgaug，如python的multiprocessing库或一些深度学习库的多进程支持。 这样做会产生很大的风险，意外地在每个子进程中应用相同的增强（仅适用于不同的图像）。 如果仍然决定构建自定义实现，请确保使用每个子进程的不同种子调用imgaug.seed(value)和augmenter.reseed(value)。 然后还建议为每个子进程生成调试输出。 弄清楚这很容易，但很难注意到错误的发生。

二、示例: augment_batches(…, background=True)

在imgaug中进行多核增强的最简单方法是调用augment_batches(..., background=True)。 它类似于augment_images()。

不同之处在于它需要一个imgaug.Batch实例列表。这些实例中的每一个都包含批次的数据，例如图像或边界框。创建批处理是很简单，可以通过batch = imgaug.Batch(images=, bounding_boxes=)创建。

与augment_images()的另一个区别是augment_batches()返回一个生成器，它从子进程接收数据的同时不断产生增强的批处理。

最后也是重要的差异是augment_batches()当前不使用增强器中设置的随机状态，而是选择一个新的。不这样做的话，所有子进程都会应用相同的增强（仅适用于不同的图像）。

如果你需要更多地控制随机状态使用pool()或imgaug.multicore.Pool（参见下面的内容）。

让我们尝试使用augment_batches()。首先，我们定义一些示例数据。

import numpy as np
import imgaug as ia
%matplotlib inline

BATCH_SIZE = 16
NB_BATCHES = 100

image = ia.quokka_square(size=(256, 256))
images = [np.copy(image) for _ in range(BATCH_SIZE)]

现在我们将图像组合到imgaug.Batch实例：

batches = [ia.Batch(images=images) for _ in range(NB_BATCHES)]

我们的增强序列包含PiecewiseAffine，它往往是一个非常缓慢的增强器。 在图像上使用更密集的点网格会加剧减慢速度。 每个这样的点将导致更多的局部仿射变换被应用。

from imgaug import augmenters as iaa

aug = iaa.Sequential([
    iaa.PiecewiseAffine(scale=0.05, nb_cols=6, nb_rows=6),  # very slow
    iaa.Fliplr(0.5),  # very fast
    iaa.CropAndPad(px=(-10, 10))  # very fast
])

现在我们增强批次。 让我们先在没有多核增强的情况下进行增强，看看单个CPU核心需要多长时间。 augment_batches()返回imgaug.Batch实例的生成器。 然后，我们可以通过属性imgaug.Batch.images_aug访问增强图像。

import time

time_start = time.time()
batches_aug = list(aug.augment_batches(batches, background=False))  # list() converts generator to list
time_end = time.time()

print("Augmentation done in %.2fs" % (time_end - time_start,))
ia.imshow(batches_aug[0].images_aug[0])

Augmentation done in 134.91s

100批大约130秒，每批包含16个大小为256x256的图像。 每张图片约为0.08秒。 不是很快，GPU很可能比这更快地训练。 让我们尝试使用多核增强。

time_start = time.time()
batches_aug = list(aug.augment_batches(batches, background=True))  # background=True for multicore aug
time_end = time.time()

print("Augmentation done in %.2fs" % (time_end - time_start,))
ia.imshow(batches_aug[0].images_aug[0])

Augmentation done in 28.07s

减少到不到30秒——大约是单核时间的四分之一。 已经好多了。请注意，这是一个过时的CPU，有4个内核和8个线程。 现代的8核CPU应该会快很多。

三、增强批量非图像数据

上面的例子只展示了如何增强图像。 通常情况下，你还会想要增强关键点或边界框。 通过在创建imgaug.Batch对象时进行简单的更改便可实现。 在这种情况下，您不必担心随机状态或随机/确定性模式。 imgaug将自动处理并确保图像和相关数据之间的扩充对齐。

让我们增强之前的示例数据中的关键点。

BATCH_SIZE = 16
NB_BATCHES = 100
image = ia.quokka(size=0.2)
images = [np.copy(image) for _ in range(BATCH_SIZE)]
keypoints = ia.quokka_keypoints(size=0.2)
keypoints = [keypoints.deepcopy() for _ in range(BATCH_SIZE)]

batches = [ia.Batch(images=images, keypoints=keypoints) for _ in range(NB_BATCHES)]

现在以与以前相同的方式增加数据：

time_start = time.time()
batches_aug = list(aug.augment_batches(batches, background=True))  # background=True for multicore aug
time_end = time.time()

print("Augmentation done in %.2fs" % (time_end - time_start,))
ia.imshow(
    batches_aug[0].keypoints_aug[0].draw_on_image(
        batches_aug[0].images_aug[0]
    )
)

Augmentation done in 83.81s

就是这样。 只需在实例化imgaug.Batch()实例时添加keypoints=，其余部分由库处理。 对于边界框（bounding_boxes=），热图（heatmaps=）或分割图（segmentation_maps=）也可以这样做。 只需确保列表具有相同的长度，并且具有相同索引的条目实际上彼此属于（例如图像0014059.jpg和图像0014059.jpg的关键点）。

您可能已经注意到，这里的增强时间从大约30秒增加到大约80秒——仅仅添加了关键点。 这是因为当将关键点转换为坐标时，由于使用不当，PiecewiseAffine基于图像的方法进行关键点增强。 它是目前库中最慢的关键点增强器（因此在扩充关键点或边界框时避免使用PiecewiseAffine）。

四、使用池

augment_batches()易于使用，但是没有足够的自定义。 如果你想要更多自定义，例如 控制使用的CPU核心数或随机数种子，augmenter.pool()是一个简单的替代方案（它是augment_batches()使用的后端）。

这次使用pool()再次增加了先前定义的批次。 我们将池配置为使用除1之外的所有CPU核心（processes=-1），在20个任务之后重启子进程（maxtasksperchild=20）并以1为随机数种子开始。如果你处理内存泄漏的话，参数maxtasksperchild可能很有用，因为随着时间的推移会消耗越来越多的内存。 如果您没有这个问题，则没有理由使用该参数（使用它会消耗性能）。

with aug.pool(processes=-1, maxtasksperchild=20, seed=1) as pool:
    batches_aug = pool.map_batches(batches)
ia.imshow(batches_aug[0].images_aug[0])

五、使用池生成器

前两个示例显示了如何使用imgaug池中的列表。 对于大型数据集，为避免将整个数据集存储在内存中，使用生成器可能更合适。 通过用imap_batches()替换map_batches()可以轻而易举地完成这项工作。 该函数的输出也是一个生成器。

def create_generator(lst):
    for list_entry in lst:
        yield list_entry

my_generator = create_generator(batches)

with aug.pool(processes=-1, maxtasksperchild=20, seed=1) as pool:
    batches_aug = pool.imap_batches(my_generator)

    for i, batch_aug in enumerate(batches_aug):
        if i == 0:
            ia.imshow(batch_aug.images_aug[0])
        # do something else with the batch here

六、总结

使用imgaug进行多核增强，请执行以下操作：

• 将您的数据转换为imgaug.Batch的实例。 确保相应的数据在批次中具有相同的列表索引，如图像及其相应的关键点。
• 调用augmenter.augment_batches(batches, background=True)。 这将返回一个生成器。
• 如果您需要更多控制或想要使用生成器作为输入，请使用augmenter.pool([processes], [maxtasksperchild], [seed])。 在池上调用pool.map_batches(list)或pool.imap_batches(generator)。
• 避免实现自己的多核系统或使用其他库，因为它很容易搞砸。

下一章：imgaug数据增强神器：第四章 增强关键点/界标