Spaces:
Sleeping
Sleeping
File size: 2,006 Bytes
3e25ded |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
"""huggingface_example: A Flower / Hugging Face app."""
import warnings
import torch
from flwr.client import Client, ClientApp, NumPyClient
from flwr.common import Context
from transformers import logging
from .task import (
train,
test,
load_data,
set_params,
get_params,
get_model,
)
warnings.filterwarnings("ignore", category=FutureWarning)
# To mute warnings reminding that we need to train the model to a downstream task
# This is something this example does.
logging.set_verbosity_error()
# Flower client
class IMDBClient(NumPyClient):
def __init__(self, model_name, trainloader, testloader) -> None:
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.trainloader = trainloader
self.testloader = testloader
self.net = get_model(model_name)
self.net.to(self.device)
def fit(self, parameters, config) -> tuple[list, int, dict]:
set_params(self.net, parameters)
train(self.net, self.trainloader, epochs=1, device=self.device)
return get_params(self.net), len(self.trainloader), {}
def evaluate(self, parameters, config) -> tuple[float, int, dict[str, float]]:
set_params(self.net, parameters)
loss, accuracy = test(self.net, self.testloader, device=self.device)
return float(loss), len(self.testloader), {"accuracy": float(accuracy)}
def client_fn(context: Context) -> Client:
"""Construct a Client that will be run in a ClientApp."""
# Read the node_config to fetch data partition associated to this node
partition_id = context.node_config["partition-id"]
num_partitions = context.node_config["num-partitions"]
# Read the run config to get settings to configure the Client
model_name = context.run_config["model-name"]
trainloader, testloader = load_data(partition_id, num_partitions, model_name)
return IMDBClient(model_name, trainloader, testloader).to_client()
app = ClientApp(client_fn=client_fn)
|