Image Processing 必備(三)：Imgaug 之調用多核CPU

Date： 2020-09-02

官方教程：https://nbviewer.jupyter.org/github/aleju/imgaug-doc/blob/master/notebooks/A03%20-%20Multicore%20Augmentation.ipynb

多核處理概述

圖像增強通常是一個緩慢的過程，特別是當在大量的圖像數據下進行多種不同的增強技術時。

提升性能的一種方式就是在多個CPU內核上同時進行數據增強。imgaug中提供了一個實現多線程的原生系統，大致遵循下述步驟：

1. 將數據集拆分為批次。每批包含一個或多個圖像以及與它們相關的附加數據，如，邊界框或分割圖。

2. 啟動一個或多個子進程。它們中的每一個都在自己的CPU核心上運行。

3. 將批次發送到子進程。嘗試在子進程上平均分配它們，以便它們中的每一個都有相似的工作量。

4. 讓子進程增強數據；

5. 從子進程接收增強的批處理。

Important： imgaug中提供了多核功能，並且也建議使用提供的功能進行多核處理。不建議使用python的multiprocess庫的多進程處理，可能會產生不可預估的錯誤。

示例：augment_batches(..., background=True)

在imgaug中進行多核增強最簡單的方式就是調用augment_batches(..., background=True).其工作原理相似於augment_images()。區別在於以下幾點：

1. 其需要一個imgaug.augmentables.batches.Batch或者imgaug.augmentables.batches.UnnormalizedBatch實例構成的列表。這些實例中均包含批次的數據，例如圖像或者是bbox信息。而實例中包含的批次batch可以通過batch = UnnormalizedBatch(images=<list of numpy arrays>, bounding_boxes=<list of imgaug.BoundingBoxOnImages>)生成；

2. 另一個不同在於augment_batches()返回的是一個生成器，當從子進程接收到時會不斷產生批處理；

3. 最后一各卻別在於augment_batches()當前不使用增強器中設置的隨機狀態，而是選擇一個新的。不這樣做的話，所有子進程都會應用相同的增強。

---

Tip: 可還記得augment_images()是干什么用的？

答：augment_images()是通過aug.augment_images(images=images)進行使用的。也就是說，是使用某個圖像增強方法的實例對圖像進行圖像增強時使用!

---

使用單核處理

1. 首先，生成一些實例數據：

import numpy as np
import imgaug as ia
import imageio
from imgaug import augmenters as iaa 
%matplotlib inline
​
BATCH_SIZE = 16
NB_BATCHES = 100  # number of batch
​
image = imageio.imread("./pick1.jpg")
images = [np.copy(image) for _ in range(BATCH_SIZE)]

　　2.使用UnnormalizedBatch將圖像生成批次batch：

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

　　3.定義增強序列

這里選用處理時長較長的PiecewiseAffine,在圖像上使用更密集的點網格會加劇減慢速度。 每個這樣的點將導致更多的局部仿射變換被應用。

seq = 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
])

從上述代碼中可以看出，增強序列中選擇的3個增強方法具有不同的速度。

4.對生成的batches進行增強：

首先我們不使用多核增強的方式，看看使用單個CPU核進行處理需要多長時間。

augment_batches()返回一個Batch示例的生成器，然后我們可以通過UnnormalizedBatch.images_aug訪問增強后的圖像

import time
​
time_start = time.time()
batches_aug = list(seq.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 398.36s

通過上述運行結果的反饋，對於100個batch每個batch16張圖片(500*313)進行數據增強，居然用了398.36s。

太慢了，官方教程上說，使用GPU會再快一些。

使用多核處理

從下述代碼中可以看出， 使用多核處理和單核處理的區別  在於augment_batches(..., background=)中background參數為True還是False。

time_start = time.time()
batches_aug = list(seq.augment_batches(batches, background=True))  # 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 214.00s

使用多核處理，處理時間從398.36s降低到214.00s，節省用時184.36s。

如下圖所示，本節教程使用的CUP具有2核，4個線程。對於更高核數的CPU和線程數，多核處理會有更好的效果。

---

Tip：如何查看自己電腦的核數和線程數？

答：在cmd命令中輸入"wmic",然后在生成的新一行中輸入"cpu get *"。

• NumberOfCores: 表示CPU核心數；

• NumberOfLogicalProcessors:表示CPU線程數。

---

處理非圖像數據

之前的示例只展示了如何增強圖像，通常還需要增強關鍵點，bbox等非圖像的數據，以保持增強前后的一致性。

對非數據類型進行多核處理與圖像的區別在於使用UnnormalizedBatch生成batch實例時的一些小的不同。並且在對非圖像數據進行多核增強處理時，不必再像之前那樣考慮隨機/確定模式的情況。imgaug會自動處理，並確保圖像和相關數據之間的增強一致。

關鍵點示例

生成關鍵點

import numpy as np
import imgaug as ia
import imageio
from imgaug import augmenters as iaa 
%matplotlib inline
​
BATCH_SIZE = 16
NB_BATCHES = 100
image = imageio.imread("./pick1.jpg")
images = [np.copy(image) for _ in range(BATCH_SIZE)]
​
keypoint = ia.KeypointsOnImage([
                    ia.Keypoint(x=151, y=106),
                    ia.Keypoint(x=247, y=78),
                    ia.Keypoint(x=179, y=140),
                    ia.Keypoint(x=206, y=134)
                    ], shape=image.shape)  # KeyPointsOnImage class  
keypoints = [keypoint.deepcopy() for _ in range(BATCH_SIZE)] # has deepcopy func
​
seq = 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
])

關鍵點增強處理

之前在UnnormalizedBatch()中只添加了images，此時由於具有keypoints，也同時添加了keypoint字段。

多核的利用分為兩步走， 其一，對自己生成的批次圖像、真值數據使用UnnormalizedBatch()生成可使用的batches數據， 其二，用數據增強方法中包含的augment_batches()方法（background要設置為True）

from imgaug.augmentables.batches import UnnormalizedBatch
import time
​
batches = [UnnormalizedBatch(images=images, keypoints=keypoints) for _ in range(NB_BATCHES)]
​
time_start = time.time()
batched_aug = list(seq.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(
    batched_aug[0].keypoints_aug[0].draw_on_image(
        batched_aug[0].images_aug[0]
    )
)

Augmentation done in 370.16s

從結果可以看出，僅僅增加了keypoint，運行時間便從214.00s增加至370.16s，增加了155.84s，近一倍的時間。

同理，UnnormalizedBatch()函數同樣可以使用

• bbox，即bounding_boxes=[list of imgaug.augmentables.bbs.BoundingBoxesOnImage]

• polygons(多邊形)，即polygons=[list of imgaug.augmentables.polygons.PolygonsOnImage]

• heatmaps(熱力圖)，即heatmaps=[list of imgaug.augmentables.heatmaps.HeatmapsOnImage]

• segmentation maps(分割圖)， 即segmentation_maps=[list of imgaug.augmentables.segmaps.SegmentationMapOnImage] 等圖像中使用的真值類型。

關於augment_batches()的使用感覺還比較簡單，但是有些無腦。例如，關於控制使用CPU的核數仿佛並沒有自定義設置的方式。感覺一切都交給它，私人設置的少，缺乏可控性和安全感。

---

Tip: 基於坐標的數據(如關鍵點或邊界框)時，避免使用PiecewiseAffine,時間代價會增加很多。

---

Pool

通過前兩個案例的使用來看，augment_batches()雖然使用簡單，但是沒有提供更多定制化的設置，比如控制使用CPU的核心數量。

augmenter.pool()可以輕松的替代augment_batches()的使用，並且具有以下兩種優點：

• 增加可控性，可以具體設置CPU的使用；

• 可以使用生成器避免大數據量一次性讀入造成內存不足的問題。

pool()的使用也是使用之前定義的batches。 pool()方法是通過增強方法或者增強序列產生的。

(一) 增加可控性

將pool配置為使用除一個(processes=-1)之外的所有CPU核心，在執行20個任務后重啟子進程(maxtasksperchild=20)，並以隨機數種子1開始。

如果處理導致越來越多的內存泄漏問題，參數maxtasksperchild可能會很有用。

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

如果電腦上安裝有任意殺毒軟件，在內存使用懸浮窗上可以看到兩者在內存使用的區別。

注意，這里只調用了一次map_batch()來增加輸入批。在實踐中，可以使用不同的輸入批對每個生成的pool多次調用該命令——建議這樣做，因為創建新pool需要重新生成子進程，這會花費一些時間。

augmeter.pool()是創建imgaug.multicore.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])

(二) 數據批次讀入內存

在處理圖像數據時，往往會遇到一個問題，即圖像數據量較大，難以一次性將圖像數據讀入內存。

在深度學習各種框架(Tensorflow,Pytorch)都有關於圖像分批讀入的方式。圖像增強imgaug同樣需要這樣的操作。

在imgaug的Pool中，使用方式也相對容易，就是將map_batches([list])替換為imap_batches([generator])。此時的輸出batches_aug將會是一個生成器。

1. batches to generator

由於imap_batches()的輸入是生成器，則首先需要將batches構造成一個generator。這里會使用到yield。

def create_generator(lst):
    for list_entry in lst:
        yield list_entry
​
my_generator = create_generator(batches)

　　2.使用pool中的imap_batches()接收生成器

以下代碼由於涉及到多進程的處理，如果在pycharm中編寫，需要將下屬代碼放到"if name == 'main':"之后。否則會報

The "freeze_support()" line can be omitted if the program is not going to be frozen to produce an executable.

with seq.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

通過上述的示例，個人認為pool中包含的map_batches(<list>)和imap_batches(<generator>)實際的作用是將seq(增強方法/序列)作用於圖像，建立其圖像-進程-增強操作之間的橋梁。

完整示例

import imageio
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.batches import UnnormalizedBatch
import numpy as np
import time
​
BATCH_SIZE = 10
BATCH_NB = 8
img = imageio.imread("./pick1.jpg")
images = [np.copy(img) for _ in range(BATCH_SIZE)]
​
aug = iaa.Affine(rotate=(-20, 20))
​
batches = [UnnormalizedBatch(images=images) for _ in range(BATCH_NB)]  # 形成batch
​
​
def create_generator(batches):
    for i in batches:
        yield i
​
​
gen = create_generator(batches)
​
# if __name__ == '__main__':   # pycharm中需要添加！
​
with aug.pool(processes=-1, seed=1) as pool:
    batches_aug = pool.imap_batches(gen)
​
    for i, batch in enumerate(batches_aug):
        if i == 0:
            ia.imshow(batch.images_aug[0])

圖像增強與模型訓練步調不一致

在上一小節介紹的pool.imap_batches()在解決數據同時讀入內存的問題。看起來這樣處理之后已經近乎完美，但實際上仍然存在問題。

這個問題體現在圖像增強后會使用模型進行訓練，由於數據增強處理速度相對較快，模型訓練的速度相對較慢，就會存在二者步調不一致的問題。 同時，pool總是貪婪的從生成器中導入batch，步調不一致的問題會造成pool將數據搶先都讀入到RAM中，造成RAM資源緊張的問題。

解決這一問題的方法：在imap_batches()中添加output_buffer_size 參數，限制pool從gen中最多導入batch的數量。也就是說，會讓pipeline中最多存在output_buffer_size個batch，不再是沒上限的讀取。

import time
​
# We use a single, very fast augmenter here to show that batches
# are only loaded once there is space again in the buffer.
pipeline = iaa.Fliplr(0.5)
​
def create_generator(lst):
    for list_entry in lst:
        print("Loading next unaugmented batch...")
        yield list_entry
​
# only use 25 batches here, which is enough to show the effect
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 here for a while to simulate a slowly training model
        time.sleep(0.1)
​
        if i < len(batches)-1:
            print("Requesting next augmented batch...")

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...
Requesting next augmented batch...
Requesting next augmented batch...
Requesting next augmented batch...
Requesting next augmented batch...

整理總結

本節主要介紹了不可控多核augment_batches()和可控多核pool()的使用。二者都是基於UnnormalizedBatch()方法生成的batches(該函數可以通過參數將image，各種真值組成batch)。

• augment_batches()不可控是因為其只能通過background=True/False控制是否進行多核處理，並不能調節使用核的數量。功能較簡單。

• pool中具有map_batches和imap_batches兩種選項，前者輸入為list列表，后者輸入為生成器generator。由於輸入的不同，后者在功能上還可以實現數據分批次的讀入內存，並且通過output_buffer_size參數控制圖像增強與模型訓練步調不一致的問題（限制pipeline中batch的數量，從而抑制pool貪婪的從生成器中獲取batch）。

值得注意的是上述兩種方法都是通過增強方法/序列點出來的。實際上是在多核上建立圖像和增強技術間的橋梁。