{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "#!/usr/bin/env python\n", "# -*- coding: utf-8 -*-\n", "# Python version: 3.6\n", "\n", "\n", "import os\n", "import copy\n", "import time\n", "import pickle\n", "import numpy as np\n", "from tqdm import tqdm\n", "\n", "import torch\n", "from tensorboardX import SummaryWriter\n", "\n", "from options import args_parser\n", "from update import LocalUpdate, test_inference\n", "from models import MLP, CNNMnist, CNNFashion_Mnist, CNNCifar\n", "from utils import get_dataset, average_weights, exp_details" ] }, { "cell_type": "code", "execution_count": 83, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Namespace(dataset='mnist', epochs=2, frac=0.1, gpu=None, iid=1, local_bs=10, local_ep=10, lr=0.01, model='mlp', num_classes=10, num_users=100, optimizer='sgd', verbose=1)" ] }, "execution_count": 83, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import argparse\n", "\n", "parser = argparse.ArgumentParser()\n", "# parser.add_argument('--name', '-n', default='foo', help='foo')\n", "parser.add_argument('--model', type=str, default='mlp', help='model name(mlp or cnn)')\n", "parser.add_argument('--dataset', type=str, default='mnist', help=\"name of dataset(mnist or cifar)\")\n", "parser.add_argument('--epochs', type=int, default=2, help=\"number of rounds of training(10)\")\n", "parser.add_argument('--iid', type=int, default=1, help='Default set to IID. Set to 0 for non-IID.')\n", "parser.add_argument('--num_users', type=int, default=100, help=\"number of users: K\")\n", "parser.add_argument('--num_classes', type=int, default=10, help=\"number of classes\")\n", "parser.add_argument('--gpu', default=None, help=\"To use cuda, set to a specific GPU ID. Default set to use CPU.\")\n", "parser.add_argument('--frac', type=float, default=0.1, help='the fraction of clients: C')\n", "parser.add_argument('--local_ep', type=int, default=10, help=\"the number of local epochs: E\")\n", "parser.add_argument('--local_bs', type=int, default=10, help=\"local batch size: B\")\n", "parser.add_argument('--lr', type=float, default=0.01, help='learning rate')\n", "parser.add_argument('--optimizer', type=str, default='sgd', help=\"type of optimizer\")\n", "parser.add_argument('--verbose', type=int, default=1, help='verbose')\n", "# parser.add_argument('--seed', type=int, default=1, help='random seed')\n", "\n", "\n", "\n", "\n", "# args = parser.parse_args([])\n", "args, _ = parser.parse_known_args()\n", "\n", "args\n", "# --dataset=cifar --gpu=0 --iid=0 --epochs=10" ] }, { "cell_type": "code", "execution_count": 84, "metadata": {}, "outputs": [], "source": [ "# if __name__ == '__main__':\n", "start_time = time.time()\n", "\n", "# define paths\n", "path_project = os.path.abspath('..')\n", "logger = SummaryWriter('../logs')\n", "\n", "# args = args_parser()\n", "# exp_details(args)\n", "\n", "if args.gpu:\n", " torch.cuda.set_device(args.gpu)\n", "device = 'cuda' if args.gpu else 'cpu'" ] }, { "cell_type": "code", "execution_count": 85, "metadata": {}, "outputs": [], "source": [ "# load dataset and user groups\n", "train_dataset, test_dataset, user_groups = get_dataset(args)\n", "\n", "# BUILD MODEL\n", "if args.model == 'cnn':\n", " # Convolutional neural netork\n", " if args.dataset == 'mnist':\n", " global_model = CNNMnist(args=args)\n", " elif args.dataset == 'fmnist':\n", " global_model = CNNFashion_Mnist(args=args)\n", " elif args.dataset == 'cifar':\n", " global_model = CNNCifar(args=args)\n", "\n", "elif args.model == 'mlp':\n", " # Multi-layer preceptron\n", " img_size = train_dataset[0][0].shape\n", " len_in = 1\n", " for x in img_size:\n", " len_in *= x\n", " global_model = MLP(dim_in=len_in, dim_hidden=64,\n", " dim_out=args.num_classes)\n", "else:\n", " exit('Error: unrecognized model')" ] }, { "cell_type": "code", "execution_count": 86, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "MLP(\n", " (layer_input): Linear(in_features=784, out_features=64, bias=True)\n", " (relu): ReLU()\n", " (dropout): Dropout(p=0.5, inplace=False)\n", " (layer_hidden): Linear(in_features=64, out_features=10, bias=True)\n", " (softmax): Softmax(dim=1)\n", ")\n" ] } ], "source": [ "# Set the model to train and send it to device.\n", "global_model.to(device)\n", "global_model.train()\n", "print(global_model)\n", "\n", "# copy weights\n", "global_weights = global_model.state_dict()" ] }, { "cell_type": "code", "execution_count": 87, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/2 [00:00