minor fix + structuring

experiments/hetro/dataset.py → experiments/dataset.py

@@ -166,7 +166,7 @@ def gen_random_loaders(data_name, data_path, num_users, bz, classes_per_user):
     dataloaders = []
     datasets = get_datasets(data_name, data_path, normalize=True)
     for i, d in enumerate(datasets):
-        # we want the same partition for train/test/val
+        # ensure same partition for train/test/val
         if i == 0:
             cls_partitions = gen_classes_per_node(d, num_users, classes_per_user)
             loader_params['shuffle'] = True

experiments/hetro/models.py → experiments/pfedhn/models.py

@@ -1,9 +1,5 @@
-import random
-from abc import abstractmethod
 from collections import OrderedDict
-from typing import List
 
-import torch
 import torch.nn.functional as F
 from torch import nn
 from torch.nn.utils import spectral_norm
@@ -73,9 +69,9 @@ def forward(self, idx):
         return weights
 
 
-class CNNTargetLook(nn.Module):
+class CNNTarget(nn.Module):
     def __init__(self, in_channels=3, n_kernels=16, out_dim=10):
-        super(CNNTargetLook, self).__init__()
+        super(CNNTarget, self).__init__()
 
         self.conv1 = nn.Conv2d(in_channels, n_kernels, 5)
         self.pool = nn.MaxPool2d(2, 2)

experiments/node.py → experiments/pfedhn/node.py

@@ -1,6 +1,4 @@
-import os
-
-from experiments.hetro.dataset import gen_random_loaders
+from experiments.dataset import gen_random_loaders
 
 
 class BaseNodes:

experiments/hetro/trainer.py → experiments/pfedhn/trainer.py

@@ -10,8 +10,8 @@
 import torch.utils.data
 from tqdm import trange
 
-from experiments.hetro.models import CNNHyper, CNNTargetLook
-from experiments.node import BaseNodes
+from experiments.pfedhn.models import CNNHyper, CNNTarget
+from experiments.pfedhn.node import BaseNodes
 from experiments.utils import get_device, set_logger, set_seed, str2bool
 
 
@@ -78,11 +78,11 @@ def train(data_name: str, data_path: str, classes_per_user: int, num_nodes: int,
 
     if data_name == "cifar10":
         hnet = CNNHyper(num_nodes, embed_dim, hidden_dim=hyper_hid, n_hidden=n_hidden, n_kernels=n_kernels)
-        net = CNNTargetLook(n_kernels=n_kernels)
+        net = CNNTarget(n_kernels=n_kernels)
     elif data_name == "cifar100":
         hnet = CNNHyper(num_nodes, embed_dim, hidden_dim=hyper_hid,
                         n_hidden=n_hidden, n_kernels=n_kernels, out_dim=100)
-        net = CNNTargetLook(n_kernels=n_kernels, out_dim=100)
+        net = CNNTarget(n_kernels=n_kernels, out_dim=100)
     else:
         raise ValueError("choose data_name from ['cifar10', 'cifar100']")
 
@@ -214,7 +214,10 @@ def train(data_name: str, data_path: str, classes_per_user: int, num_nodes: int,
         _, val_avg_loss, val_avg_acc, _ = eval_model(nodes, num_nodes, hnet, net, criteria, device, split="val")
         step_results, avg_loss, avg_acc, all_acc = eval_model(nodes, num_nodes, hnet, net, criteria, device, split="test")
         logging.info(f"\nStep: {step + 1}, AVG Loss: {avg_loss:.4f},  AVG Acc: {avg_acc:.4f}")
-        results[step + 1] = step_results
+
+        results['test_avg_loss'].append(avg_loss)
+        results['test_avg_acc'].append(avg_acc)
+
         if best_acc < val_avg_acc:
             best_acc = val_avg_acc
             best_step = step
@@ -234,7 +237,7 @@ def train(data_name: str, data_path: str, classes_per_user: int, num_nodes: int,
 
     save_path = Path(save_path)
     save_path.mkdir(parents=True, exist_ok=True)
-    with open(str(save_path / f"results_{inner_steps}_la_steps_seed_{seed}.json"), "w") as file:
+    with open(str(save_path / f"results_{inner_steps}_inner_steps_seed_{seed}.json"), "w") as file:
         json.dump(results, file, indent=4)
 
 
@@ -252,7 +255,7 @@ def train(data_name: str, data_path: str, classes_per_user: int, num_nodes: int,
     )
     parser.add_argument("--data-path", type=str, default="data", help="dir path for MNIST dataset")
     parser.add_argument("--classes_per_user", type=int, default=2, help="N classes assigned to each user")
-    parser.add_argument("--num-nodes", type=int, default=20, help="number of simulated nodes")
+    parser.add_argument("--num-nodes", type=int, default=50, help="number of simulated nodes")
 
     ##################################
     #       Optimization args        #
@@ -281,8 +284,8 @@ def train(data_name: str, data_path: str, classes_per_user: int, num_nodes: int,
     #       General args        #
     #############################
     parser.add_argument("--gpu", type=int, default=0, help="gpu device ID")
-    parser.add_argument("--eval-every", type=int, default=10, help="eval every X selected epochs")
-    parser.add_argument("--save-path", type=str, default="fhn_hetro", help="dir path for output file")
+    parser.add_argument("--eval-every", type=int, default=30, help="eval every X selected epochs")
+    parser.add_argument("--save-path", type=str, default="pfedhn_hetro_res", help="dir path for output file")
     parser.add_argument("--seed", type=int, default=42, help="seed value")
 
     args = parser.parse_args()

experiments/hetro/local_layers/models.py → experiments/pfedhn_pc/models.py

@@ -19,7 +19,7 @@ def forward(self, x):
         return self.layer(x)
 
 
-class CNNHyper(nn.Module):
+class CNNHyperPC(nn.Module):
     def __init__(
             self, n_nodes, embedding_dim, in_channels=3, out_dim=10, n_kernels=16, hidden_dim=100,
             spec_norm=False, n_hidden=1):
@@ -49,8 +49,6 @@ def __init__(
         self.l1_bias = nn.Linear(hidden_dim, 120)
         self.l2_weights = nn.Linear(hidden_dim, 84 * 120)
         self.l2_bias = nn.Linear(hidden_dim, 84)
-        # self.l3_weights = nn.Linear(hidden_dim, self.out_dim * 84)
-        # self.l3_bias = nn.Linear(hidden_dim, self.out_dim)
 
         if spec_norm:
             self.c1_weights = spectral_norm(self.c1_weights)
@@ -61,8 +59,6 @@ def __init__(
             self.l1_bias = spectral_norm(self.l1_bias)
             self.l2_weights = spectral_norm(self.l2_weights)
             self.l2_bias = spectral_norm(self.l2_bias)
-            # self.l3_weights = spectral_norm(self.l3_weights)
-            # self.l3_bias = spectral_norm(self.l3_bias)
 
     def forward(self, idx):
         emd = self.embeddings(idx)
@@ -77,28 +73,24 @@ def forward(self, idx):
             "fc1.bias": self.l1_bias(features).view(-1),
             "fc2.weight": self.l2_weights(features).view(84, 120),
             "fc2.bias": self.l2_bias(features).view(-1),
-            # "fc3.weight": self.l3_weights(features).view(self.out_dim, 84),
-            # "fc3.bias": self.l3_bias(features).view(-1),
         }
         return weights
 
 
-class CNNTarget(nn.Module):
-    def __init__(self, in_channels=3, n_kernels=16, out_dim=10):
+class CNNTargetPC(nn.Module):
+    def __init__(self, in_channels=3, n_kernels=16):
         super().__init__()
 
         self.conv1 = nn.Conv2d(in_channels, n_kernels, 5)
         self.pool = nn.MaxPool2d(2, 2)
         self.conv2 = nn.Conv2d(n_kernels, 2 * n_kernels, 5)
         self.fc1 = nn.Linear(2 * n_kernels * 5 * 5, 120)
         self.fc2 = nn.Linear(120, 84)
-        # self.fc3 = nn.Linear(84, out_dim)
 
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
         x = self.pool(F.relu(self.conv2(x)))
         x = x.view(x.shape[0], -1)
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
-        # x = self.fc3(x)
         return x

experiments/hetro/local_layers/node.py → experiments/pfedhn_pc/node.py

@@ -1,4 +1,4 @@
-from experiments.hetro.dataset import gen_random_loaders
+from experiments.dataset import gen_random_loaders
 
 
 class BaseNodesForLocal:

experiments/hetro/local_layers/trainer.py → experiments/pfedhn_pc/trainer.py

@@ -10,9 +10,8 @@
 import torch.utils.data
 from tqdm import trange
 
-from experiments.hetro.local_layers.models import (CNNHyper, CNNTarget,
-                                                   LocalLayer)
-from experiments.hetro.local_layers.node import BaseNodesForLocal
+from experiments.pfedhn_pc.models import CNNHyperPC, CNNTargetPC, LocalLayer
+from experiments.pfedhn_pc.node import BaseNodesForLocal
 from experiments.utils import get_device, set_logger, set_seed, str2bool
 
 
@@ -65,7 +64,7 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
           steps: int, inner_steps: int, optim: str, lr: float, inner_lr: float,
           embed_lr: float, wd: float, inner_wd: float, embed_dim: int, hyper_hid: int,
           n_hidden: int, n_kernels: int, bs: int, device, eval_every: int, save_path: Path,
-          seed: int) -> None:
+          ) -> None:
 
     ###############################
     # init nodes, hnet, local net #
@@ -88,11 +87,11 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
         logging.info("auto embedding size")
         embed_dim = int(1 + num_nodes / 4)
 
-    hnet = CNNHyper(
+    hnet = CNNHyperPC(
         num_nodes, embed_dim, hidden_dim=hyper_hid, n_hidden=n_hidden,
         n_kernels=n_kernels
     )
-    net = CNNTarget(n_kernels=n_kernels)
+    net = CNNTargetPC(n_kernels=n_kernels)
 
     hnet = hnet.to(device)
     net = net.to(device)
@@ -227,7 +226,10 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
         _, val_avg_loss, val_avg_acc, _ = eval_model(nodes, num_nodes, hnet, net, criteria, device, split="val")
         step_results, avg_loss, avg_acc, all_acc = eval_model(nodes, num_nodes, hnet, net, criteria, device, split="test")
         logging.info(f"\nStep: {step + 1}, AVG Loss: {avg_loss:.4f},  AVG Acc: {avg_acc:.4f}")
-        results[step + 1] = step_results
+
+        results['test_avg_loss'].append(avg_loss)
+        results['test_avg_acc'].append(avg_acc)
+
         if best_acc < val_avg_acc:
             best_acc = val_avg_acc
             best_step = step
@@ -269,7 +271,7 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
     ##################################
     #       Optimization args        #
     ##################################
-    parser.add_argument("--num-steps", type=int, default=2000)
+    parser.add_argument("--num-steps", type=int, default=5000)
     parser.add_argument("--batch-size", type=int, default=64)
     parser.add_argument("--inner-steps", type=int, default=50, help="number of inner steps")
     parser.add_argument("--optim", type=str, default='sgd', choices=['adam', 'sgd'], help="learning rate")
@@ -279,7 +281,7 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
     ################################
     parser.add_argument("--n-hidden", type=int, default=3, help="num. hidden layers")
     parser.add_argument("--inner-lr", type=float, default=5e-3, help="learning rate for inner optimizer")
-    parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
+    parser.add_argument("--lr", type=float, default=5e-2, help="learning rate")
     parser.add_argument("--wd", type=float, default=1e-3, help="weight decay")
     parser.add_argument("--inner-wd", type=float, default=5e-5, help="inner weight decay")
     parser.add_argument("--embed-dim", type=int, default=-1, help="embedding dim")
@@ -292,8 +294,8 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
     #       General args        #
     #############################
     parser.add_argument("--gpu", type=int, default=0, help="gpu device ID")
-    parser.add_argument("--eval-every", type=int, default=10, help="eval every X selected epochs")
-    parser.add_argument("--save-path", type=str, default="cifar_local_layers", help="dir path for output file")
+    parser.add_argument("--eval-every", type=int, default=30, help="eval every X selected epochs")
+    parser.add_argument("--save-path", type=str, default="pfedhn_pc_cifar_res", help="dir path for output file")
     parser.add_argument("--seed", type=int, default=42, help="seed value")
 
     args = parser.parse_args()
@@ -325,5 +327,4 @@ def train(data_name: str, data_path: str, classes_per_node: int, num_nodes: int,
         device=device,
         eval_every=args.eval_every,
         save_path=args.save_path,
-        seed=args.seed
     )