Add ResNet18 and ResNet50

main_fed.py

@@ -144,11 +144,11 @@
 
  best_save_path = os.path.join(base_dir, algo_dir, 'best_model.pt')
 
-#  torch.save(net_local_list[0].state_dict(), best_save_path)
+ torch.save(net_local_list[0].state_dict(), best_save_path)
 
- for user_idx in range(args.num_users):
- best_save_path = os.path.join(base_dir, algo_dir, 'best_local_{}.pt'.format(user_idx))
- torch.save(net_local_list[user_idx].state_dict(), best_save_path)
+#  for user_idx in range(args.num_users):
+#  best_save_path = os.path.join(base_dir, algo_dir, 'best_local_{}.pt'.format(user_idx))
+#  torch.save(net_local_list[user_idx].state_dict(), best_save_path)
 
  results.append(np.array([iter, loss_avg, loss_test, acc_test, best_acc]))
  final_results = np.array(results)

models/ResNet.py

@@ -0,0 +1,122 @@
+'''ResNet in PyTorch.
+For Pre-activation ResNet, see 'preact_resnet.py'.
+Reference:
+[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
+ Deep Residual Learning for Image Recognition. arXiv:1512.03385
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, in_planes, planes, stride=1):
+ super(BasicBlock, self).__init__()
+ self.conv1 = nn.Conv2d(
+ in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+ stride=1, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+
+ self.shortcut = nn.Sequential()
+ if stride != 1 or in_planes != self.expansion*planes:
+ self.shortcut = nn.Sequential(
+ nn.Conv2d(in_planes, self.expansion*planes,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(self.expansion*planes)
+ )
+
+ def forward(self, x):
+ out = F.relu(self.bn1(self.conv1(x)))
+ out = self.bn2(self.conv2(out))
+ out += self.shortcut(x)
+ out = F.relu(out)
+ return out
+
+
+class Bottleneck(nn.Module):
+ expansion = 4
+
+ def __init__(self, in_planes, planes, stride=1):
+ super(Bottleneck, self).__init__()
+ self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+ stride=stride, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.conv3 = nn.Conv2d(planes, self.expansion *
+ planes, kernel_size=1, bias=False)
+ self.bn3 = nn.BatchNorm2d(self.expansion*planes)
+
+ self.shortcut = nn.Sequential()
+ if stride != 1 or in_planes != self.expansion*planes:
+ self.shortcut = nn.Sequential(
+ nn.Conv2d(in_planes, self.expansion*planes,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(self.expansion*planes)
+ )
+
+ def forward(self, x):
+ out = F.relu(self.bn1(self.conv1(x)))
+ out = F.relu(self.bn2(self.conv2(out)))
+ out = self.bn3(self.conv3(out))
+ out += self.shortcut(x)
+ out = F.relu(out)
+ return out
+
+
+class ResNet(nn.Module):
+ def __init__(self, block, num_blocks, num_classes=10):
+ super(ResNet, self).__init__()
+ self.in_planes = 64
+
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=3,
+ stride=1, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+ self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+ self.linear = nn.Linear(512*block.expansion, num_classes)
+
+ def _make_layer(self, block, planes, num_blocks, stride):
+ strides = [stride] + [1]*(num_blocks-1)
+ layers = []
+ for stride in strides:
+ layers.append(block(self.in_planes, planes, stride))
+ self.in_planes = planes * block.expansion
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ out = F.relu(self.bn1(self.conv1(x)))
+ out = self.layer1(out)
+ out = self.layer2(out)
+ out = self.layer3(out)
+ out = self.layer4(out)
+ out = F.avg_pool2d(out, 4)
+ out = out.view(out.size(0), -1)
+ out = self.linear(out)
+ return out
+
+
+def ResNet18(num_classes=10):
+ return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes)
+
+
+def ResNet34(num_classes=10):
+ return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes)
+
+
+def ResNet50(num_classes=10):
+ return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes)
+
+
+def ResNet101(num_classes=10):
+ return ResNet(Bottleneck, [3, 4, 23, 3], num_classes=num_classes)
+
+
+def ResNet152(num_classes=10):
+ return ResNet(Bottleneck, [3, 8, 36, 3], num_classes=num_classes)

utils/train_utils.py

@@ -1,5 +1,6 @@
 from torchvision import datasets, transforms
 from models.Nets import CNNCifar, MobileNetCifar
+from models.ResNet import ResNet18, ResNet50
 from utils.sampling import iid, noniid, iid_unbalanced, noniid_unbalanced
 
 trans_mnist = transforms.Compose([transforms.ToTensor(),
@@ -94,6 +95,10 @@ def get_model(args):
  net_glob = CNNCifar(args=args).to(args.device)
  elif args.model == 'mobile' and args.dataset in ['cifar10', 'cifar100']:
  net_glob = MobileNetCifar(num_classes=args.num_classes).to(args.device)
+ elif args.model == 'resnet18' and args.dataset in ['cifar10', 'cifar100']:
+ net_glob = ResNet18(num_classes=args.num_classes).to(args.device)
+ elif args.model == 'resnet50' and args.dataset in ['cifar10', 'cifar100']:
+ net_glob = ResNet50(num_classes=args.num_classes).to(args.device)
  elif args.model == 'cnn' and args.dataset == 'mnist':
  net_glob = CNNMnist(args=args).to(args.device)
  elif args.model == 'mlp' and args.dataset == 'mnist':