0_get_started.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "def get_predictions(model, inputs, optimizer):\n",
    "    optimizer.zero_grad()\n",
    "    return model(inputs)\n",
    "\n",
    "def update_model(data, model, criterion, optimizer):\n",
    "    inputs, labels = data\n",
    "    preds = get_predictions(model, inputs, optimizer)\n",
    "    loss = criterion(preds, labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss.item()\n",
    "\n",
    "def get_transforms(norm=0.5):\n",
    "    return transforms.Compose([transforms.ToTensor(),\n",
    "                               transforms.Normalize((norm, norm, norm), (norm, norm, norm))])\n",
    "\n",
    "def get_data(transforms, batch_size=4):\n",
    "    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                            download=True, transform=transforms)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
    "                                            shuffle=True, num_workers=2)\n",
    "    return trainloader\n",
    "\n",
    "def save_model(model, path):\n",
    "    torch.save(model.state_dict(), path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    config = {\n",
    "        'norm': 0.5,\n",
    "        'batch_size': 8,\n",
    "        'lr': 0.05,\n",
    "        'momentum': 0.9,\n",
    "        'epochs': 1\n",
    "    }\n",
    "\n",
    "    # setup training\n",
    "    #if you want to collaborate with your team, please specifiy with entity = \"\" in wandb.init\n",
    "    with wandb.init(project='get_started1', config=config):\n",
    "        config = wandb.config\n",
    "\n",
    "        transforms = get_transforms(config.norm)\n",
    "        data = get_data(transforms, config.batch_size)\n",
    "        model = Net()\n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "        optimizer = optim.SGD(model.parameters(), lr=config.lr, momentum=config.momentum)\n",
    "\n",
    "        # train model\n",
    "        for epoch in range(config.epochs):\n",
    "            for i, batch in enumerate(data, 0):\n",
    "                loss = update_model(batch, model, criterion, optimizer)\n",
    "\n",
    "                # log results\n",
    "                wandb.log({'epoch': epoch, 'loss': loss, \"lr\":config.lr})\n",
    "        path = './cifar_net.pth'\n",
    "        save_model(model, path)\n",
    "\n",
    "train()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}