delta_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4149fa67-895f-4129-a47c-c84edfbc7af2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from delta import DeltaNode\n",
    "from delta.task import LearningTask\n",
    "\n",
    "\n",
    "class LeNet(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = torch.nn.Conv2d(1, 16, 5, padding=2)\n",
    "        self.pool1 = torch.nn.AvgPool2d(2, stride=2)\n",
    "        self.conv2 = torch.nn.Conv2d(16, 16, 5)\n",
    "        self.pool2 = torch.nn.AvgPool2d(2, stride=2)\n",
    "        self.dense1 = torch.nn.Linear(400, 100)\n",
    "        self.dense2 = torch.nn.Linear(100, 10)\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        x = self.conv1(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.pool1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.pool2(x)\n",
    "\n",
    "        x = x.view(-1, 400)\n",
    "        x = self.dense1(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.dense2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "model = LeNet()\n",
    "loss = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(\n",
    "    model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-3, nesterov=True\n",
    ")\n",
    "\n",
    "\n",
    "def train_step(batch, model=model, loss=loss, optimizer=optimizer):\n",
    "    x, y = batch\n",
    "    y_pred = model(x)\n",
    "    loss_val = loss(y_pred, y)\n",
    "    optimizer.zero_grad()\n",
    "    loss_val.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "\n",
    "def preprocess(x: np.ndarray, y: str):\n",
    "    x /= 255.0\n",
    "    x *= 2\n",
    "    x -= 1\n",
    "    x = x.reshape((1, 28, 28))\n",
    "    return torch.from_numpy(x), torch.tensor(int(y), dtype=torch.long)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    logging.basicConfig(level=logging.DEBUG)\n",
    "    task = LearningTask(\n",
    "        name=\"test\",\n",
    "        dataset=\"mnist\",\n",
    "        preprocess=preprocess,\n",
    "        train_step=train_step,\n",
    "        dataloader={\"batch_size\": 256, \"shuffle\": True, \"drop_last\": True},\n",
    "        total_epoch=1,\n",
    "        members=[\n",
    "            \"55ZUVDZ4FIUAR6TFSH6I7CP6XPXWYC5AKMYDRJIIUONGLJ6E\",\n",
    "            \"Y7L6HUG5BCXWCXR5ADW2ATHA5PX4KNG74I2CBWULOEFVBWH5\",\n",
    "        ],\n",
    "        merge_iter=20,\n",
    "    )\n",
    "    delta_node = DeltaNode(\"http://127.0.0.1:6701\")\n",
    "    delta_node.create_task(task)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}