联邦学习：在不移动数据的情况下跨10家医院训练ML模型 | WellAlly康心伴

pip install syft numpy pandas torch torchvision

# src/data_simulation.py
import pandas as pd
import numpy as np

def generate_hospital_data(num_samples=100, hospital_id=1):
    """
    为单个医院生成合成健康数据。
    """
    np.random.seed(42 + hospital_id)
    data = {
        'age': np.random.randint(20, 80, num_samples),
        'bmi': np.random.uniform(18.5, 40, num_samples),
        'blood_pressure': np.random.randint(80, 180, num_samples),
        'has_condition': np.random.randint(0, 2, num_samples)
    }
    return pd.DataFrame(data)

# 模拟两个医院的数据
hospital_a_data = generate_hospital_data(num_samples=150, hospital_id=1)
hospital_b_data = generate_hospital_data(num_samples=200, hospital_id=2)

print("医院A数据：")
print(hospital_a_data.head())
print("\n医院B数据：")
print(hospital_b_data.head())

# src/federated_setup.py
import torch
import syft as sy
from data_simulation import hospital_a_data, hospital_b_data
from sklearn.model_selection import train_test_split

# 创建hook以扩展PyTorch的联邦学习能力
hook = sy.TorchHook(torch)

# 为我们的医院创建虚拟工作者
hospital_a = sy.VirtualWorker(hook, id="hospital-a")
hospital_b = sy.VirtualWorker(hook, id="hospital-b")

# 准备训练数据
def prepare_data(df):
    X = torch.tensor(df[['age', 'bmi', 'blood_pressure']].values).float()
    y = torch.tensor(df['has_condition'].values).float().unsqueeze(1)
    return X, y

X_a, y_a = prepare_data(hospital_a_data)
X_b, y_b = prepare_data(hospital_b_data)

# 将数据发送到各自的虚拟工作者
X_a_ptr = X_a.send(hospital_a)
y_a_ptr = y_a.send(hospital_a)
X_b_ptr = X_b.send(hospital_b)
y_b_ptr = y_b.send(hospital_b)

print("数据已发送到虚拟工作者：")
print("医院A数据指针：", X_a_ptr)
print("医院B数据指针：", X_b_ptr)

# src/model_and_training.py
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的逻辑回归模型
class HealthClassifier(nn.Module):
    def __init__(self, input_features=3):
        super(HealthClassifier, self).__init__()
        self.linear = nn.Linear(input_features, 1)

    def forward(self, x):
        return torch.sigmoid(self.linear(x))

# 训练逻辑
def train_on_client(model, optimizer, X_train, y_train, epochs=10):
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = nn.BCELoss()(outputs, y_train)
        loss.backward()
        optimizer.step()
    return model

# src/main.py
import torch
import syft as sy
from model_and_training import HealthClassifier, train_on_client
from federated_setup import hospital_a, hospital_b, X_a_ptr, y_a_ptr, X_b_ptr, y_b_ptr

# 初始化全局模型
global_model = HealthClassifier()

# 定义训练参数
learning_rate = 0.01
num_rounds = 5

for round in range(num_rounds):
    print(f"\n--- 第 {round + 1} 轮 ---")

    # 1. 将全局模型发送到每个医院
    model_a = global_model.copy().send(hospital_a)
    model_b = global_model.copy().send(hospital_b)

    # 2. 在本地数据上训练模型
    optimizer_a = optim.SGD(model_a.parameters(), lr=learning_rate)
    optimizer_b = optim.SGD(model_b.parameters(), lr=learning_rate)

    print("在医院A上训练...")
    trained_model_a = train_on_client(model_a, optimizer_a, X_a_ptr, y_a_ptr)
    print("在医院B上训练...")
    trained_model_b = train_on_client(model_b, optimizer_b, X_b_ptr, y_b_ptr)

    # 3. 从医院获取更新后的模型
    trained_model_a.get()
    trained_model_b.get()

    # 4. 平均模型权重以更新全局模型
    with torch.no_grad():
        global_model.linear.weight.set_(((trained_model_a.linear.weight.data + trained_model_b.linear.weight.data) / 2))
        global_model.linear.bias.set_(((trained_model_a.linear.bias.data + trained_model_b.linear.bias.data) / 2))

    print("全局模型已更新。")

print("\n联邦训练完成！")
print("最终全局模型权重：", global_model.state_dict())