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

模塊串講將對 MegEngine 的 Python 層相關(guān)模塊分別進行更加深入的介紹溃列，會涉及到一些原理的解釋和代碼解讀面殖。Python 層模塊串講共分為上、中哭廉、下三個部分脊僚，本文介紹 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)建一個 Dataset 對象拳喻；
2. 按照訓(xùn)練場景的需求可能需要對數(shù)據(jù)做一些變換或合并的處理，這里可能需要創(chuàng)建 Sampler猪腕、Transform冗澈、Collator 等對象來完成相應(yīng)操作；
3. 創(chuàng)建一個 DataLoader 對象陋葡；
4. 將數(shù)據(jù)分批加載到 DataLoader 里亚亲，迭代 DataLoader 對象進行訓(xùn)練。

下面我們看一下這幾個對象的實現(xiàn)腐缤。

Dataset

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

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

Dataset

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

下面是一個使用 Dataset 生成一個由 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 最大的特點就是可以根據(jù)給定索引隨機訪問數(shù)據(jù)集中對應(yīng)下標(biāo)的數(shù)據(jù)质帅。

ArrayDataset

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

下面的例子隨機生成了一個具有 100 個樣本嫉嘀、每張樣本為 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)

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

StreamDataset

當(dāng)數(shù)據(jù)集規(guī)模較大時，使用流失數(shù)據(jù)迭代訪問數(shù)據(jù)對象是比較主流的做法瓣俯。從類的定義可以看到：由于無法根據(jù)索引獲取數(shù)據(jù)杰标，因此 StreamDataset 無需實現(xiàn) __getitem__() 方法和 __len__() 方法，但是需要實現(xiàn)一個 __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 適用的場景主要是：

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

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

下面的例子展示了如何生成一個由 0 到 5 這五個數(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ù)有時候需要規(guī)定規(guī)模的大小邓了，還有定義抽樣規(guī)則等需求毡泻，使用 Sampler 可以對每批數(shù)據(jù)的抽樣規(guī)則進行自定義胜茧。

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

dataloader = DataLoader(dataset, sampler=RandomSampler)

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

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

SequentialSampler

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

對一個含有 100 個數(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, 則會得到 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绒净，表示逐個遍歷數(shù)據(jù)集中的樣本见咒，drop_last 為 False。

RandomSampler

RandomSampler 用來對數(shù)據(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īng)抽樣過的數(shù)據(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 的情況下缤言，只能劃分為有限個 batch宝当。 這意味著抽樣所能得到的數(shù)據(jù)批數(shù)是有限的，想要重復(fù)利用數(shù)據(jù)胆萧， 最常見的做法是循環(huán)多個周期 epochs 來反復(fù)遍歷數(shù)據(jù)集：

for epoch in epochs:
    for batch_data in dataloader:

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

>>> 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 的使用方法，有時候?qū)τ跀?shù)據(jù)集中的數(shù)據(jù)還需要做一些變換以滿足業(yè)務(wù)需要盾鳞，這就是我們接下來要說的 transform犬性。

Transform

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

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

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

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

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

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

下面舉個例子如何實現(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]])

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

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

>>> 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實現(xiàn)應(yīng)當(dāng)被應(yīng)用于DataLoader：

dataloader = DataLoader(dataset, transform=composed_transform)

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

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

數(shù)據(jù)變換操作也是有計算開銷的违施，且該流程通常在 CPU 設(shè)備上進行互纯，以及有些操作會調(diào)用類似 OpenCV 的庫。 如果我們對每個樣本進行多次加載（比如訓(xùn)練多個周期）磕蒲，那么變換操作也會被執(zhí)行多次留潦，這可能會帶來額外的開銷。 因此在有些時候辣往，我們會選擇將預(yù)處理操作在更早的流程中進行兔院，即直接對原始數(shù)據(jù)先進行一次預(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ù)的整個流程中屹堰， Collator 負(fù)責(zé)合并樣本肛冶，最終得到批數(shù)據(jù)街氢。

Collator 僅適用于 Map-style 的數(shù)據(jù)集，因為 Iterable-style 數(shù)據(jù)集的批數(shù)據(jù)必然是逐個合并的睦袖。

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

• 如果你的 Dataset 子類的 __getitem__ 方法返回的是單個元素，則 Collator 得到一個普通列表；
• 如果你的 Dataset 子類的 __getitem__ 方法返回的是一個元組伦乔，則 Collator 得到一個元組列表厉亏。

MegEngine 中使用 Collator 作為默認(rèn)實現(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ù)招刹，比如每個數(shù)據(jù)樣本是 shape 為 (C, H, W) 的圖片恬试，如果我們在 Sampler 中指定了 batch_size 為 N。那么 Collator 就會將獲得的樣本列表合并成一個 shape 為 (N, C, H, W) 的批樣本結(jié)構(gòu)疯暑。

我們可以模擬得到這樣一個 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)單進程運行 DataLoader 時（設(shè)置 num_workers=0）越锈，每當(dāng)我們向 DataLoader 索要一批數(shù)據(jù)時仗嗦，DataLoader 將從 Sampler 獲得下一批數(shù)據(jù)的索引， 根據(jù) Dataset 提供的 __getitem__() 方法將對應(yīng)的數(shù)據(jù)逐個加載到內(nèi)存瞪浸， 加載進來的數(shù)據(jù)可以通過指定的 Transform 做一些處理儒将，再通過 Collator 將單獨的數(shù)據(jù)組織成批數(shù)據(jù)。

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

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

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

1再芋、假設(shè)圖像數(shù)據(jù)按照一定的規(guī)則放置于同一目錄下（通常數(shù)據(jù)集主頁會對目錄組織和文件命名規(guī)則進行介紹）菊霜。 要創(chuàng)建對應(yīng)的數(shù)據(jù)加載器，首先需要一個繼承自 Dataset 的類济赎。 我們可以創(chuàng)建一個自定義的數(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)建一個數(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)建了能夠返回一個樣本及其標(biāo)簽的類CustomImageDataset, 但僅依賴Dataset本身還無法實現(xiàn)自動分批、亂序勾徽、并行等功能滑凉； 我們必須接著創(chuàng)建DataLoader, 它通過其它的參數(shù)配置項圍繞這個類“包裝”， 可以按照我們的要求從數(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ù)并進行訓(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
        # ...