多工學習中的資料分佈問題(二)

orion發表於2021-12-01

在上一篇博文《多工學習中的資料分佈問題(一)》(連結：https://www.cnblogs.com/orion-orion/p/15621953.html)中我們提到論文[1]在聯邦學習的情景下引入了多工學習，其採用的手段是使每個client/Synthetic節點的訓練資料分佈不同，從而使各任務節點學習到不同的模型。

該論文的實驗手段是先後用FEMNIST、CIFAR10、Shakespare、Synthetic等資料集對模型進行測試，這些資料集包括CV、NLP、普通分類/迴歸這三種不同的任務。但是，該論文在給定用一組資料集進行測試的過程中，所有client節點上要解決的任務種類和執行的模型是一樣的(比如採用CIFAR10資料集，則所有client節點都採用MobileNet-v2網路；採用Shakespare資料集，所有任務節點採用Stacked-LSTM網路)。那麼此時疑惑就來了，既然單次實驗的資料集一樣，網路也一樣，那麼談何多工呢？文章採用的手段是，在單次實驗時，對原始資料集進行非獨立同分布(non_idd)
的隨機取樣，為\(T\)個不同非任務生成\(T\)個不同分佈的資料集，以做到每個任務節點訓練出的模型不同。

接下來，我們就來仔細研究論文的資料集劃分與隨機取樣演算法。該論文的程式碼已經開源在Github上[2]。以下我們以CIFAR10資料集的生成為例，來詳細地對該論文的資料集取樣演算法進行分析。

1.資料集匯入

首先，從torchvision中匯入訓練和測試資料集，並統一拼接成一個dataset。

from torchvision.datasets import CIFAR10
from torchvision.transforms import Compose, ToTensor, Normalize
from torch.utils.data import ConcatDataset
# 相對於本檔案的相對目錄
RAW_DATA_PATH = "raw_data/"
transform = Compose([
        ToTensor(),
        Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])

# 這個dataset物件可以用torch.utils.data.DataLoader並行載入
dataset =\
    ConcatDataset([
        # transform對input處理(targettransform對target處理)
        # download為true，會自動下載到引數root對應的目錄，如果已經有了，就不會下載
        # download為false，不會自動下載。
        # train = True，從訓練集create資料
        CIFAR10(root=RAW_DATA_PATH, download=True, train=True, transform=transform),
        # test = False，從測試集create資料
        CIFAR10(root=RAW_DATA_PATH, download=False, train=False, transform=transform)
    ])

2. 資料集拆分到client

然後有兩種劃分方式，一種是按照病態非獨立同分布來劃分資料（其實最開始是論文[2]提出的劃分方式，此時傳入的命令列引數args.pathological_non_iid_split=True），一種按照標籤對資料進行劃分（若沒有設定命令列引數args.pathological_non_iid_split ，則預設按照標籤劃分）。

2.1 病態獨立同分佈劃分方式((pathological non iid split))

我們先來看按照病態獨立同分布來劃分資料。如果選擇這種劃分方式，則每個client會受到\(n\)個shard（碎片）的資料集，每個碎片最多包含兩個類別。此時我們可以選擇傳入引數args.n_shard，該參數列示每個client/task的碎片數量（預設值為2）。（當然，如果沒有按照病態非獨立同分布來劃分資料，則不需要設定args.n_shard引數）
然後，我們將資料集劃分到各client上，此時我們需要將這個功能編寫成一個函式並進行呼叫：

clients_indices = \
    clients_indices =\
        pathological_non_iid_split(
            dataset=dataset,
            n_classes=N_CLASSES,
            n_clients=args.n_tasks,
            n_classes_per_client=args.n_shards,
            frac=args.s_frac,
            seed=args.seed
        )

接下來我們來看這個函式如何設計。先看函式原型：

def pathological_non_iid_split(dataset, n_classes, n_clients, n_classes_per_client, frac=1, seed=1234):

我們解釋一下函式的引數，這裡dataset是torch.utils.Dataset型別的資料集，n_classes表示資料集裡樣本分類數，n_client表示client節點的數量，n_client_per_client表示每個client中的類別數，frac是使用資料集的比例(預設是1，即使用全部資料)，seed是傳入的隨機數種子。該函式返回一個由n_client個subgroup組成的列表client_indices，每個subgroup對應某個client所需的樣本索引組成的列表。

接下來我們看這個函式的內容。該函式完成的功能可以概括為：先將樣本按照標籤進行排序；再將樣本劃分為n_client * n_classes_per_client個shards（每個shard大小相等），對n_clients中的每一個client分配n_classes_per_client個shards（分配到client後，每個client中的shards要合併）。

首先，我們根據frac獲取資料集的子集。

    rng_seed = (seed if (seed is not None and seed >= 0) else int(time.time())) 
    rng = random.Random(rng_seed)
    np.random.seed(rng_seed)

    # get subset
    n_samples = int(len(dataset) * frac)
    selected_indices = rng.sample(list(range(len(dataset))), n_samples)

然後從被選出的資料集索引selected_indices建立一個key為類別\(\{0,1,...,n\_classes-1\}\)，value為對應樣本集索引列表的字典，這在實際上這就相當於按照label對樣本進行排序了。

    label2index = {k: [] for k in range(n_classes)}
    for idx in selected_indices:
        _, label = dataset[idx]
        label2index[label].append(idx)

    sorted_indices = []
    for label in label2index:
        sorted_indices += label2index[label]

然後該函式將資料分為n_clients * n_classes_per_client 個獨立同分布的shards，每個shards大小相等。然後給n_clients中的每一個client分配n_classes_per_client個shards（分配到client後，每個client中的shards要合併），程式碼如下：

    def iid_divide(l, g):
        """
        將列表`l`分為`g`個獨立同分布的group（其實就是直接劃分）
        每個group都有 `int(len(l)/g)` 或者 `int(len(l)/g)+1` 個元素
        返回由不同的groups組成的列表
        """
        num_elems = len(l)
        group_size = int(len(l) / g)
        num_big_groups = num_elems - g * group_size
        num_small_groups = g - num_big_groups
        glist = []
        for i in range(num_small_groups):
            glist.append(l[group_size * i: group_size * (i + 1)])
        bi = group_size * num_small_groups
        group_size += 1
        for i in range(num_big_groups):
            glist.append(l[bi + group_size * i:bi + group_size * (i + 1)])
        return glist


    n_shards = n_clients * n_classes_per_client
    # 一共分成n_shards個獨立同分布的shards
    shards = iid_divide(sorted_indices, n_shards)
    random.shuffle(shards)
    # 然後再將n_shards拆分為n_client份
    tasks_shards = iid_divide(shards, n_clients)

    clients_indices = [[] for _ in range(n_clients)]
    for client_id in range(n_clients):
        for shard in tasks_shards[client_id]:
            # 這裡shard是一個shard的資料索引(一個列表)
            # += shard 實質上是在列表裡併入列表
            clients_indices[client_id] += shard

最後，返回clients_indices

    return clients_indices

2.2 按照標籤劃分劃分方式(split dataset by labels)

現在我們來看按照標籤來劃分資料。如果選擇這種劃分方式，則不再傳入引數args.n_shard進行shard的劃分。我們只需要將資料集標籤進行排序後直接劃分到各client上，此時我們需要將這個功能編寫成一個函式並進行呼叫：

clients_indices = \
    split_dataset_by_labels(
        dataset=dataset,
        n_classes=N_CLASSES,
        n_clients=args.n_tasks,
        n_clusters=args.n_components,
        alpha=args.alpha,
        frac=args.s_frac,
        seed=args.seed
    )

接下來我們來看這個函式如何設計。先看函式原型：

def split_dataset_by_labels(dataset, n_classes, n_clients, n_clusters, alpha, frac, seed=1234):

我們解釋一下函式的引數，這裡dataset是torch.utils.Dataset型別的資料集，n_classes表示資料集裡樣本分類數，n_clusters是簇的個數（後面會解釋其含義，如果設定為-1，則就預設n_clusters=n_classes），alpha 用於控制clients之間的資料diversity（多樣性），frac是使用資料集的比例(預設是1，即使用全部資料)，seed是傳入的隨機數種子。該函式返回一個由n_client個subgroup組成的列表client_indices，每個subgroup對應某個client所需的樣本索引組成的列表。

接下來我們看這個函式的內容。這個函式的內容可以概括為：先將所有類別分組為n_clusters個簇；再對每個簇c，將樣本劃分給不同的clients（每個client的樣本數量按照dirichlet分佈來確定）。

首先，我們判斷n_clusters的數量，如果為-1，則預設每一個cluster對應一個資料class：

    if n_clusters == -1:
        n_clusters = n_classes

然後得到隨機數生成器（簡稱rng）：

    rng_seed = (seed if (seed is not None and seed >= 0) else int(time.time()))
    rng = random.Random(rng_seed)
    np.random.seed(rng_seed)

然後將打亂後的標籤集合\(\{0,1,...,n\_classes-1\}\)分為n_clusters個獨立同分布的簇。

    all_labels = list(range(n_classes))
    rng.shuffle(all_labels)
    clusters_labels = iid_divide(all_labels, n_clusters)

然後再建立根據上面劃分為簇的標籤(clusters_labels)建立key為label, value為簇id(group_idx)的字典，

    label2cluster = dict()  # maps label to its cluster
    for group_idx, labels in enumerate(clusters_labels):
        for label in labels:
            label2cluster[label] = group_idx

接著獲取資料集的子集

    n_samples = int(len(dataset) * frac)
    selected_indices = rng.sample(list(range(len(dataset))), n_samples)

之後，我們

    # 記錄每個cluster大小的向量
    clusters_sizes = np.zeros(n_clusters, dtype=int)
    # 儲存每個cluster對應的資料索引
    clusters = {k: [] for k in range(n_clusters)}
    for idx in selected_indices:
        _, label = dataset[idx]
        # 由樣本資料的label先找到其cluster的id
        group_id = label2cluster[label]
        # 再將對應cluster的大小+1
        clusters_sizes[group_id] += 1
        # 將樣本索引加入其cluster對應的列表中
        clusters[group_id].append(idx)

    # 將每個cluster對應的樣本索引列表打亂
    for _, cluster in clusters.items():
        rng.shuffle(cluster)

接著，我們按照dirichlet分佈設定每一個cluster的樣本個數。

    # 記錄來自每個cluster的client的樣本數量
    clients_counts = np.zeros((n_clusters, n_clients), dtype=np.int64) 

    # 遍歷每一個cluster
    for cluster_id in range(n_clusters):
        # 對每個cluster中的每個client賦予一個滿足dirichlet分佈的權重
        weights = np.random.dirichlet(alpha=alpha * np.ones(n_clients))
        # np.random.multinomial 表示投擲骰子clusters_sizes[cluster_id]次，落在各client上的權重依次是weights
        # 該函式返回落在各client上各多少次，也就對應著各client應該分得的樣本數
        clients_counts[cluster_id] = np.random.multinomial(clusters_sizes[cluster_id], weights)

    # 對每一個cluster上的每一個client的計數次數進行字首（累加）求和，
    # 相當於最終返回的是每一個cluster中按照client進行劃分的樣本分界點下標
    clients_counts = np.cumsum(clients_counts, axis=1)

然後，我們根據每一個cluster中的每一個client分得的樣本情況（我們已經得到了每一個cluster中按照client進行劃分的樣本分界點下標），合併歸納得到每一個client中分得的樣本情況。

    def split_list_by_indices(l, indices):
        """
        將列表`l` 劃分為長度為 `len(indices)` 的子列表
        第`i`個子列表從下標 `indices[i]` 到下標`indices[i+1]`
        （從下標0到下標`indices[0]`的子列表另算）
        返回一個由多個子列表組成的列表
        """
        res = []
        current_index = 0
        for index in indices: 
            res.append(l[current_index: index])
            current_index = index

        return res
    
    clients_indices = [[] for _ in range(n_clients)]
    for cluster_id in range(n_clusters):
        # cluster_split為一個cluster中按照client劃分好的樣本
        cluster_split = split_list_by_indices(clusters[cluster_id], clients_counts[cluster_id])

        # 將每一個client的樣本累加上去
        for client_id, indices in enumerate(cluster_split):
            clients_indices[client_id] += indices

最後，我們返回每個client對應的樣本索引：

    return clients_indices

3. 總結

按照病態獨立同分佈劃分和按照樣本標籤劃分兩種方式，其實本質上都是要使每個client的分佈不同，而這也是我們進行多工學習的前提。

參考文獻

  • [1] Marfoq O, Neglia G, Bellet A, et al. Federated multi-task learning under a mixture of distributions[J]. Advances in Neural Information Processing Systems, 2021, 34.
  • [2] McMahan B, Moore E, Ramage D, et al. Communication-efficient learning of deep networks from decentralized data[C]//Artificial intelligence and statistics. PMLR, 2017: 1273-1282.

相關文章