Modify the data loading mechanism and separate validation class

Author: maduc238

README.md

@@ -102,7 +102,6 @@ learning:
   momentum: 0.5
   batch-size: 256
   control-count: 3    # control count on client
-  validation: False   # run validate on client side
 ```
 
 This configuration is use for server.
@@ -152,7 +151,7 @@ If the `*.pth` file exists, the server will read the file and send the parameter
 
 ---
 
-Version 1.8.0
+Version 2.0.0
 
 The application is under development...
 

client.py

@@ -13,6 +13,7 @@
 parser = argparse.ArgumentParser(description="Split learning framework")
 parser.add_argument('--layer_id', type=int, required=True, help='ID of layer, start from 1')
 parser.add_argument('--device', type=str, required=False, help='Device of client')
+parser.add_argument('--event_time', type=bool, default=False, required=False, help='Log event time for debug mode')
 
 args = parser.parse_args()
 
@@ -46,7 +47,7 @@
 if __name__ == "__main__":
     src.Log.print_with_color("[>>>] Client sending registration message to server...", "red")
     data = {"action": "REGISTER", "client_id": client_id, "layer_id": args.layer_id, "message": "Hello from Client!"}
-    scheduler = Scheduler(client_id, args.layer_id, channel, device)
+    scheduler = Scheduler(client_id, args.layer_id, channel, device, args.event_time)
     client = RpcClient(client_id, args.layer_id, address, username, password, scheduler.train_on_device, device)
     client.send_to_server(data)
     client.wait_response()

config.yaml

@@ -26,4 +26,3 @@ learning:
   momentum: 0.5
   batch-size: 128
   control-count: 3
-  validation: False

src/Model.py

@@ -1,12 +1,5 @@
 import torch
 import torch.nn as nn
-import numpy as np
-import math
-from tqdm import tqdm
-
-import torchvision
-import torchvision.transforms as transforms
-import torch.nn.functional as F
 
 if torch.cuda.is_available():
     device = "cuda"
@@ -464,44 +457,3 @@ def forward(self, x):
         out83 = self.layer83(out82)
         out84 = self.layer84(out83)
         return out84
-
-
-def test(model_name, state_dict_full, logger):
-    transform_test = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
-    ])
-
-    testset = torchvision.datasets.CIFAR10(
-        root='./data', train=False, download=True, transform=transform_test)
-    test_loader = torch.utils.data.DataLoader(
-        testset, batch_size=100, shuffle=False, num_workers=2)
-
-    klass = globals().get(model_name)
-    if klass is None:
-        raise ValueError(f"Class '{model_name}' does not exist.")
-    model = klass()
-    model = nn.Sequential(*nn.ModuleList(model.children()))
-    model.load_state_dict(state_dict_full)
-    # evaluation mode
-    model.eval()
-    test_loss = 0
-    correct = 0
-    for data, target in tqdm(test_loader):
-        output = model(data)
-        test_loss += F.nll_loss(output, target, reduction='sum').item()
-        pred = output.data.max(1, keepdim=True)[1]  # get the index of the max log-probability
-        correct += pred.eq(target.data.view_as(pred)).long().cpu().sum()
-
-    test_loss /= len(test_loader.dataset)
-    accuracy = 100.0 * correct / len(test_loader.dataset)
-    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
-        test_loss, correct, len(test_loader.dataset), accuracy))
-
-    if np.isnan(test_loss) or math.isnan(test_loss) or abs(test_loss) > 10e5:
-        return False
-    else:
-        logger.log_info('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
-            test_loss, correct, len(test_loader.dataset), accuracy))
-
-    return True

src/RpcClient.py

@@ -1,7 +1,13 @@
 import time
 import pickle
 import pika
+import random
+import torch
+import torchvision
+import torchvision.transforms as transforms
+
 from torch import nn
+from collections import defaultdict
 
 import src.Log
 import src.Model
@@ -23,6 +29,24 @@ def __init__(self, client_id, layer_id, address, username, password, train_func,
         self.model = None
         self.connect()
 
+        self.train_set = None
+        self.label_to_indices = None
+        if self.layer_id == 1:
+            # Read and load dataset
+            transform_train = transforms.Compose([
+                transforms.RandomCrop(32, padding=4),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+            ])
+
+            self.train_set = torchvision.datasets.CIFAR10(
+                root='./data', train=True, download=True, transform=transform_train)
+
+            self.label_to_indices = defaultdict(list)
+            for idx, (_, label) in enumerate(self.train_set):
+                self.label_to_indices[label].append(idx)
+
     def wait_response(self):
         status = True
         reply_queue_name = f'reply_{self.client_id}'
@@ -66,12 +90,20 @@ def response_message(self, body):
             batch_size = self.response["batch_size"]
             lr = self.response["lr"]
             momentum = self.response["momentum"]
-            validation = self.response["validation"]
             control_count = self.response["control_count"]
 
             # Start training
-            result, size = self.train_func(self.model, control_count, batch_size, lr, momentum,
-                                           validation, label_count, num_layers)
+            if self.layer_id == 1:
+                selected_indices = []
+                for label, count in enumerate(label_count):
+                    selected_indices.extend(random.sample(self.label_to_indices[label], count))
+
+                subset = torch.utils.data.Subset(self.train_set, selected_indices)
+                train_loader = torch.utils.data.DataLoader(subset, batch_size=batch_size, shuffle=True)
+
+                result, size = self.train_func(self.model, lr, momentum, num_layers, control_count, train_loader)
+            else:
+                result, size = self.train_func(self.model, lr, momentum, num_layers, control_count)
 
             # Stop training, then send parameters to server
             model_state_dict = self.model.state_dict()