Quickstart JAX

이 튜토리얼에서는 Flower를 사용하여 기존 JAX 워크로드의 연합 버전을 구축하는 방법을 보여드립니다. JAX를 사용해 scikit-learn 데이터 세트에서 선형 회귀 모델을 훈련하고 있습니다. 예제는 ‘파이토치 - Centralized에서 Federated으로 <https://github.com/adap/flower/blob/main/examples/pytorch-from-centralized-to-federated>`_ 워크스루와 유사하게 구성하겠습니다. 먼저, JAX를 사용한 선형 회귀 튜토리얼`을 기반으로 centralized 학습 접근 방식을 구축합니다. 그런 다음 centralized 트레이닝 코드를 기반으로 federated 방식으로 트레이닝을 실행합니다.

JAX 예제 빌드를 시작하기 전에 jax, jaxlib, scikit-learn, flwr 패키지를 설치해야 합니다:

$ pip install jax jaxlib scikit-learn flwr

JAX를 사용한 선형 회귀

먼저 선형 회귀 모델을 기반으로 하는 중앙 집중식 훈련 코드에 대한 간략한 설명부터 시작하겠습니다. 더 자세한 설명을 원하시면 공식 `JAX 문서 <https://jax.readthedocs.io/>`_를 참조하세요.

전통적인(중앙 집중식) 선형 회귀 훈련에 필요한 모든 구성 요소가 포함된 jax_training.py`라는 새 파일을 생성해 보겠습니다. 먼저, JAX 패키지인 :code:`jax`와 :code:`jaxlib`를 가져와야 합니다. 또한 데이터 세트에 :code:`make_regression`을 사용하고 데이터 세트를 학습 및 테스트 세트로 분할하기 위해 :code:`train_test_split`을 사용하므로 :code:`sklearn`을 가져와야 합니다. 연합 학습을 위해 아직 :code:`flwr 패키지를 가져오지 않은 것을 볼 수 있습니다. 이 작업은 나중에 수행됩니다.

from typing import Dict, List, Tuple, Callable
import jax
import jax.numpy as jnp
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

key = jax.random.PRNGKey(0)

code:load_data() 함수는 앞서 언급한 트레이닝 및 테스트 세트를 로드합니다.

def load_data() -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], List[np.ndarray]]:
    # create our dataset and start with similar datasets for different clients
    X, y = make_regression(n_features=3, random_state=0)
    X, X_test, y, y_test = train_test_split(X, y)
    return X, y, X_test, y_test

모델 아키텍처(매우 간단한 선형 회귀 모델)는 :code:`load_model()`에 정의되어 있습니다.

def load_model(model_shape) -> Dict:
    # model weights
    params = {
        'b' : jax.random.uniform(key),
        'w' : jax.random.uniform(key, model_shape)
    }
    return params

이제 훈련 집합을 반복하고 각 훈련 예제 배치에 대해 손실을 측정하는(함수 loss_fn()) 훈련(함수 train())을 정의해야 합니다. JAX는 grad() 함수(main() 함수에 정의되고 :code:`train()`에서 호출됨)로 파생물을 취하므로 손실 함수는 분리되어 있습니다.

def loss_fn(params, X, y) -> Callable:
    err = jnp.dot(X, params['w']) + params['b'] - y
    return jnp.mean(jnp.square(err))  # mse

def train(params, grad_fn, X, y) -> Tuple[np.array, float, int]:
    num_examples = X.shape[0]
    for epochs in range(10):
        grads = grad_fn(params, X, y)
        params = jax.tree_multimap(lambda p, g: p - 0.05 * g, params, grads)
        loss = loss_fn(params,X, y)
        # if epochs % 10 == 9:
        #     print(f'For Epoch {epochs} loss {loss}')
    return params, loss, num_examples

모델의 평가는 evaluation() 함수에 정의되어 있습니다. 이 함수는 모든 테스트 예제를 가져와 선형 회귀 모델의 손실을 측정합니다.

def evaluation(params, grad_fn, X_test, y_test) -> Tuple[float, int]:
    num_examples = X_test.shape[0]
    err_test = loss_fn(params, X_test, y_test)
    loss_test = jnp.mean(jnp.square(err_test))
    # print(f'Test loss {loss_test}')
    return loss_test, num_examples

데이터 로딩, 모델 아키텍처, 훈련 및 평가를 정의했으므로 이제 모든 것을 종합하여 JAX를 사용 모델을 훈련할 수 있습니다. 이미 언급했듯이 jax.grad() 함수는 :code:`main()`에 정의되어 :code:`train()`에 전달됩니다.

def main():
    X, y, X_test, y_test = load_data()
    model_shape = X.shape[1:]
    grad_fn = jax.grad(loss_fn)
    print("Model Shape", model_shape)
    params = load_model(model_shape)
    params, loss, num_examples = train(params, grad_fn, X, y)
    evaluation(params, grad_fn, X_test, y_test)


if __name__ == "__main__":
    main()

이제 (중앙 집중식) JAX 선형 회귀 워크로드를 실행할 수 있습니다:

python3 jax_training.py

지금까지는 JAX를 사용해 본 적이 있다면 이 모든 것이 상당히 익숙해 보일 것입니다. 다음 단계로 넘어가서 우리가 구축한 것을 사용하여 하나의 서버와 두 개의 클라이언트로 구성된 간단한 연합 학습 시스템을 만들어 보겠습니다.

JAX와 Flower의 만남

기존 워크로드를 연합하는 개념은 항상 동일하고 이해하기 쉽습니다. 서버*를 시작한 다음 *서버*에 연결된 *클라이언트*에 대해 :code:`jax_training.py`의 코드를 사용해야 합니다. *서버*는 모델 파라미터를 클라이언트로 전송합니다. 클라이언트는 학습을 실행하고 파라미터를 업데이트합니다. 업데이트된 파라미터는 *서버*로 다시 전송되며, 수신된 모든 파라미터 업데이트의 평균을 구합니다. 이는 연합 학습 프로세스의 한 라운드를 설명하며, 이 과정을 여러 라운드에 걸쳐 반복합니다.

이 예제는 하나의 *서버*와 두 개의 *클라이언트*로 구성됩니다. 먼저 server.py`를 설정해 보겠습니다. *server*는 Flower 패키지 :code:`flwr`를 가져와야 합니다. 다음으로, :code:`start_server 함수를 사용하여 서버를 시작하고 세 차례의 연합 학습을 수행하도록 지시합니다.

import flwr as fl

if __name__ == "__main__":
    fl.server.start_server(server_address="0.0.0.0:8080", config=fl.server.ServerConfig(num_rounds=3))

이미 *서버*를 시작할 수 있습니다:

python3 server.py

마지막으로, client.py`에서 *client* 로직을 정의하고 :code:`jax_training.py`에서 이전에 정의한 JAX 교육을 기반으로 빌드합니다. *클라이언트*는 :code:`flwr`을 가져와야 하며, JAX 모델의 파라미터를 업데이트하기 위해 :code:`jax 및 :code:`jaxlib`도 가져와야 합니다:

from typing import Dict, List, Callable, Tuple

import flwr as fl
import numpy as np
import jax
import jax.numpy as jnp

import jax_training

Flower *클라이언트*를 구현한다는 것은 기본적으로 flwr.client.Client 또는 :code:`flwr.client.NumPyClient`의 서브클래스를 구현하는 것을 의미합니다. 구현은 :code:`flwr.client.NumPyClient`를 기반으로 하며, 이를 :code:`FlowerClient`라고 부를 것입니다. :code:`NumPyClient`는 필요한 일부 보일러플레이를 피할 수 있기 때문에 NumPy 상호 운용성이 좋은 프레임워크(예: JAX)를 사용하는 경우 :code:`Client`보다 구현하기가 약간 더 쉽습니다. code:`FlowerClient`는 모델 매개변수를 가져오거나 설정하는 메서드 2개, 모델 학습을 위한 메서드 1개, 모델 테스트를 위한 메서드 1개 등 총 4개의 메서드를 구현해야 합니다:

1. set_parameters (선택사항)

   • 서버에서 수신한 로컬 모델의 모델 파라미터를 설정합니다

   • 매개 변수를 NumPy :code:`ndarray`로 변환

   • (신경망 레이어 목록으로 생각하면 됩니다) NumPy :code:`ndarray`로 받은 모델 파라미터 목록에 대해 반복합니다

2. get_parameters

   • 모델 매개변수를 가져와서 NumPy :code:`ndarray`의 목록으로 반환합니다(이는 :code:`flwr.client.NumPyClient`가 기대하는 바와 같습니다)

3. fit

   • 서버에서 받은 파라미터로 로컬 모델의 파라미터를 업데이트합니다

   • 로컬 훈련 세트에서 모델을 훈련합니다

   • 업데이트된 로컬 모델 파라미터를 가져와 서버로 반환합니다

4. evaluate

   • 서버에서 받은 파라미터로 로컬 모델의 파라미터를 업데이트합니다

   • 로컬 테스트 세트에서 업데이트된 모델을 평가합니다

   • 로컬 손실을 서버로 반환합니다

어려운 부분은 JAX 모델 매개변수를 :code:`DeviceArray`에서 :code:`NumPy ndarray`로 변환하여 `NumPyClient`와 호환되도록 하는 것입니다.

두 개의 NumPyClient 메서드인 fit`과 :code:`evaluate`는 이전에 :code:`jax_training.py`에 정의된 함수 :code:`train()`과 :code:`evaluate()`를 사용합니다. 따라서 여기서 우리가 실제로 하는 일은 이미 정의된 함수 중 훈련과 평가를 위해 호출할 함수를 :code:`NumPyClient 서브클래스를 통해 Flower에게 알려주는 것입니다. 전달되는 데이터 유형을 더 잘 이해할 수 있도록 유형 type annotation을 포함했습니다.

class FlowerClient(fl.client.NumPyClient):
    """Flower client implementing using linear regression and JAX."""

    def __init__(
        self,
        params: Dict,
        grad_fn: Callable,
        train_x: List[np.ndarray],
        train_y: List[np.ndarray],
        test_x: List[np.ndarray],
        test_y: List[np.ndarray],
    ) -> None:
        self.params= params
        self.grad_fn = grad_fn
        self.train_x = train_x
        self.train_y = train_y
        self.test_x = test_x
        self.test_y = test_y

    def get_parameters(self, config) -> Dict:
        # Return model parameters as a list of NumPy ndarrays
        parameter_value = []
        for _, val in self.params.items():
            parameter_value.append(np.array(val))
        return parameter_value

    def set_parameters(self, parameters: List[np.ndarray]) -> Dict:
        # Collect model parameters and update the parameters of the local model
        value=jnp.ndarray
        params_item = list(zip(self.params.keys(),parameters))
        for item in params_item:
            key = item[0]
            value = item[1]
            self.params[key] = value
        return self.params


    def fit(
        self, parameters: List[np.ndarray], config: Dict
    ) -> Tuple[List[np.ndarray], int, Dict]:
        # Set model parameters, train model, return updated model parameters
        print("Start local training")
        self.params = self.set_parameters(parameters)
        self.params, loss, num_examples = jax_training.train(self.params, self.grad_fn, self.train_x, self.train_y)
        results = {"loss": float(loss)}
        print("Training results", results)
        return self.get_parameters(config={}), num_examples, results

    def evaluate(
        self, parameters: List[np.ndarray], config: Dict
    ) -> Tuple[float, int, Dict]:
        # Set model parameters, evaluate the model on a local test dataset, return result
        print("Start evaluation")
        self.params = self.set_parameters(parameters)
        loss, num_examples = jax_training.evaluation(self.params,self.grad_fn, self.test_x, self.test_y)
        print("Evaluation accuracy & loss", loss)
        return (
            float(loss),
            num_examples,
            {"loss": float(loss)},
        )

연합 프로세스를 정의했으면 이제 실행할 수 있습니다.

def main() -> None:
    """Load data, start MNISTClient."""

    # Load data
    train_x, train_y, test_x, test_y = jax_training.load_data()
    grad_fn = jax.grad(jax_training.loss_fn)

    # Load model (from centralized training) and initialize parameters
    model_shape = train_x.shape[1:]
    params = jax_training.load_model(model_shape)

    # Start Flower client
    client = FlowerClient(params, grad_fn, train_x, train_y, test_x, test_y)
    fl.client.start_client(server_address="0.0.0.0:8080", client=client.to_client())

if __name__ == "__main__":
    main()

여기까지입니다. 이제 두 개의 터미널 창을 추가로 열고 다음을 실행할 수 있습니다

python3 client.py

를 입력하고(그 전에 서버가 계속 실행 중인지 확인하세요) 두 클라이언트에서 연합 학습을 실행하는 JAX 프로젝트를 확인합니다. 축하합니다!

다음 단계

이 예제의 소스 코드는 시간이 지남에 따라 개선되었으며 여기에서 확인할 수 있습니다: ‘Quickstart JAX <https://github.com/adap/flower/blob/main/examples/quickstart-jax>`_. 두 클라이언트가 동일한 데이터 세트를 로드하기 때문에 이 예제는 다소 단순화되어 있습니다.

이제 이 주제를 더 자세히 살펴볼 준비가 되었습니다. 더 정교한 모델을 사용하거나 다른 데이터 집합을 사용해 보는 것은 어떨까요? 클라이언트를 더 추가하는 것은 어떨까요?