在前面的文章中，我們簡單介紹了在 MegEngine imperative 中的各模塊以及它們的作用。對于新用戶而言可能不太了解各個(gè)模塊的使用方法，對于模塊的結(jié)構(gòu)和原理也是一頭霧水。Python 作為現(xiàn)在深度學(xué)習(xí)領(lǐng)域的主流編程語言，其相關(guān)的模塊自然也是深度學(xué)習(xí)框架的重中之重。

模塊串講將對 MegEngine 的 Python 層相關(guān)模塊分別進(jìn)行更加深入的介紹，會(huì)涉及到一些原理的解釋和代碼解讀。Python 層模塊串講共分為上、中、下三個(gè)部分，本文介紹 Python 層的 data 模塊。讀者將通過本文了解到要構(gòu)建數(shù)據(jù) pipeline 所需要的對象，以及如何高效地構(gòu)建 pipeline。

構(gòu)建數(shù)據(jù)處理 pipeline —— data 模塊

神經(jīng)網(wǎng)絡(luò)需要數(shù)據(jù)才可以訓(xùn)練，數(shù)據(jù)源文件可能是各種格式，讀取數(shù)據(jù)需要定義采樣規(guī)則、數(shù)據(jù)變換規(guī)則、數(shù)據(jù)合并策略等，這些和數(shù)據(jù)相關(guān)的模塊都封裝在 data 下。

在 MegEngine 中訓(xùn)練模型讀取數(shù)據(jù)的 pipeline 一般是：

1. 創(chuàng)建一個(gè) Dataset 對象；
2. 按照訓(xùn)練場景的需求可能需要對數(shù)據(jù)做一些變換或合并的處理，這里可能需要?jiǎng)?chuàng)建 Sampler、Transform、Collator 等對象來完成相應(yīng)操作；
3. 創(chuàng)建一個(gè) DataLoader 對象；
4. 將數(shù)據(jù)分批加載到 DataLoader 里，迭代 DataLoader 對象進(jìn)行訓(xùn)練。

下面我們看一下這幾個(gè)對象的實(shí)現(xiàn)。

Dataset

在 MegEngine 中，數(shù)據(jù)集是一個(gè)可迭代的對象，所有的 Dataset 對象都繼承自 class Dataset，都需要實(shí)現(xiàn)自己的 __getitem__() 方法和 __len__() 方法，這兩個(gè)方法分別是用來獲取給定索引的對應(yīng)的數(shù)據(jù)樣本和返回?cái)?shù)據(jù)集的大小。

根據(jù)對數(shù)據(jù)集訪問方式的區(qū)別，MegEngine 中的數(shù)據(jù)集類型主要分為兩種：ArrayDataset 和 StreamDataset，前者支持隨機(jī)訪問數(shù)據(jù)樣本，而后者只可以順序訪問。二者的主要區(qū)別見下表：

Dataset

Dataset 支持對數(shù)據(jù)集的隨機(jī)訪問，訪問類型是 Map-style 的，也就是可以從索引映射到數(shù)據(jù)樣本，使用時(shí)需要實(shí)現(xiàn) __getitem__() 方法和 __len__() 方法。

下面是一個(gè)使用 Dataset 生成一個(gè)由 0 到 5 的數(shù)組成的數(shù)據(jù)集的例子：

from megengine.data.dataset import Dataset

class CustomMapDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, idx):
        return self.data[idx]

    def __len__(self):
        return len(self.data)

使用起來如下：

>>> data = list(range(0, 5))
>>> map_dataset = CustomMapDataset(data)
>>> print(len(map_dataset))
5
>>> print(map_dataset[2])
2

可以發(fā)現(xiàn) Dataset 最大的特點(diǎn)就是可以根據(jù)給定索引隨機(jī)訪問數(shù)據(jù)集中對應(yīng)下標(biāo)的數(shù)據(jù)。

ArrayDataset

對于 Numpy ndarray 類型的數(shù)據(jù)，MegEngine 中對 Dataset 進(jìn)一步封裝實(shí)現(xiàn)了 ArrayDataset，使用 ArrayDataset 無需實(shí)現(xiàn) __getitem__() 方法和 __len__() 方法。

下面的例子隨機(jī)生成了一個(gè)具有 100 個(gè)樣本、每張樣本為 32 × 32 像素的 RGB 圖片的數(shù)據(jù)集：

import numpy as np
from megengine.data.dataset import ArrayDataset

data = np.random.random((100, 3, 32, 32))
target = np.random.random((100, 1))
dataset = ArrayDataset(data, target)
>>> print(len(dataset))
100
>>> print(type(dataset[0]), len(dataset[0]))
<class 'tuple'> 2
>>> print(dataset[0][0].shape)
(3, 32, 32)

由于需要支持隨機(jī)訪問，因此對于支持順序訪問的 Dataset 需要將索引等信息加載進(jìn)內(nèi)存，如果數(shù)據(jù)集規(guī)模較大導(dǎo)致內(nèi)存無法存放從而發(fā)生 OOM（Out Of Memory），我們需要考慮使用流式數(shù)據(jù) StreamDataset。

StreamDataset

當(dāng)數(shù)據(jù)集規(guī)模較大時(shí)，使用流失數(shù)據(jù)迭代訪問數(shù)據(jù)對象是比較主流的做法。從類的定義可以看到：由于無法根據(jù)索引獲取數(shù)據(jù)，因此 StreamDataset 無需實(shí)現(xiàn) __getitem__() 方法和 __len__() 方法，但是需要實(shí)現(xiàn)一個(gè) __iter__() 方法定義流式獲取數(shù)據(jù)的規(guī)則：

class StreamDataset(Dataset):
    r"""An abstract class for stream data.
    __iter__ method is aditionally needed.
    """

    @abstractmethod
    def __init__(self):
        pass

    @abstractmethod
    def __iter__(self):
        pass

    def __getitem__(self, idx):
        raise AssertionError("can not get item from StreamDataset by index")

    def __len__(self):
        raise AssertionError("StreamDataset does not have length")

StreamDataset 適用的場景主要是：

• 隨機(jī)讀取成本過高，或者數(shù)據(jù)規(guī)模太大，無法支持；
• 必須根據(jù)流數(shù)據(jù)才能判斷當(dāng)前批是否已經(jīng)完整。

可以使用流數(shù)據(jù)返回從數(shù)據(jù)庫、遠(yuǎn)程服務(wù)器甚至實(shí)時(shí)生成的日志中讀取的數(shù)據(jù)流。

下面的例子展示了如何生成一個(gè)由 0 到 5 這五個(gè)數(shù)組成的數(shù)據(jù)集：

from megengine.data.dataset import StreamDataset

class CustomIterableDataset(StreamDataset):
    def __init__(self, data):
        self.data = data

    def __iter__(self):
        return iter(self.data)
>>> data = list(range(0, 5))
>>> iter_dataset = CustomIterableDataset(data)
>>> it = iter(iter_dataset)
>>> print(type(it))
list_iterator
>>> print(next(it))
0
>>> print(next(it))
1

Sampler

有了 DataSet 之后，DataLoader 可以從數(shù)據(jù)集加載數(shù)據(jù)到內(nèi)存，但是對每批數(shù)據(jù)有時(shí)候需要規(guī)定規(guī)模的大小，還有定義抽樣規(guī)則等需求，使用 Sampler 可以對每批數(shù)據(jù)的抽樣規(guī)則進(jìn)行自定義。

準(zhǔn)確來說，抽樣器的職責(zé)是決定數(shù)據(jù)的獲取順序，方便為 DataLoader 提供一個(gè)可供迭代的多批數(shù)據(jù)的索引：

dataloader = DataLoader(dataset, sampler=RandomSampler)

在 MegEngine 中，Sampler 是所有抽樣器的抽象基類，在大部分情況下用戶無需對抽樣器進(jìn)行自定義實(shí)現(xiàn)， 因?yàn)樵?MegEngine 中已經(jīng)實(shí)現(xiàn)了常見的各種抽樣器，比如上面示例代碼中的 RandomSampler 抽樣器。

下面介紹 MegEngine 中幾種常見的 Sampler。

SequentialSampler

SequentialSampler 也叫 MapSampler， 顧名思義就是對數(shù)據(jù)集進(jìn)行順序抽樣的抽樣器。

對一個(gè)含有 100 個(gè)數(shù)據(jù)樣本的數(shù)據(jù)集，batch_size 為 10，可以得到 10 批順序索引：

>>> from megengine.data import SequentialSampler
>>> sampler = SequentialSampler(image_dataset, batch_size=10)
>>> print(len(list(sampler)))
10
如果將 batch_size 修改為 30, 則會(huì)得到 4 批順序索引，最后一批長度為 10：

>>> sampler = SequentialSampler(image_dataset, batch_size=30)
>>> for batch_id, indices in enumerate(sampler):
...     print(batch_id, len(indices))
0 30
1 30
2 30
3 10
我們可以通過設(shè)置 drop_last=True 丟掉最后一批不完整的索引：

>>> sampler = SequentialSampler(image_dataset, 30, drop_last=True)
>>> for batch_id, indices in enumerate(sampler):
....    print(batch_id, len(indices))
0 30
1 30

默認(rèn)情況下 batch_size 為 1，表示逐個(gè)遍歷數(shù)據(jù)集中的樣本，drop_last 為 False。

RandomSampler

RandomSampler 用來對數(shù)據(jù)集進(jìn)行無放回隨機(jī)抽樣（也叫簡單隨機(jī)抽樣）。

直接看例子：

>>> from megengine.data import RandomSampler
>>> sampler = RandomSampler(image_dataset, batch_size=10)
>>> for batch_id, indices in enumerate(sampler):
...     print(batch_id, indices)
0 [78, 20, 74, 6, 45, 65, 99, 67, 88, 57]
1 [81, 0, 94, 98, 71, 30, 66, 10, 85, 56]
2 [51, 87, 62, 42, 7, 75, 11, 12, 39, 95]
3 [73, 15, 77, 72, 89, 13, 55, 26, 49, 33]
4 [9, 8, 64, 3, 37, 2, 70, 29, 34, 47]
5 [22, 18, 93, 4, 40, 92, 79, 36, 84, 25]
6 [83, 90, 68, 58, 50, 48, 32, 54, 35, 1]
7 [14, 44, 17, 63, 60, 97, 96, 23, 52, 38]
8 [80, 59, 53, 19, 46, 43, 24, 61, 16, 5]
9 [86, 82, 31, 76, 28, 91, 27, 21, 69, 41]

ReplacementSampler

ReplacementSampler 是有放回隨機(jī)抽樣，也就是可能抽樣到之前已經(jīng)抽樣過的數(shù)據(jù)。

使用方法和無放回隨機(jī)抽樣類似：

>>> from megengine.data import ReplacementSampler
>>> sampler = ReplacementSampler(image_dataset, batch_size=10)
>>> for batch_id, indices in enumerate(sampler):
...     print(batch_id, indices)
0 [58, 29, 42, 79, 91, 73, 86, 46, 85, 23]
1 [42, 33, 61, 8, 22, 10, 98, 56, 59, 96]
2 [38, 72, 26, 0, 40, 33, 30, 59, 1, 25]
3 [71, 95, 89, 88, 29, 97, 97, 46, 42, 0]
4 [42, 22, 28, 82, 49, 52, 88, 68, 46, 66]
5 [47, 62, 26, 17, 68, 31, 70, 69, 26, 4]
6 [43, 18, 17, 91, 99, 96, 91, 7, 24, 39]
7 [50, 55, 86, 65, 93, 38, 39, 4, 6, 60]
8 [92, 82, 61, 36, 67, 56, 24, 18, 70, 60]
9 [91, 63, 95, 99, 19, 47, 9, 9, 68, 37]

Infinite

通常數(shù)據(jù)集在給定 batch_size 的情況下，只能劃分為有限個(gè) batch。 這意味著抽樣所能得到的數(shù)據(jù)批數(shù)是有限的，想要重復(fù)利用數(shù)據(jù)， 最常見的做法是循環(huán)多個(gè)周期 epochs 來反復(fù)遍歷數(shù)據(jù)集：

for epoch in epochs:
    for batch_data in dataloader:

但在一些情況下，我們希望能夠直接從數(shù)據(jù)集中無限進(jìn)行抽樣， 因此MegEngine提供了 Infinite 包裝類用來進(jìn)行無限抽樣：

>>> from megengine.data import Infinite
>>> sampler = Infinite(SequentialSampler(image_dataset, batch_size=10))
>>> sample_queue = iter(sampler)
>>> for step in range(20):
...     indice = next(sample_queue)
...     print(step, indice)
0 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1 [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
2 [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
3 [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
4 [40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
5 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59]
6 [60, 61, 62, 63, 64, 65, 66, 67, 68, 69]
7 [70, 71, 72, 73, 74, 75, 76, 77, 78, 79]
8 [80, 81, 82, 83, 84, 85, 86, 87, 88, 89]
9 [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]
10 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
11 [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
12 [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
13 [30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
14 [40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
15 [50, 51, 52, 53, 54, 55, 56, 57, 58, 59]
16 [60, 61, 62, 63, 64, 65, 66, 67, 68, 69]
17 [70, 71, 72, 73, 74, 75, 76, 77, 78, 79]
18 [80, 81, 82, 83, 84, 85, 86, 87, 88, 89]
19 [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]

以上就是常見的 Sampler 的使用方法，有時(shí)候?qū)τ跀?shù)據(jù)集中的數(shù)據(jù)還需要做一些變換以滿足業(yè)務(wù)需要，這就是我們接下來要說的 transform。

Transform

在深度學(xué)習(xí)中對數(shù)據(jù)進(jìn)行變換（Transformation）以滿足業(yè)務(wù)需求和增強(qiáng)模型性能是很常見的操作。

在 megengine.data.transform 中提供的各種數(shù)據(jù)變換都是基于 Transform 抽象類實(shí)現(xiàn)的，其中：

• apply 抽象方法可用于單個(gè)的數(shù)據(jù)樣本， 需要在子類中實(shí)現(xiàn)；
• 各種變換操作可以通過 Compose 進(jìn)行組合，這樣使用起來更加方便。

我們能夠很方便地在 DataLoader 加載數(shù)據(jù)時(shí)進(jìn)行相應(yīng)地變換操作。例如：

dataloader = DataLoader(dataset, transform=Compose([Resize(32), ToMode('CHW')]))

上面就是將兩個(gè) transform 操作 Resize() 和 ToMode() 組合起來對數(shù)據(jù)進(jìn)行變換。

下面舉個(gè)例子如何實(shí)現(xiàn)自己的 Transform：

>>> from megengine.data.transform import Transform
>>> class AddOneTransform(Transform):
...     def apply(self, input):
...         return input + 1
>>> AddOneTransform().apply(data)
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

上面這個(gè) Transform 實(shí)現(xiàn)了自己的 apply() 方法，對數(shù)據(jù)集中的所有樣本做了一個(gè) +1 操作。

可以使用 Compose 對數(shù)據(jù)變換進(jìn)行組合：

>>> from megengine.data.transform import Compose
>>> composed_transform = Compose([AddOneTransform(), AddOneTransform()])
>>> composed_transform.apply(data)
array([[ 2,  3,  4],
       [ 5,  6,  7],
       [ 8,  9, 10]])

最終，我們的各種Transform實(shí)現(xiàn)應(yīng)當(dāng)被應(yīng)用于DataLoader：

dataloader = DataLoader(dataset, transform=composed_transform)

實(shí)際使用時(shí)對數(shù)據(jù)做的操作往往比上面的例子要復(fù)雜許多，MegEngine 在 VisionTransform 中已經(jīng)實(shí)現(xiàn)了很多轉(zhuǎn)換方法供用戶使用。用戶也可以根據(jù)需要實(shí)現(xiàn)自己的數(shù)據(jù)變換方法。

當(dāng)我們從 DataLoader 中獲取批數(shù)據(jù)時(shí)，如果定義了 Transform, 則會(huì)在每次加載完樣本后立即對其進(jìn)行變換。

數(shù)據(jù)變換操作也是有計(jì)算開銷的，且該流程通常在 CPU 設(shè)備上進(jìn)行，以及有些操作會(huì)調(diào)用類似 OpenCV 的庫。 如果我們對每個(gè)樣本進(jìn)行多次加載（比如訓(xùn)練多個(gè)周期），那么變換操作也會(huì)被執(zhí)行多次，這可能會(huì)帶來額外的開銷。 因此在有些時(shí)候，我們會(huì)選擇將預(yù)處理操作在更早的流程中進(jìn)行，即直接對原始數(shù)據(jù)先進(jìn)行一次預(yù)處理操作， 這樣在 DataLoader 中獲取的輸入便已經(jīng)是經(jīng)過預(yù)處理的數(shù)據(jù)了，這樣可以盡可能地減少 Transform 操作。

用戶應(yīng)當(dāng)考慮到，原始數(shù)據(jù)相關(guān)的 I/O 和處理也有可能成為模型訓(xùn)練整體流程中的瓶頸。

Collator

在使用 DataLoader 獲取批數(shù)據(jù)的整個(gè)流程中， Collator 負(fù)責(zé)合并樣本，最終得到批數(shù)據(jù)。

Collator 僅適用于 Map-style 的數(shù)據(jù)集，因?yàn)?Iterable-style 數(shù)據(jù)集的批數(shù)據(jù)必然是逐個(gè)合并的。

經(jīng)過 DataSet 和 Transform 的處理后， Collator 通常會(huì)接收到一個(gè)列表：

• 如果你的 Dataset 子類的 __getitem__ 方法返回的是單個(gè)元素，則 Collator 得到一個(gè)普通列表；
• 如果你的 Dataset 子類的 __getitem__ 方法返回的是一個(gè)元組，則 Collator 得到一個(gè)元組列表。

MegEngine 中使用 Collator 作為默認(rèn)實(shí)現(xiàn)，通過調(diào)用 apply 方法來將列表數(shù)據(jù)合并成批數(shù)據(jù)：

from megengine.data import Collator
collator = Collator()

默認(rèn)的 Collator 支持 NumPy ndarray, Numbers, Unicode strings, bytes, dicts 或 lists 數(shù)據(jù)類型。 要求輸入必須包含至少一種上述數(shù)據(jù)類型，否則用戶需要使用自己定義的 Collator。

Collator 的作用是合并數(shù)據(jù)，比如每個(gè)數(shù)據(jù)樣本是 shape 為 (C, H, W) 的圖片，如果我們在 Sampler 中指定了 batch_size 為 N。那么 Collator 就會(huì)將獲得的樣本列表合并成一個(gè) shape 為 (N, C, H, W) 的批樣本結(jié)構(gòu)。

我們可以模擬得到這樣一個(gè) image_list 數(shù)據(jù)，并借助 Collator 得到 batch_image：

>>> N, C, H, W = 5, 3, 32, 32
>>> image_list = []
>>> for i in range(N):
...     image_list.append(np.random.random((C, H, W)))
>>> print(len(image_list), image_list[0].shape)
5 (3, 32, 32)
>>> batch_image = collator.apply(image_list)
>>> batch_image.shape
(5, 3, 32, 32)

DataLoader

前面介紹的 Dataset、Sampler、Transform、Collator 等對象都是為了更靈活地配置 DataLoader 對象的。

當(dāng)單進(jìn)程運(yùn)行 DataLoader 時(shí)（設(shè)置 num_workers=0），每當(dāng)我們向 DataLoader 索要一批數(shù)據(jù)時(shí)，DataLoader 將從 Sampler 獲得下一批數(shù)據(jù)的索引， 根據(jù) Dataset 提供的 __getitem__() 方法將對應(yīng)的數(shù)據(jù)逐個(gè)加載到內(nèi)存， 加載進(jìn)來的數(shù)據(jù)可以通過指定的 Transform 做一些處理，再通過 Collator 將單獨(dú)的數(shù)據(jù)組織成批數(shù)據(jù)。

DataLoader 也支持多進(jìn)程加載以提升數(shù)據(jù)加載處理速度（提高 num_workers 數(shù)量）。 一般 worker 數(shù)量越多，數(shù)據(jù)加載處理的速度會(huì)越快。不過如果 worker 數(shù)過多， 并大大超出了系統(tǒng)中 cpu 的數(shù)量，這些子進(jìn)程可能會(huì)存在競爭 cpu 資源的情況，反而導(dǎo)致效率的降低。

一般來說，我們建議根據(jù)系統(tǒng)中 cpu 的數(shù)量設(shè)置 worker 的值。 比如在一臺(tái) 64 cpu, 8 gpu 的機(jī)器上，預(yù)期中每個(gè) gpu 會(huì)對應(yīng) 8 個(gè) cpu, 那么我們在使用時(shí)對應(yīng)的把 worker 數(shù)設(shè)置在 8 左右就是個(gè)不錯(cuò)的選擇。

下面以一個(gè)加載圖像分類數(shù)據(jù)的流程來舉例說明如何創(chuàng)建一個(gè)加載數(shù)據(jù)的 pipeline。

1、假設(shè)圖像數(shù)據(jù)按照一定的規(guī)則放置于同一目錄下（通常數(shù)據(jù)集主頁會(huì)對目錄組織和文件命名規(guī)則進(jìn)行介紹）。 要?jiǎng)?chuàng)建對應(yīng)的數(shù)據(jù)加載器，首先需要一個(gè)繼承自 Dataset 的類。 我們可以創(chuàng)建一個(gè)自定義的數(shù)據(jù)集：

import cv2
import numpy as np
import megengine
from megengine.data.dataset import Dataset

class CustomImageDataset(Dataset):
    def __init__(self, image_folder):
        # get all mapping indice
        self.image_folder = image_folder
        self.image_list = os.listdir(image_folder)

    # get the sample
    def __getitem__(self, idx):
        # get the index
        image_file = self.image_list[idx]

        # get the data
        # in this case we load image data and convert to ndarray
        image = cv2.imread(self.image_folder + image_file, cv2.IMREAD_COLOR)
        image = np.array(image)

        # get the label
        # in this case the label was noted in the name of the image file
        # ie: 1_image_28457.png where 1 is the label
        # and the number at the end is just the id or something
        target = int(image_file.split("_")[0])

        return image, target

    def __len__(self):
        return len(self.images)

要獲取示例圖像，可以創(chuàng)建一個(gè)數(shù)據(jù)集對象，并將示例索引傳遞給__getitem__()方法， 然后將返回圖像數(shù)組和對應(yīng)的標(biāo)簽，例如：

dataset = CustomImageDataset("/path/to/image/folder")
data, sample = dataset.__getitem__(0) # dataset[0]

2、現(xiàn)在我們已經(jīng)預(yù)先創(chuàng)建了能夠返回一個(gè)樣本及其標(biāo)簽的類CustomImageDataset, 但僅依賴Dataset本身還無法實(shí)現(xiàn)自動(dòng)分批、亂序、并行等功能； 我們必須接著創(chuàng)建DataLoader, 它通過其它的參數(shù)配置項(xiàng)圍繞這個(gè)類“包裝”， 可以按照我們的要求從數(shù)據(jù)集類中返回整批樣本。

from megengine.data.transform import ToMode
from megengine.data import DataLoader, RandomSampler

dataset = YourImageDataset("/path/to/image/folder")

# you can implement the function to randomly split your dataset
train_set, val_set, test_set = random_split(dataset)

# B is your batch-size, ie. 128
train_dataloader = DataLoader(train_set,
      sampler=RandomSampler(train_set, batch_size=B),
      transform=ToMode('CHW'),
)

3、現(xiàn)在可以加載數(shù)據(jù)并進(jìn)行訓(xùn)練了：

for epoch in range(epochs):

    for images, targets in train_dataloder:
        # now 'images' is a batch containing B samples
        # and 'targets' is a batch containing B targets
        # (of the images in 'images' with the same index

        # remember to convert data to tensor
        images = megengine.Tensor(images)
        targets = megengine.Tensor(targets)

        # train function
        # ...