연합학습 (Federated Learning) LEAF 데이터셋 사용법 (1) - FEMNIST 데이터셋, pytorch loader 구현

Federated Learning Dataset?

최근 FL 공부를 다시 좀 하면서 직접 framework를 구현하고 있다. 관련하여서 가장 귀찮은게 일반적인 데이터셋을 그대로 사용하기가 어렵다는 점이다. 실제 FL의 목적은 여러 device에서 학습을 진행하고 이를 server에 보내서 aggregation을 하는 과정을 거치지만, 많은 연구들에서는 하나의 플랫폼(일반적으로 충분한 연산 능력을 갖춘 서버나 데스크톱)에서 시뮬레이션과 같은 형태로 FL을 구현한다.

 

따라서 하나의 플랫폼에서 데이터를 적절히 분류하여서 FL에 적합한 형태로 정리할 필요가 있다. 여기서 이 '적절히'라는게 정말 애매한 조건이다.

 

FL을 제안한 논문에서 FL에서 기본적인 FL을 위한 데이터셋의 구성에 대해서 언급을 하고 있다.

이중 Non-IID, Unbalanced, Massively distributed이 세가지가 데이터셋 구성에 있어서 중요한 부분이라고 할 수 있다.  IID (Independent and Identically Distribution)는 한국어로는 독립 항등분포로 FL의 측면에서는 client들이 가지는 데이터가 독립적이고, 동시에 모든 client들의 데이터 분포가 동일하다는 말이 된다. 그러나 실제로 모든 client들이 가지는 데이터의 분포가 동일하리라는 법은 없다. 오히려 희박하다. 따라서 Non-IID라는 가정이 필요하다. 다음은 Unbalanced인데, 각 cllient들이 반드시 동일한 양의 데이터를 가지리란 법이 없으므로 data 개수의 분포가 unbalance하다는 가정을 깔고 가야한다. 마지막 Massively distributed, 여러 client가(꽤나 많이) 존재해야 한다는 것이다. FL이 기본적으로 해결하고자 하는 문제를 고려하면 당연한 이야기이다.

 

LEAF!

문제는 non-iid는 유일하지 않고, unbalance의 정도와 client의 수 모두 임의로 설정할 수 있는 parameter이다. 따라서 FL 관련 연구들에서 사용하는 데이터셋에 대한 공정한 평가가 쉽지 않은 것이 현실이다. 이런 문제를 해결하기 위해서 CMU와 Google, Determined AI에서 LEAF라는 FL을 위한 Benchmark를 공개했다. LEAF는 데이터셋 뿐만 아니라 기본적인 Federated Learning Framework 또한 제공하여 연구에 활용하기 간편한 도구라고 볼 수 있으시겠다.

 

Homepage Link, github, paper

구체적으로 코드를 까보시거나 관련하여서 디테일하게 궁금하신 부분이 있으면 위의 링크들을 확인해보시면 된다.

홈페이지에 가보면 FEMNIST, Shakespeare, Twitter, Celeba, Synthetic Dataset, Reddit 6개의 dataset을 Federated setting에 맞추어서 제공하고 있음을 확인할 수 있다. Image 뿐만 아니라 NLP domain 데이터셋도 다루고 있어 FL 방법론 다루는 연구에서 쓰기 적합해보인다. 본인은 해당 게시글에서 간단하게 FEMNIST 데이터셋에 대해서만 어떻게 데이터셋을 다운로드 및 전처리를 수행하고 dataloader를 작성하였는지만 다루겠다.

 

LEAF 사용법

우선 https://github.com/TalwalkarLab/leaf에 방문하여서 github 레포를 적당히 구경하고 git clone 혹은 fork 하여서 가져오자. 

 

뭐 요런 폴더가 생긴다. 일단 FEMNIST 데이터셋을 준비하자. './data/femnist' 폴더로 이동하자. 해당 폴더에는 다음과 같은 파일들이 들어있다.

 

뭐 별거 없어 보이는데, 실제로 별거 없다. ./preprocess.sh만 실행하면 성공적으로 FEMNIST dataset이 다운로드 되고 적절한 형태로 preprocessing 된다! 간단하지만, 앞서 이야기하였듯 FL dataset의 구성은 자유도가 꽤나 높아서 내가 원하는 setting을 위해서 여러 parameter를 설정해줄 필요가 있다. 'README.md'를 열어서 내용을 읽어주자.

 

기본적인 requirements는 미리미리 설치해두는 치밀함을 가지자. 독자분들은 이미 환경정도는 미리 세팅해두셨으리라 믿는다. preprocess.sh는 나와있는 것 같이 8개의 인자를 받는다. 개인적으로 이 인자들을 잘 이해하고 데이터셋을 구성하는 것이 중요해보인다.

 

-s: sampling 방식

가장 먼저 나오는 -s 는 sampling 방식에 관한 것이다. 선택지는 'iid'와 'niid'가 있다. iid로 설정하는 경우 user들에 대한 data가 iid로 추출된다. distribution이 거의 유사하게 추출될 듯하다. 본인은 non-iid 상황의 dataset을 구성하고 싶으니 niid로 해주었다.

 

--iu: user의 수

다음은 --iu인데 얘는 user의 수다. 기본적으로 LEAF의 paper를 읽어보면, 각 데이터셋의 total device수가 나와있다. FEMNIST의 경우 3,550.

이 device수를 기반으로 비율을 설정하여 본인이 원하는 수의 device 수로 나눌 수 있다. 기본값은 0.01이고 36 (3550*0.01=35.5)개의 file이 생성되는 것으로 확인했다. 적절히 공부하는 수준으로 하고 싶으면 이 값을 줄이고 extensive한 device 수를 경험하고 싶으면 값을 키우면 되시겠다.

 

--sf: sampling할 데이터의 비율

--sf는 sampling할 fraction인데, 꽤나 혼란스러울 수 있는 값이라고 개인적으로 생각한다. 전체 805,263개의 sample이 존재하는데, 이 값을 0.2로 설정하면 805,263=161,052.6 이니 대충 161,503개 정도의 데이터가 추출될 것이라고 생각했는데, 결론적으로 163558개의 데이터가 추출되었다(?!). 우선 대충 '해당 비율은 각 user들을 기준으로 돌아가는 것이 아니라, 전체 데이터셋을 기준으로 돌아가는 비율이다!' 정도로 해놓고 넘어갔다가 다시 돌아오자.

 

-k: 한 user가 최소한 가질 sample의 수

정말 운이 안좋은 경우 한 user가 data를 몇개 못가지는 상황이 벌어질 수 있다. 이런 상황을 고려하여서 최소한 50개 뭐 100개 이렇게 설정해줄 수 있는 값이다.

 

--tf: training set과 test set fraction

학습과 추론 데이터를 구분하는 비율. 길게 설명하지 않겠다.

 

--smplseed, spltseed: sampling과 split을 위한 seed

데이터셋 sampling과 split에 사용되는 random seed를 고정하는 작업. 재현 가능한 연구를 위해 반드시 필요한 부분인데 감사히도 구현해주셨다.

 

--iu와 --sf...

본인은 이 두 친구 때문에 한참 시간을 썼다. --sf와 --iu를 같이 고려해보자, --iu는 최종적으로 나오는 .json 파일의 수를 결정한다. 3,550 명의 유저(=device)에 대해서 0.01이니 .json file이 36개가 나오는 것을 확인했다. 각 .json file에는 평균적으로 21~22명정도의 유저가 들어있다(user 안의 user... 몽중몽 구조... 액자식 구성...). 이를 통해 논문에서 언급하는 device가 README.md의 user임을 알 수 있다. 

 

그리고 모든 sample수를 모두 합하니 164,558개다. 이 개수는 어디서 온 것인가? 일단 전체 데이터셋의 개수는 805,263개, 이 데이터셋은 기본적으로 3,550개의 device로 나뉘어져 있다. 한 유저가 가지는  그런데 우리가 36개의 device를 원한다고 지정했다. 이 순간 어느 수준에서 device개수와 data의 수 사이에 dependency(?) 같은 것이 생긴다. non-iid 환경이다보니 정확한 비율로 그 값을 맞추지는 못하고, 대강 그 값에 가까운 데이터셋을 만들어주는 것이 아닐까 생각된다. 암튼 그렇다. 본인이 멍청해서 이해를 잘 못하는 것 같으니 멋진 독자분이 계시면 댓글로 설명해주시라.

 

몇번 데이터셋으로 연합 학습을 돌려보니 한 떨어지는 .json file 하나(위에서 언급한 device)를 유저 한 명으로 봐줘야 reporting된 성능을 얻을 수 있는 것 같다.

 

Dataset 구성 확인하기

여튼저튼쨋든 json file들이 이쁘게 떨어졌다.

이제 json 하나를 까서 내용물을 살펴보자. 다음과 같은 코드로 내용물을 불러오고 key 값들을 찍어봤다.

import json

file_path = "DATASET PATH"

with open(file_path, 'r') as file:
    data = json.load(file)

print(data.keys()
#dict_keys(['users', 'num_samples', 'user_data'])

key로는 'users', 'num_samples', 'user_data' 세가지가 있다. 각각 접근해보면 user_id, user가 가지는 sample의 수, 각 유저들이 가지는 data 이렇게 구성되어 있다. 야호 이제 다 왔다. 이제 FL용 loader만 구성해보면 되겠다.

 

Dataset 및 Dataloader 구현

우선 본인은 tensorflow와 친하지 않아서 다루는 코드는 Pytorch를 기반으로 함에 양해를 구한다. 

class DataManager():
    def __init__(self, files:list, is_train:bool=True):
        self.files = files
        self.is_train = is_train					   # train을 위한 데이터인 경우 True 아니면 False
        self.users, self.data = [], {}
        if not self.is_train:
            self.global_test_data = {'x':[], 'y':[]}   # server에서 global test를 위한 데이터셋 생성
        
        for idx, file in enumerate(self.files):
            idx = str(idx)
            self.data[idx] = {'x':[], 'y':[]}
            with open(file) as f:
                data = json.load(f)
                self.users.append(idx)

                for user in data['users']:              # 각 유저의 data 저장
                    self.data[idx]['x'] = self.data[idx]['x'] + data['user_data'][user]['x']
                    self.data[idx]['y'] = self.data[idx]['y'] + data['user_data'][user]['y']
                    
                if not self.is_train:                   # test dataset인 경우 global evaluation 위해서 모든 데이터셋 저장
                    for user in data['users']:
                        self.global_test_data['x'] = self.global_test_data['x'] + data['user_data'][user]['x']
                        self.global_test_data['y'] = self.global_test_data['y'] + data['user_data'][user]['y']

코드도 그리 잘 짠 코드는 아니니 본인이 개량하시길 권장한다. 입력으로 모든 json file들의 path를 list 형태로 넘겨주면 작동한다. server에서 global 모델을 평가하기 위해서 나뉘어져 있는 모든 test 데이터셋을 수하는 부분이 들어가 있다. 각 user의 index를 key로하는 dictionay안에 다시 각각 데이터와 레이블을 의미하는 'x', 'y' key를 가지는 dictionary가 들어 있는 형태로 반환된다.

 

class FEMNIST(Dataset):
    def __init__(self, data:dict):
        self.data = data
        self.data['x'] = np.array(data['x'])
        self.data['y'] = np.array(data['y'])
        
    def __getitem__(self, idx):
        # self.X = torch.tensor(self.data['x'][idx,:].reshape(1, 28, 28)).float() # CNN 모델의 경우 이미지의 원 형태로 복원
        self.X = torch.tensor(self.data['x'][idx,:]).float()
        self.Y = torch.tensor(self.data['y'][idx]).long()
        return self.X, self.Y

    def __len__(self):
        return len(self.data['y'])

user 한명을 특정해서 내부의 dictionary를 넘겨주면 이를 tensor로 바꿔서 pytorch dataset으로 구성하는 코드이다. 기본적으로 FEMNIST의 데이터는 28*28 형태의 이미지를 784로 flattening 해두었기 때문에 CNN 모델에 사용하기 위해서는 reshape을 해줄 필요가 있으니 혹시 CNN에 집어넣으실 계획이라면 channel 1을 살려서 (1,28,28) 형태로 reshape하는 코드 주석을 지워서 넣어주시면 된다.

 

def get_files(PATH:str)->dict:
    file_dict = {'train':None, 'test':None}
    
    files = [os.path.join(PATH+'/train', file) \
        for file in os.listdir(PATH+'/train') if file.endswith('.json')]
    files.sort()
    file_dict['train'] = files
    
    files = [os.path.join(PATH+'/test', file) \
        for file in os.listdir(PATH+'/test') if file.endswith('.json')]
    files.sort()
    file_dict['test'] = files
    return file_dict

PATH = cfg.DATAPATH['femnist']
        
file_dict = get_files(PATH)
        
TRAIN_DM = DataManager(file_dict['train'], is_train=True)
TEST_DM = DataManager(file_dict['test'], is_train=False)
DM_dict = {'train':TRAIN_DM,
            'test':TEST_DM}
print("DATA READY")

뭐 이런식으로 사용하시면 된다. 본인은 configuration을 따로 빼서 dataset 별로 path를 지정해두고, get_files()라는 함수에 넘겨서 trainset과 testset의 파일들을 모두 가지는 dictionary를 가지고 오는 방식을 택했다. 뭐 그리고 그냥 아까만든 DataManager를 통해서 각 user별로 parsing한 이후 나중에 client를 정의할 때 그 client의 data만 FEMNIST class로 dataset을 할당해주면 된다.

 

마무리하며...

여러모로 괜찮은 데이터셋이긴 하나 사용법 같은 것들이 제대로 정리된 글을 찾기 힘들어서 그냥 맨땅에 헤딩하면서 대충 정리해봤다. 늘 FL에서 논문마다 다들 자기 마음대로 dataset 구성하고 구현도 조금씩 다르게 하고 그래서 fair한 비교가 어려웠는데 LEAF같은 Benchmark가 있다는 것은 좋은 일이라고 생각한다. 그래도 MNIST나 CIFAR-10, 100 같은 dataset을 non-iid로 다시 FL 용으로 예쁘게 정리해주는 레포는 아직 못찾아봐서 이 부분도 시간이 나면 한 번 구현해둬야겠다.

 

그리고 (1)이라고 해둔 이후는 FedAvg 구현한 코드에 대해서 정리를 해둬야겠다는 생각이 들어서인데 언제 다 정리해서 올릴지, 사실 못할지도 모르겠지만 일단 가능성은 열어두려고 그런거다. 아무도 그러지 않겠지만 기다리지 말라는 뜻