paper-figures/appendix_figure_9_distributed_mean_estimation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format='retina'\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "from algorithms import ChainTopology, RingTopology, QuadraticsTask, relaysum_model, relaysum_grad, relaysum_mix, all_reduce, gossip, d2, gradient_tracking, BinaryTreeTopology, relaysum_grad_overlap, Iterate, StarTopology\n",
    "from random_quadratics import RandomQuadraticsTask\n",
    "from collections import defaultdict\n",
    "import torch\n",
    "import math\n",
    "import tuning\n",
    "import pandas as pd\n",
    "from warehouse import Warehouse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_default_dtype(torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relaysum_dme(world, num_steps):\n",
    "    state = torch.zeros([world.num_workers])\n",
    "    local_numbers = torch.zeros([world.num_workers])\n",
    "\n",
    "    # Initialize messages between connected workers to 0\n",
    "    messages: Mapping[Edge, float] = defaultdict(float)\n",
    "\n",
    "    num_messages_sent = 0\n",
    "\n",
    "    for step in range(num_steps):\n",
    "        yield Iterate(step, state, num_messages_sent)\n",
    "        samples = torch.randn_like(state) + target\n",
    "\n",
    "        local_numbers_update = (samples - local_numbers) / (step + 1)\n",
    "        local_numbers += local_numbers_update\n",
    "        \n",
    "        new_messages = {}\n",
    "        for worker in world.workers:\n",
    "            neighbors = world.neighbors(worker)\n",
    "            for neighbor in neighbors:\n",
    "                new_messages[worker, neighbor] = local_numbers_update[worker] + sum(\n",
    "                    messages[n, worker] for n in neighbors if n != neighbor\n",
    "                )\n",
    "\n",
    "        messages = new_messages\n",
    "\n",
    "        for worker in world.workers:\n",
    "            neighbors = world.neighbors(worker)\n",
    "            sum_updates = local_numbers_update[worker] + sum(new_messages[n, worker] for n in neighbors)\n",
    "            state[worker] += sum_updates / world.num_workers\n",
    "\n",
    "        num_messages_sent += world.max_degree\n",
    "\n",
    "    yield Iterate(step, state, num_messages_sent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relaysum_dme_corr(world, num_steps):\n",
    "    state = torch.zeros([world.num_workers])\n",
    "    counts = torch.zeros([world.num_workers])\n",
    "\n",
    "    # Initialize messages between connected workers to 0\n",
    "    messages: Mapping[Edge, float] = defaultdict(float)\n",
    "\n",
    "    num_messages_sent = 0\n",
    "\n",
    "    for step in range(num_steps):\n",
    "        yield Iterate(step, state / counts, num_messages_sent)\n",
    "        samples = torch.randn_like(state) + target\n",
    "        \n",
    "        new_messages = {}\n",
    "        for worker in world.workers:\n",
    "            neighbors = world.neighbors(worker)\n",
    "            for neighbor in neighbors:\n",
    "                new_messages[worker, neighbor, \"sample\"] = samples[worker] + sum(\n",
    "                    messages[n, worker, \"sample\"] for n in neighbors if n != neighbor\n",
    "                )\n",
    "                new_messages[worker, neighbor, \"count\"] = 1 + sum(\n",
    "                    messages[n, worker, \"count\"] for n in neighbors if n != neighbor\n",
    "                )\n",
    "\n",
    "        messages = new_messages\n",
    "\n",
    "        for worker in world.workers:\n",
    "            neighbors = world.neighbors(worker)\n",
    "            sum_updates = samples[worker] + sum(new_messages[n, worker, \"sample\"] for n in neighbors)\n",
    "            count_updates = 1 + sum(new_messages[n, worker, \"count\"] for n in neighbors)\n",
    "            state[worker] += sum_updates\n",
    "            counts[worker] += count_updates\n",
    "\n",
    "        num_messages_sent += world.max_degree\n",
    "\n",
    "    yield Iterate(step, state / counts, num_messages_sent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gossip_dme(world, num_steps):\n",
    "    state = torch.zeros([world.num_workers])\n",
    "\n",
    "    W = world.gossip_matrix()\n",
    "\n",
    "    num_messages_sent = 0\n",
    "    for step in range(num_steps):\n",
    "        yield Iterate(step, state, num_messages_sent)\n",
    "        samples = torch.randn_like(state) + target\n",
    "\n",
    "        state += (samples - state) / (step + 1)\n",
    "\n",
    "        # gossip\n",
    "        state = W @ state\n",
    "\n",
    "        num_messages_sent += world.max_degree\n",
    "\n",
    "    yield Iterate(step, state, num_messages_sent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def error(state):\n",
    "    return torch.mean((state - target)**2).item()\n",
    "\n",
    "w = Warehouse()\n",
    "\n",
    "w.clear(\"error\", {\"method\": \"RelaySum\"})\n",
    "for n in [8, 16, 32, 64, 128]:\n",
    "    for seed in range(40):\n",
    "        for iterate in relaysum_dme_corr(ChainTopology(n), 400):\n",
    "            if iterate.step % 5 == 0:\n",
    "                w.log_metric(\"error\", {\"value\": error(iterate.state), \"step\": iterate.step}, {\"n\": n, \"method\": \"RelaySum\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in [8, 16, 32, 64, 128]:\n",
    "    for seed in range(40):\n",
    "        for iterate in gossip_dme(RingTopology(n), 400):\n",
    "            if iterate.step % 5 == 0:\n",
    "                w.log_metric(\"error\", {\"value\": error(iterate.state), \"step\": iterate.step}, {\"n\": n, \"method\": \"Gossip\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in [8, 16, 32, 64, 128]:\n",
    "    for iterate in gossip_dme(RingTopology(n), 400):\n",
    "        if iterate.step % 5 == 0:\n",
    "            w.log_metric(\"error\", {\"value\": 1/(n * (iterate.step + 1)), \"step\": iterate.step}, {\"n\": n, \"method\": \"1 / nT\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "sns.set_theme(\"paper\")\n",
    "sns.set_style(\"whitegrid\")\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib\n",
    "matplotlib.rcParams['text.usetex'] = True\n",
    "plt.rcParams.update({\n",
    "    \"text.usetex\": True,\n",
    "    \"font.family\": \"serif\",\n",
    "    \"font.serif\": [\"Times\"],\n",
    "    'text.latex.preamble' : r'\\usepackage{amsmath}\\usepackage{amssymb}\\usepackage{newtxmath}'\n",
    "})\n",
    "import matplotlib.ticker as mtick\n",
    "colors = [sns.color_palette(\"tab10\")[i] for i in [7, 0, 1, 9, 3]]\n",
    "sns.set_palette(colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = w.query(\"error\")\n",
    "df[\"method\"] = pd.Categorical(df.method, categories=[\"1 / nT\", \"RelaySum\", \"Gossip\"])\n",
    "g = sns.FacetGrid(col=\"method\", data = df, hue=\"n\")\n",
    "g.map(sns.lineplot, \"step\", \"value\")\n",
    "g.set(yscale=\"log\")\n",
    "g.set_ylabels(\"Mean squared error\")\n",
    "g.set_xlabels(\"Steps\")\n",
    "g.add_legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g.savefig(\"distributed-mean-estimation-results.pdf\", bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.4 64-bit ('anaconda3': conda)",
   "language": "python",
   "name": "python37464bitanaconda3conda7403c56cbe6b465c807157f843c10752"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}