How can I apply federated learning to YOLOv8 model?

Question

I made a model for detecting seven objects that detect people using YOLOv8 and saved it as .pt.

The dataset is a jpg and txt file consisting of images and labels.

How can I apply Federated Learning to YOLOv8n in a Google colab environment?

I'm a beginner in CV field.

import openfl.native as fx
import torch
from torch.optim import SGD
import torchvision.transforms as T
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from ultralytics import YOLO

# 로컬 모델 파일 경로
model_files = [
    "/content/drive/MyDrive/runs/detect/local_1_train_result/weights/best.pt",
    "/content/drive/MyDrive/runs/detect/local_2_train_result/weights/best.pt",
    "/content/drive/MyDrive/runs/detect/local_3_train_result/weights/best.pt",
    "/content/drive/MyDrive/runs/detect/local_4_train_result/weights/best.pt",
    "/content/drive/MyDrive/runs/detect/local_5_train_result/weights/best.pt",
    "/content/drive/MyDrive/runs/detect/local_6_train_result/weights/best.pt",
    "/content/drive/MyDrive/runs/detect/local_7_train_result/weights/best.pt"
]

# 데이터 경로
data_path = '/content/drive/MyDrive/local_1_final_all/train'

# YOLOv5 모델 로드
def load_yolov8_model():
    model = YOLO(model='yolov8n.pt')
    return model

# 로컬 모델 로드
local_models = []
global_model = load_yolov8_model()

for i, model_file in enumerate(model_files):
    local_model = load_yolov8_model()
    local_model.load_state_dict(torch.load(model_file))
    local_models.append(local_model)
    global_model.features[i] = local_model.features[i]

# 데이터 로딩 및 전처리
transform = T.Compose([T.ToTensor()])

def custom_collate_fn(batch):
    return tuple(zip(*batch))

dataset = ImageFolder(data_path, transform=transform)
data_loader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=custom_collate_fn)

# OpenFL 초기화
fx.init('torch', 'openfl', local_model=global_model)

# 연합 학습을 위한 옵티마이저 설정
optimizer = SGD(global_model.parameters(), lr=0.001)

# 연합 학습 실행
fx.federated_averaging(
    model=global_model,
    data={'data': dataset},
    server_optimizer=optimizer,
    aggregation='mean',
    rounds=10
)

# 연합 모델 저장
torch.save(global_model.state_dict(), "/content/drive/MyDrive/federated_model.pt")

I tried the above, but I failed all of them.