78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
from utils.fed_util import init_model
|
|
from fed_algo_cs.server_base import test
|
|
import os
|
|
import yaml
|
|
from utils.args import args_parser # args parser
|
|
from fed_algo_cs.client_base import FedYoloClient # FedYoloClient
|
|
from fed_algo_cs.server_base import FedYoloServer # FedYoloServer
|
|
from utils import Dataset # Dataset
|
|
|
|
if __name__ == "__main__":
|
|
# model structure test
|
|
model = init_model("yolo_v11_n", num_classes=1)
|
|
|
|
with open("model.txt", "w", encoding="utf-8") as f:
|
|
print(model, file=f)
|
|
|
|
# loop over model key and values
|
|
with open("model_key_value.txt", "w", encoding="utf-8") as f:
|
|
for k, v in model.state_dict().items():
|
|
print(f"{k}: {v.shape}", file=f)
|
|
|
|
# test agg function
|
|
from fed_algo_cs.server_base import FedYoloServer
|
|
import torch
|
|
import yaml
|
|
|
|
with open("./config/coco128_cfg.yaml", "r", encoding="utf-8") as f:
|
|
cfg = yaml.safe_load(f)
|
|
params = dict(cfg)
|
|
|
|
server = FedYoloServer(client_list=["c1", "c2", "c3"], model_name="yolo_v11_n", params=params)
|
|
state1 = {k: torch.ones_like(v) for k, v in server.model.state_dict().items()}
|
|
state2 = {k: torch.ones_like(v) * 2 for k, v in server.model.state_dict().items()}
|
|
state3 = {k: torch.ones_like(v) * 3 for k, v in server.model.state_dict().items()}
|
|
server.rec("c1", state1, n_data=20, loss=0.1)
|
|
server.rec("c2", state2, n_data=30, loss=0.2)
|
|
server.rec("c3", state3, n_data=50, loss=0.3)
|
|
server.select_clients(connection_ratio=1.0)
|
|
model_state, avg_loss, n_data = server.agg()
|
|
with open("agg_model.txt", "w", encoding="utf-8") as f:
|
|
for k, v in model_state.items():
|
|
print(f"{k}: {v.float().mean()}", file=f)
|
|
print(f"avg_loss: {avg_loss}, n_data: {n_data}")
|
|
|
|
# test single client training (should be the same as standalone training)
|
|
args = args_parser()
|
|
with open(args.config, "r", encoding="utf-8") as f:
|
|
cfg = yaml.safe_load(f)
|
|
# params = dict(cfg)
|
|
client = FedYoloClient(name="c1", params=cfg, model_name="yolo_v11_n")
|
|
|
|
filenames = []
|
|
data_dir = "/home/image1325/ssd1/dataset/COCO128"
|
|
with open(f"{data_dir}/train.txt") as f:
|
|
for filename in f.readlines():
|
|
filename = os.path.basename(filename.rstrip())
|
|
filenames.append(f"{data_dir}/images/train2017/" + filename)
|
|
|
|
client.load_trainset(train_dataset=filenames)
|
|
model_state, n_data, avg_loss = client.train(args=args)
|
|
|
|
model = init_model("yolo_v11_n", num_classes=80)
|
|
model.load_state_dict(model_state)
|
|
|
|
valset = Dataset(
|
|
filenames=filenames,
|
|
input_size=640,
|
|
params=cfg,
|
|
augment=False,
|
|
)
|
|
if valset is not None:
|
|
precision, recall, map50, map = test(valset=valset, params=cfg, model=model, batch_size=128)
|
|
print(
|
|
f"precision: {precision}, recall: {recall}, map50: {map50}, map: {map}, loss: {avg_loss}, n_data: {n_data}"
|
|
)
|
|
else:
|
|
raise ValueError("valset is None, please provide a valid valset in config file.")
|