更新联邦学习

FED.py

@@ -12,7 +12,7 @@ from model.repvit import repvit_m1_1
 from model.mobilenetv3 import MobileNetV3
 
 # 配置参数
-NUM_CLIENTS = 4
+NUM_CLIENTS = 2
 NUM_ROUNDS = 3
 CLIENT_EPOCHS = 5
 BATCH_SIZE = 32
@@ -22,25 +22,27 @@ TEMP = 2.0  # 蒸馏温度
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # 数据准备
-def prepare_data(num_clients):
+import os
+from torchvision.datasets import ImageFolder
+
+def prepare_data():
     transform = transforms.Compose([
-            transforms.Resize((224, 224)),  # 将图像调整为 224x224
-            transforms.Grayscale(num_output_channels=3),
-            transforms.ToTensor()
-        ])
-    train_set = datasets.MNIST("./data", train=True, download=True, transform=transform)
+        transforms.Resize((224, 224)),
+        transforms.ToTensor()
+    ])
     
-    # 非IID数据划分（每个客户端2个类别）
-    client_data = {i: [] for i in range(num_clients)}
-    labels = train_set.targets.numpy()
-    for label in range(10):
-        label_idx = np.where(labels == label)[0]
-        np.random.shuffle(label_idx)
-        split = np.array_split(label_idx, num_clients//2)
-        for i, idx in enumerate(split):
-            client_data[i*2 + label%2].extend(idx)
+    # Load datasets
+    dataset_A = ImageFolder(root='./dataset_A/train', transform=transform)
+    dataset_B = ImageFolder(root='./dataset_B/train', transform=transform)
+    dataset_C = ImageFolder(root='./dataset_C/train', transform=transform)
     
-    return [Subset(train_set, ids) for ids in client_data.values()]
+    # Assign datasets to clients
+    client_datasets = [dataset_B, dataset_C]
+    
+    # Server dataset (A) for public updates
+    public_loader = DataLoader(dataset_A, batch_size=BATCH_SIZE, shuffle=True)
+    
+    return client_datasets, public_loader
 
 # 客户端训练函数
 def client_train(client_model, server_model, dataset):
@@ -189,63 +191,47 @@ def test_model(model, test_loader):
 
 # 主训练流程
 def main():
-    # 初始化模型
+    # Initialize models
     global_server_model = repvit_m1_1(num_classes=10).to(device)
     client_models = [MobileNetV3(n_class=10).to(device) for _ in range(NUM_CLIENTS)]
-
-    round_progress = tqdm(total=NUM_ROUNDS, desc="Federated Rounds", unit="round")
     
-    # 准备数据
-    client_datasets = prepare_data(NUM_CLIENTS)
-    public_loader = DataLoader(
-        datasets.MNIST("./data", train=False, download=True, 
-                      transform= transforms.Compose([
-                        transforms.Resize((224, 224)),  # 将图像调整为 224x224
-                        transforms.Grayscale(num_output_channels=3),
-                        transforms.ToTensor()           # 将图像转换为张量
-                    ])),
-        batch_size=100, shuffle=True)
+    # Prepare data
+    client_datasets, public_loader = prepare_data()
     
-    test_dataset = datasets.MNIST(
-        "./data", 
-        train=False,
-        transform= transforms.Compose([
-                        transforms.Resize((224, 224)),  # 将图像调整为 224x224
-                        transforms.Grayscale(num_output_channels=3),
-                        transforms.ToTensor()           # 将图像转换为张量
-                    ])
-    )
+    # Test dataset (using dataset A's test set for simplicity)
+    test_dataset = ImageFolder(root='./dataset_A/test', transform=transform)
     test_loader = DataLoader(test_dataset, batch_size=100, shuffle=False)
     
+    round_progress = tqdm(total=NUM_ROUNDS, desc="Federated Rounds", unit="round")
+    
     for round in range(NUM_ROUNDS):
         print(f"\n{'#'*50}")
         print(f"Federated Round {round+1}/{NUM_ROUNDS}")
         print(f"{'#'*50}")
         
-        # 客户端选择
+        # Client selection (only 2 clients)
         selected_clients = np.random.choice(NUM_CLIENTS, 2, replace=False)
         print(f"Selected Clients: {selected_clients}")
         
-        # 客户端本地训练
+        # Client local training
         client_params = []
         for cid in selected_clients:
             print(f"\nTraining Client {cid}")
             local_model = copy.deepcopy(client_models[cid])
             local_model.load_state_dict(client_models[cid].state_dict())
-            
             updated_params = client_train(local_model, global_server_model, client_datasets[cid])
             client_params.append(updated_params)
         
-        # 模型聚合
+        # Model aggregation
         global_client_params = aggregate(client_params)
         for model in client_models:
             model.load_state_dict(global_client_params)
         
-        # 服务器知识更新
+        # Server knowledge update
         print("\nServer Updating...")
         server_update(global_server_model, client_models, public_loader)
         
-        # 测试模型性能
+        # Test model performance
         server_acc = test_model(global_server_model, test_loader)
         client_acc = test_model(client_models[0], test_loader)
         print(f"\nRound {round+1} Performance:")
@@ -253,25 +239,22 @@ def main():
         print(f"Client Model Accuracy: {client_acc:.2f}%")
         
         round_progress.update(1)
-        
         print(f"Round {round+1} completed")
     
     print("Training completed!")
-
-     # 保存训练好的模型
+    
+    # Save trained models
     torch.save(global_server_model.state_dict(), "server_model.pth")
     torch.save(client_models[0].state_dict(), "client_model.pth")
     print("Models saved successfully.")
-
-    # 创建测试数据加载器
-
-    # 测试服务器模型
+    
+    # Test server model
     server_model = repvit_m1_1(num_classes=10).to(device)
     server_model.load_state_dict(torch.load("server_model.pth"))
     server_acc = test_model(server_model, test_loader)
     print(f"Server Model Test Accuracy: {server_acc:.2f}%")
-
-    # 测试客户端模型
+    
+    # Test client model
     client_model = MobileNetV3(n_class=10).to(device)
     client_model.load_state_dict(torch.load("client_model.pth"))
     client_acc = test_model(client_model, test_loader)

README.md

@@ -0,0 +1,4 @@
+# 1.相关知识
+```
+    https://github.com/Hao840/OFAKD
+```

config.py

@@ -8,7 +8,7 @@ class Config:
     
     # 训练参数
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    batch_size = 32
+    batch_size = 128
     epochs = 150
     learning_rate = 0.001
     save_path = "checkpoints/best_model.pth"
@@ -22,4 +22,6 @@ class Config:
     checkpoint_path = "checkpoints/last_checkpoint.pth"
     output_path = "runs/"
 
+    cache = 'RAM'
+
 config = Config()

data_loader.py

@@ -1,31 +1,68 @@
 import os
+from logger import logger
 from PIL import Image
-import numpy as np
 import torch
-from torchvision import datasets, transforms
 from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
 
-class ClassifyDataset(Dataset):
-    def __init__(self, data_dir,transforms = None):
-        self.data_dir = data_dir
-        # Assume the dataset is structured with subdirectories for each class
-        self.transform = transforms
-        self.dataset = datasets.ImageFolder(self.data_dir, transform=self.transform)
-        self.image_size = (3, 224, 224)
+class ImageClassificationDataset(Dataset):
+    def __init__(self, root_dir, transform=None,Cache=False):
+        self.root_dir = root_dir
+        self.transform = transform
+        self.classes = sorted(os.listdir(root_dir))
+        self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)}
+        self.image_paths = []
+        self.image = []
+        self.labels = []
+        self.Cache = Cache
+
+        logger.log("info",
+            "init the dataloader"
+        )
+
+        for cls_name in self.classes:
+            cls_dir = os.path.join(root_dir, cls_name)
+            for img_name in os.listdir(cls_dir):
+                try:
+                    img_path = os.path.join(cls_dir, img_name)
+                    imgs = Image.open(img_path).convert('RGB')
+                    if Cache == 'RAM':
+                        if self.transform:
+                            imgs = self.transform(imgs)
+                        self.image.append(imgs)
+                    else:
+                        self.image_paths.append(img_path)
+                    self.labels.append(self.class_to_idx[cls_name])
+                except:
+                    logger.log("info",
+                        "read image error " + 
+                        img_path
+                    )
 
     def __len__(self):
-        return len(self.dataset)
+        return len(self.labels)
 
     def __getitem__(self, idx):
-        try:
-            image, label = self.dataset[idx]
-            return image, label
-        except Exception as e:
-            black_image = np.zeros((224, 224, 3), dtype=np.uint8)
-            return self.transform(Image.fromarray(black_image)), 0  # -1 作为默认标签
-    
-def create_data_loaders(data_dir,batch_size=64):
-    # Define transformations for training data augmentation and normalization
+        label = self.labels[idx]
+        if self.Cache == 'RAM':
+            image = self.image[idx]
+        else: 
+            img_path = self.image_paths[idx]
+            image = Image.open(img_path).convert('RGB')
+            if self.transform:
+                image = self.transform(image)
+
+        return image, label
+
+def get_data_loader(root_dir, batch_size=64, num_workers=4, pin_memory=True,Cache=False):
+    # Define the transform for the training data and for the validation data
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),  # Resize images to 224x224
+        transforms.ToTensor(),          # Convert PIL Image to Tensor
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # Normalize the images
+    ])
+
+     # Define transformations for training data augmentation and normalization
     train_transforms = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
@@ -40,17 +77,38 @@ def create_data_loaders(data_dir,batch_size=64):
     ])
 
     # Load the datasets with ImageFolder
-    train_dir = data_dir + '/train'
-    valid_dir = data_dir + '/val'
-    test_dir = data_dir + '/test'
+    train_dir = root_dir + '/train'
+    valid_dir = root_dir + '/val'
+    test_dir = root_dir + '/test'
 
-    train_data = ClassifyDataset(train_dir, transforms=train_transforms)
-    valid_data = ClassifyDataset(valid_dir, transforms=valid_test_transforms)
-    test_data = ClassifyDataset(test_dir, transforms=valid_test_transforms)
+    train_data = ImageClassificationDataset(train_dir, transform=train_transforms,Cache=Cache)
+    valid_data = ImageClassificationDataset(valid_dir, transform=valid_test_transforms,Cache=Cache)
+    test_data = ImageClassificationDataset(test_dir, transform=valid_test_transforms,Cache=Cache)
 
-    # Create the DataLoaders with batch size 64
-    train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
-    valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=batch_size)
-    test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size)
 
-    return train_loader, valid_loader, test_loader
+    # Create the data loader
+    train_loader = DataLoader(
+        train_data,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=pin_memory
+    )
+
+     # Create the data loader
+    valid_loader = DataLoader(
+        valid_data,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=pin_memory
+    )
+
+     # Create the data loader
+    test_loader = DataLoader(
+        test_data,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=pin_memory
+    )
+
+    return train_loader, valid_loader, test_loader

dataset/recover.py

@@ -106,15 +106,15 @@ def process_images(input_folder, background_image_path, output_base):
     递归处理所有子文件夹并保持目录结构
     """
     # 预处理背景路径（只需执行一次）
-    if os.path.isfile(background_image_path):
-        background_paths = [background_image_path]
-    else:
-        valid_ext = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
-        background_paths = [
-            os.path.join(background_image_path, f) 
-            for f in os.listdir(background_image_path)
-            if os.path.splitext(f)[1].lower() in valid_ext
-        ]
+    # if os.path.isfile(background_image_path):
+    #     background_paths = [background_image_path]
+    # else:
+    #     valid_ext = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
+    #     background_paths = [
+    #         os.path.join(background_image_path, f) 
+    #         for f in os.listdir(background_image_path)
+    #         if os.path.splitext(f)[1].lower() in valid_ext
+    #     ]
     
     # 递归遍历输入目录
     for root, dirs, files in os.walk(input_folder):
@@ -136,10 +136,10 @@ def process_images(input_folder, background_image_path, output_base):
 
             try:
                 # 去背景处理
-                foreground = remove_background(input_path)
+                result = remove_background(input_path)
                 
            
-                result = edge_fill2(foreground)
+                # result = edge_fill2(result)
                 
                 # 保存结果
                 cv2.imwrite(output_path, result)
@@ -150,8 +150,8 @@ def process_images(input_folder, background_image_path, output_base):
 
 
 # 使用示例
-input_directory = 'L:/Tobacco/2023_JY/20230821/SOURCE'
+input_directory = 'L:/Grade_datasets/JY_A'
 background_image_path = 'F:/dataset/02.TA_EC/rundata/BACKGROUND/ZY_B'
-output_directory = 'L:/Test'
+output_directory = 'L:/Grade_datasets/MOVE_BACKGROUND'
 
 process_images(input_directory, background_image_path, output_directory)

dataset/test.py

@@ -0,0 +1,12 @@
+import os
+ 
+def debug_walk_with_links(input_folder):
+    for root, dirs, files in os.walk(input_folder):
+        print(f'Root: {root}')
+        print(f'Dirs: {dirs}')
+        print(f'Files: {files}')
+        print('-' * 40)
+ 
+if __name__ == "__main__":
+    input_folder = 'L:/Grade_datasets'
+    debug_walk_with_links(input_folder)

main.py

@@ -7,6 +7,7 @@ from torchvision.datasets import MNIST
 from torchvision.transforms import ToTensor
 
 from model.repvit import *
+from model.mobilenetv3 import *
 from data_loader import *
 from utils import *
 
@@ -14,11 +15,11 @@ def main():
     # 初始化组件
     initialize()
     
-    model = repvit_m1_1(num_classes=10).to(config.device)
+    model = repvit_m1_0(num_classes=9).to(config.device)
     optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)
     criterion = nn.CrossEntropyLoss()
     
-    train_loader, valid_loader, test_loader = create_data_loaders('F:/dataset/02.TA_EC/datasets/EC27',batch_size=config.batch_size)
+    train_loader, valid_loader, test_loader = get_data_loader('/home/yoiannis/deep_learning/dataset/02.TA_EC/datasets/EC27',batch_size=config.batch_size,Cache='RAM')
     
     # 初始化训练器
     trainer = Trainer(model, train_loader, valid_loader, optimizer, criterion)

trainner.py

@@ -4,6 +4,7 @@ from torch.utils.data import DataLoader
 from config import config
 from logger import logger
 from utils import save_checkpoint, load_checkpoint
+import time
 
 class Trainer:
     def __init__(self, model, train_loader, val_loader, optimizer, criterion):
@@ -21,7 +22,7 @@ class Trainer:
         self.model.train()
         total_loss = 0.0
         progress_bar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{config.epochs}")
-        
+        time_start = time.time()
         for batch_idx, (data, target) in enumerate(progress_bar):
             data, target = data.to(config.device), target.to(config.device)