【图像增广库imgaug】官方文档翻译（二）：在多个 CPU 内核上进行扩充_imgaug.multicore.pool

---

图片增广可能是一个缓慢的过程，尤其是在处理大图像和组合许多不同的增强技术时。可以查看文档了解使用过时硬件的预期单核性能的一些下限。
提高性能的一种方法是同时在多个 CPU 内核上进行增强。 提供了一个本地系统来做到这一点。它大致基于以下步骤：

1.将数据集拆分为batches。每个batch包含一个或多个图像以及与之关联的附加数据，例如边界框或分割图。（也可以使用生成器动态拆分数据。）
2.启动一个或多个子进程。它们中的每一个都在自己的CPU内核上运行。
3.将批处理发送到子进程。尝试将它们平均分布在子进程上，以便每个进程都有相似的工作量要做。
4.让子进程扩充数据。
5.从子进程接收扩充的batch。
从这些步骤中可以得出一些要点。首先，必须将数据拆分为batches。其次，将所有数据合并到一个批处理中使用多核增强是毫无意义的，因为每个单独的批处理都只由一个内核增强。第三，对少量数据使用多核增强也可能毫无意义，因为启动子进程可能比简单地在单个 CPU 内核上扩充数据集花费更多时间。（不过，您可以在epoch之间重用子进程，这样可能仍然有效。

重要提示：提供多核功能，建议将其用于多核增强。不建议使用例如python的库或某些深度学习库的多处理支持在定制的多核例程中执行。这样做会带来一个很大的风险，即意外地在每个子worker身上应用相同的图像扩充（只是对不同的图像）。如果仍然决定构建自定义实现，请确保调用每个子进程并使用不同的seed。然后最好为每个子进程生成调试输出。因为很容易搞错，甚至很难注意到。

---

一、示例：augment_batches(…, background=True)

执行多核增强的最简单方法是调用augment_batches(…, background=True) 。它的工作原理类似于示例augment_images()。不同之处在于，它需要一个imgaug.augmentables.batches.Batch列表或者imgaug.augmentables.batches.UnnormalizedBatch实例。这些实例中的每一个都包含一个batch处理的数据，例如图像或边界框。创建批处理是很简单的，可以通过batch = UnnormalizedBatch(images=< list of numpy arrays>, bounding_boxes=< list of imgaug.BoundingBoxOnImages>)完成。augment_images()另一个区别是augment_batches()返回一个生成器，该生成器在从子进程接收到增强批次时连续产生这些批次。最后（也是重要的）区别是，augment_batches()不使用增强器中设置的随机状态，而是选择一个新的状态。这是因为否则所有子进程都将应用相同的增强（仅应用于不同的图像）。如果您需要对随机状态使用进行更多控制，或者相反（请参阅下面的进一步内容）。

让我们试着使用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)]
1
2
3
4
5
6
7
8
9

现在，我们将图像合并到实例中：UnnormalizedBatch

from imgaug.augmentables.batches import UnnormalizedBatch
batches = [UnnormalizedBatch(images=images) for _ in range(NB_BATCHES)]
1
2

由于这里的数据已经很好地规范化，我们也可以使用具有相同接口的imgaug.augmentables.batches.Batch代替。例如，如果我们有keypoints，并且这些 keypoints 以xy元组列表的形式提供，UnnormalizedBatch就必须被使用，这会将这些列表转换为KeypointsOnImage实例。

我们的图像增强sequence包含PiecewiseAffine，往往是一个缓慢的增强器。我们通过在图像上使用更密集的点网格会进一步减慢速度。每个这样的点都将导致应用更多的局部仿射变换。

from imgaug import augmenters as iaa

aug = iaa.Sequential([
    iaa.PiecewiseAffine(scale=0.05, nb_cols=6, nb_rows=6),  # 很慢
    iaa.Fliplr(0.5),  # 很快
    iaa.CropAndPad(px=(-10, 10))  # 很快
])
1
2
3
4
5
6
7

现在，我们扩充生成的batch。让我们首先在不进行多核增强的情况下进行增强，以查看单个 CPU 核心需要多长时间。 augment_batches()返回一个Batch实例生成器。然后，我们可以通过属性UnnormalizedBatch.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])
1
2
3
4
5
6
7
8

Augmentation done in 151.40s

大约 150 秒，100 个batches，每个batch包含 16 张大小为 256x256 的图像。这大约是每张图像0.09秒。不是很快，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])
1
2
3
4
5
6

Augmentation done in 81.32s

低至 80 秒，这大约是单个内核时间的一半。这已经好多了。请注意，这是在具有4个内核和8个线程的过时CPU上。现代 8 核 CPU 应该会表现好很多。

---

二、包含非图像数据的Batches

上面的示例仅显示了如何增强图像。通常，您还需要增加例如这些关键点或边界框。这是通过创建UnnormalizedBatch实例时的小改变来实现的。在这种情况下，您不必担心随机状态或随机/确定性模式。 imgaug将自动处理该操作，并确保图像和关联数据之间的增强对齐。

让我们用一些keypoints扩展前面的示例数据。

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 = [UnnormalizedBatch(images=images, keypoints=keypoints) for _ in range(NB_BATCHES)]
1
2
3
4
5
6
7
8

以先前相同的方式增强数据：

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]
    )
)
1
2
3
4
5
6
7
8
9
10

Augmentation done in 138.25s

就是这样。只需在实例化UnnormalizedBatch实例时添加keypoints=< list of imgaug.augmentables.kps.KeypointsOnImage>，其余部分由库处理。对于边界框 （bounding_boxes=< list of imgaug.augmentables.bbs.BoundingBoxesOnImage>）、线串 （bounding_boxes=< list of imgaug.augmentables.lines.LineStringsOnImage）、多边形 （polygons=< list of imgaug.augmentables.polygons.PolygonsOnImage>）、热图 （heatmaps=< list of imgaug.augmentables.heatmaps.HeatmapsOnImage>） 或分割图 （segmentation_maps=< list of imgaug.augmentables.segmaps.SegmentationMapOnImage>），也可以执行相同的操作。只需确保列表具有相同的长度，并且具有相同索引的条目一一对应（例如图像0014059.jpg和对应keypoints）。请注意，所有这些输入也可以以非规范化形式提供，例如一组(N,2) float32格式数组的关键点keypoints列表，一组(N,2) float32格式数组的多边形polygons列表，又或是一组(H,W,C) int 32格式数组的分割图segmentation_maps列表。

您可能已经注意到，这里的增强时间从大约80秒增加到大约140秒 - 只是通过添加关键点keypoints。这是因为PiecewiseAffine使用基于图像的方法进行关键点增强，因为在将关键点转换为坐标时不准确。它目前是库中速度最慢的关键点增强器（因此，在扩充基于坐标的数据（如关键点或边界框）时，请避免使用PiecewiseAffine）。

---

三、使用池

augment_batches()用起来很简单，但不提供太多可自定义的设置。例如，如果您想控制已用CPU内核的数量或随机数种子，augmenter.pool()是一个简单的替代方案（它是augment_batches()用的后端）。下面的示例再次扩充了以前定义的批次，这次是用pool()。我们将池配置为使用除一个 CPU 内核（processes=-1）之外的所有 CPU 内核，在 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])
1
2
3

请注意，我们只在此处调用了map_batches()以增加输入批处理。在实际操作中，我们可以使用不同的输入批处理为每个生成的池多次调用该命令 - 也是建议这样做，因为创建新池需要重生子进程，这确实需要花费一些时间。

augmenter.pool()是一个快捷方式，它创建了 一个imgaug.multicore.Pool实例，它又是围绕 python 的multiprocessing.Pool的包装器。包装器主要处理子进程之间随机状态的正确管理。下面的示例显示了imgaug.multicore.Pool的用法，使用与上面示例中相同的种子，从而生成相同的输出。

from imgaug import multicore

with multicore.Pool(aug, processes=-1, maxtasksperchild=20, seed=1) as pool:
    batches_aug = pool.map_batches(batches)

ia.imshow(batches_aug[0].images_aug[0])
1
2
3
4
5
6

---

四、用生成器使用池

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

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

my_generator = create_generator(batches)

with aug.pool(processes=-1, 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
1
2
3
4
5
6
7
8
9
10
11
12
13

注意，如果您不需要以您提供的相同顺序返回batch，您也可以使用imap_batches_unordered()而不是imap_batches()。无序的方法速度通常更快。

---

五、限速池可降低最大内存要求

（0.3.0 版中的新功能）

默认情况下，池将贪婪地从生成器加载(并增强)尽可能多的批。没有限制最多允许多少扩增批次在管道中“等待”的速率限制。这意味着在最坏的情况下(当模型训练非常慢，而图像增强非常快)，整个增强数据集可能正在等待检索，这可能会导致高内存占用，具体取决于数据集有多大。

要解决这个问题，可以使用参数output_buffer_size。该值控制在整个扩增管道中最多允许存在多少批，即imap_batches(gen)将从gen加载新的批，直到达到output_buffer_size批，然后仅在成功生成增强批次时才加载另一个批次。

下面的代码显示了一个示例。它类似于上面的那个，但使用扩充管道，该管道生成批次的速度快于使用批次的速度。打印的消息将准确显示何时加载批处理以及何时从扩充管道请求批处理。为了限制此快速管道的内存要求，将允许的等待批次数限制为 5 个。请注意，imgaug中的批次包含增强之前和之后的图像，因此这里的有效内存要求为52I，其中I是单个图像的size。在实际操作中，该值应被视为实际内存需求的下限，例如，将数据复制到后台进程可能会暂时使需求增加一倍。

import time

# 我们在这里用一个快速的增强器来显示批次
# 只有在缓冲区中再次有空间时才加载。
pipeline = iaa.Fliplr(0.5)

def create_generator(lst):
    for list_entry in lst:
        print("Loading next unaugmented batch...")
        yield list_entry

# 这里只用25批就足够看出效果了
my_generator = create_generator(batches[0:25])

with pipeline.pool(processes=-1, seed=1) as pool:
    batches_aug = pool.imap_batches(my_generator, output_buffer_size=5)

    print("Requesting next augmented batch...")
    for i, batch_aug in enumerate(batches_aug):
        # 在这里sleep一段时间来模拟一个缓慢的训练模型
        time.sleep(0.1)

        if i < len(batches)-1:
            print("Requesting next augmented batch...")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

Requesting next augmented batch…
Loading next unaugmented batch…
Loading next unaugmented batch…
Loading next unaugmented batch…
Loading next unaugmented batch…
Loading next unaugmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Loading next unaugmented batch…
Requesting next augmented batch…
Requesting next augmented batch…
Requesting next augmented batch…
Requesting next augmented batch…
Requesting next augmented batch…
Requesting next augmented batch…

方法imap_batches_unordered()还支持 output_buffer_size。但是，map_batches()不支持该参数，并且总是扩充整个输入列表。

---

总结

因此，要使用imgaug多核增强，只需执行以下操作：

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