APS360_project_v2.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3kOR2A-rYXc9",
        "outputId": "e3a3485e-7210-405a-d21e-5fc69a6c4608"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7a22780f1750>"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ],
      "source": [
        "import os.path\n",
        "import random\n",
        "import time\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "from torch.utils.data.sampler import SubsetRandomSampler\n",
        "import torchvision.transforms as transforms\n",
        "import numpy as np\n",
        "from sklearn.model_selection import train_test_split\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "np.random.seed(1000) # Fixed numpy random seed for reproducible shuffling\n",
        "torch.manual_seed(1) # set the random seed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "xouKk6aWYeRy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cee044bb-3b58-4bd0-c8bb-cca87bcad449"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "YOFurdsRYiIp"
      },
      "outputs": [],
      "source": [
        "def select_files(directory, fraction=0.75):\n",
        "  num_files_to_return =  int(len(os.listdir(directory)) * fraction)\n",
        "  all_files = [os.path.join(directory, f) for f in os.listdir(directory)[:num_files_to_return] if f.endswith('.jpg')]\n",
        "  #all_files = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith('.jpg')]\n",
        "  random.shuffle(all_files)\n",
        "  return all_files[:num_files_to_return]\n",
        "\n",
        "root = '/content/drive/MyDrive/APS360 Project/APS360 Project Dataset/'\n",
        "classes = ['Plastic', 'Glass', 'Paper', 'Metal']\n",
        "folders = [(root+c,c) for c in classes]\n",
        "main_file_lists = [(select_files(folder,fraction=0.4),label) for (folder,label) in folders]\n",
        "main_data = [[(torchvision.io.read_image(path),label) for path in folder] for (folder,label) in main_file_lists]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "ZDpglDEJdCBY"
      },
      "outputs": [],
      "source": [
        "main_data_flat = [item for folder in main_data for item in folder]\n",
        "main_data_x = [item[0] for item in main_data_flat] #tensors\n",
        "main_data_y = [item[1] for item in main_data_flat] #labels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "SQWpxuHuY8_3"
      },
      "outputs": [],
      "source": [
        "class TrashDataset(Dataset):\n",
        "    \"\"\"Face Landmarks dataset.\"\"\"\n",
        "\n",
        "    #def __init__(self, root_dir, transform=None):\n",
        "    def __init__(self, inputs,labels, transform=None):\n",
        "        \"\"\"\n",
        "        Arguments:\n",
        "            root_dir (string): Directory with all the class folders.\n",
        "            transform (callable, optional): Optional transform to be applied\n",
        "                on a sample.\n",
        "        \"\"\"\n",
        "        self.data_x = inputs\n",
        "        self.data_y = labels\n",
        "        self.transform = transform\n",
        "\n",
        "\n",
        "    def __len__(self):\n",
        "      #classes = ['Plastic', 'Glass', 'Paper', 'Metal']\n",
        "      #folders = [self.root_dir+c for c in classes]\n",
        "      #return np.sum([os.listdir(folder) for folder in folders])\n",
        "      return len(self.data_x)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "      if self.transform:\n",
        "          sample = (self.transform(self.data_x[idx]),self.transform(self.data_y[idx]))\n",
        "          return sample\n",
        "      sample = (self.data_x[idx],self.data_y[idx])\n",
        "      return sample"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "NN9Rr1KFY_L8"
      },
      "outputs": [],
      "source": [
        "###############################################################################\n",
        "# Data Loading\n",
        "\n",
        "def get_relevant_indices(dataset, classes, target_classes):\n",
        "    \"\"\" Return the indices for datapoints in the dataset that belongs to the\n",
        "    desired target classes, a subset of all possible classes.\n",
        "\n",
        "    Args:\n",
        "        dataset: Dataset object\n",
        "        classes: A list of strings denoting the name of each class\n",
        "        target_classes: A list of strings denoting the name of desired classes\n",
        "                        Should be a subset of the 'classes'\n",
        "    Returns:\n",
        "        indices: list of indices that have labels corresponding to one of the\n",
        "                 target classes\n",
        "    \"\"\"\n",
        "    indices = []\n",
        "    for i in range(len(dataset)):\n",
        "        # Check if the label is in the target classes\n",
        "        label_index = dataset[i][1] # ex: 3\n",
        "        label_class = classes[label_index] # ex: 'cat'\n",
        "        if label_class in target_classes:\n",
        "            indices.append(i)\n",
        "    return indices\n",
        "\n",
        "def get_data_loader(batch_size):\n",
        "    \"\"\" Loads images of cats and dogs, splits the data into training, validation\n",
        "    and testing datasets. Returns data loaders for the three preprocessed datasets.\n",
        "\n",
        "    Args:\n",
        "        target_classes: A list of strings denoting the name of the desired\n",
        "                        classes. Should be a subset of the argument 'classes'\n",
        "        batch_size: A int representing the number of samples per batch\n",
        "\n",
        "    Returns:\n",
        "        train_loader: iterable training dataset organized according to batch size\n",
        "        val_loader: iterable validation dataset organized according to batch size\n",
        "        test_loader: iterable testing dataset organized according to batch size\n",
        "        classes: A list of strings denoting the name of each class\n",
        "    \"\"\"\n",
        "\n",
        "    classes = ('Plastic','Metal','Paper','Glass')\n",
        "    ########################################################################\n",
        "    # The output of torchvision datasets are PILImage images of range [0, 1].\n",
        "    # We transform them to Tensors of normalized range [-1, 1].\n",
        "    transform = None #placeholder\n",
        "    # Load CIFAR10 training data\n",
        "    trainset = TrashDataset(main_data_x,main_data_y)\n",
        "    # Get the list of indices to sample from\n",
        "    relevant_indices = np.arange(len(main_data_x))\n",
        "\n",
        "    # Split into train and validation\n",
        "    np.random.seed(1000) # Fixed numpy random seed for reproducible shuffling\n",
        "    np.random.shuffle(relevant_indices)\n",
        "    split = int(len(relevant_indices) * 0.8) #split at 80%\n",
        "\n",
        "    # split into training and validation indices\n",
        "    relevant_train_indices, relevant_val_indices = relevant_indices[:split], relevant_indices[split:]\n",
        "    train_sampler = SubsetRandomSampler(relevant_train_indices)\n",
        "    train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
        "                                               num_workers=0, sampler=train_sampler)\n",
        "    val_sampler = SubsetRandomSampler(relevant_val_indices)\n",
        "    val_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
        "                                              num_workers=0, sampler=val_sampler)\n",
        "    \"\"\"\n",
        "    # Get the list of indices to sample from\n",
        "    relevant_test_indices = get_relevant_indices(testset, classes, target_classes)\n",
        "    test_sampler = SubsetRandomSampler(relevant_test_indices)\n",
        "    test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size,\n",
        "                                             num_workers=1, sampler=test_sampler)\n",
        "    \"\"\"\n",
        "    return train_loader, val_loader, classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "_gIOJ0E9Yt5y"
      },
      "outputs": [],
      "source": [
        "def transformed_tensor(tensor):\n",
        "  #temporary, transforms:\n",
        "  data_transform = transforms.Compose([\n",
        "    transforms.ToPILImage(),\n",
        "    transforms.RandomResizedCrop((256,256)),\n",
        "    #transforms.CenterCrop((256,256)),\n",
        "    transforms.ToTensor()#,\n",
        "    #transforms.Normalize(0, 1)\n",
        "  ])\n",
        "  if False: #needs fixing #single image passed in\n",
        "    return data_transform(tensor)\n",
        "  else: #tensor of tensors passed in\n",
        "    return torch.stack([data_transform(t) for t in tensor])\n",
        "\n",
        "class LargeNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(LargeNet, self).__init__()\n",
        "        self.name = \"large\"\n",
        "        self.conv1 = nn.Conv2d(3, 5, 5, 2, 1) #convolution layer\n",
        "        self.pool = nn.MaxPool2d(2, 2) #max pooling layer\n",
        "        self.conv2 = nn.Conv2d(5, 10, 3)  #convolution layer\n",
        "        #self.fc1 = nn.Linear(10 * 14 * 14, 32) for 128x128 inputs #linear layer\n",
        "        self.fc1 = nn.Linear(10 * 30 * 30, 32)\n",
        "        self.fc2 = nn.Linear(32, 4) #linear layer\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = transformed_tensor(x)\n",
        "        x = self.pool(F.relu(self.conv1(x))) #convolution, relu activation, max pooling\n",
        "        x = self.pool(F.relu(self.conv2(x))) #convolution, relu activation, max pooling\n",
        "        #x = x.view(-1, 10 * 14 * 14) #for 128x128 inputs\n",
        "        x = x.view(-1, 10 * 30 * 30)\n",
        "        x = F.relu(self.fc1(x)) #linear, relu activation\n",
        "        x = self.fc2(x) #linear\n",
        "        #x = x.squeeze(1) # Flatten to [batch_size]\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "zeh0jv_gZCg5"
      },
      "outputs": [],
      "source": [
        "def evaluate(net, loader, criterion):\n",
        "    \"\"\" Evaluate the network on the validation set.\n",
        "\n",
        "     Args:\n",
        "         net: PyTorch neural network object\n",
        "         loader: PyTorch data loader for the validation set\n",
        "         criterion: The loss function\n",
        "     Returns:\n",
        "         err: A scalar for the avg classification error over the validation set\n",
        "         loss: A scalar for the average loss function over the validation set\n",
        "     \"\"\"\n",
        "    total_loss = 0.0\n",
        "    total_err = 0.0\n",
        "    total_epoch = 0\n",
        "    for i, data in enumerate(loader, 0):\n",
        "        inputs, labels = data\n",
        "        labels = np.array([[1.0 if c == label else 0.0 for c in classes] for label in labels])\n",
        "        outputs = net(inputs)\n",
        "        loss = criterion(outputs, torch.tensor(labels))\n",
        "\n",
        "        probabilities = F.softmax(outputs, dim=0)\n",
        "        corr = (torch.max(probabilities,dim=1)[1]).squeeze().long() != torch.max(torch.tensor(labels),dim=1)[1]\n",
        "        total_err += int(corr.sum())\n",
        "        total_loss += loss.item()\n",
        "        total_epoch += len(labels)\n",
        "    err = float(total_err) / total_epoch\n",
        "    loss = float(total_loss) / (i + 1)\n",
        "    return err, loss\n",
        "\n",
        "###############################################################################\n",
        "# Training Curve\n",
        "def plot_training_curve(path):\n",
        "    \"\"\" Plots the training curve for a model run, given the csv files\n",
        "    containing the train/validation error/loss.\n",
        "\n",
        "    Args:\n",
        "        path: The base path of the csv files produced during training\n",
        "    \"\"\"\n",
        "    import matplotlib.pyplot as plt\n",
        "    train_err = np.loadtxt(\"{}_train_err.csv\".format(path))\n",
        "    val_err = np.loadtxt(\"{}_val_err.csv\".format(path))\n",
        "    train_loss = np.loadtxt(\"{}_train_loss.csv\".format(path))\n",
        "    val_loss = np.loadtxt(\"{}_val_loss.csv\".format(path))\n",
        "    plt.title(\"Train vs Validation Error\")\n",
        "    n = len(train_err) # number of epochs\n",
        "    plt.plot(range(1,n+1), train_err, label=\"Train\")\n",
        "    plt.plot(range(1,n+1), val_err, label=\"Validation\")\n",
        "    plt.xlabel(\"Epoch\")\n",
        "    plt.ylabel(\"Error\")\n",
        "    plt.legend(loc='best')\n",
        "    plt.show()\n",
        "    plt.title(\"Train vs Validation Loss\")\n",
        "    plt.plot(range(1,n+1), train_loss, label=\"Train\")\n",
        "    plt.plot(range(1,n+1), val_loss, label=\"Validation\")\n",
        "    plt.xlabel(\"Epoch\")\n",
        "    plt.ylabel(\"Loss\")\n",
        "    plt.legend(loc='best')\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "hqJjw35uZGCH"
      },
      "outputs": [],
      "source": [
        "def train_net(net, batch_size=64, learning_rate=0.01, num_epochs=30):\n",
        "    ########################################################################\n",
        "    # Fixed PyTorch random seed for reproducible result\n",
        "    torch.manual_seed(1000)\n",
        "    ########################################################################\n",
        "    # Obtain the PyTorch data loader objects to load batches of the datasets\n",
        "    train_loader, val_loader, classes = get_data_loader(batch_size)\n",
        "    ########################################################################\n",
        "    # Define the Loss function and optimizer\n",
        "    # The loss function will be Binary Cross Entropy (BCE). In this case we\n",
        "    # will use the BCEWithLogitsLoss which takes unnormalized output from\n",
        "    # the neural network and scalar label.\n",
        "    # Optimizer will be SGD with Momentum.\n",
        "    criterion = nn.BCEWithLogitsLoss()\n",
        "    optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9)\n",
        "    #optimizer = optim.Adam(net.parameters(), lr=learning_rate)\n",
        "    ########################################################################\n",
        "    # Set up some numpy arrays to store the training/test loss/erruracy\n",
        "    train_err = np.zeros(num_epochs)\n",
        "    train_loss = np.zeros(num_epochs)\n",
        "    val_err = np.zeros(num_epochs)\n",
        "    val_loss = np.zeros(num_epochs)\n",
        "    ########################################################################\n",
        "    # Train the network\n",
        "    # Loop over the data iterator and sample a new batch of training data\n",
        "    # Get the output from the network, and optimize our loss function.\n",
        "    start_time = time.time()\n",
        "    for epoch in range(num_epochs):  # loop over the dataset multiple times\n",
        "        total_train_loss = 0.0\n",
        "        total_train_err = 0.0\n",
        "        total_epoch = 0\n",
        "        for i, data in enumerate(train_loader, 0):\n",
        "            # Get the inputs\n",
        "            inputs, labels = data\n",
        "            labels = np.array([[1.0 if c == label else 0.0 for c in classes] for label in labels])\n",
        "            # Zero the parameter gradients\n",
        "            optimizer.zero_grad()\n",
        "            # Forward pass, backward pass, and optimize\n",
        "            outputs = net(inputs)\n",
        "            loss = criterion(outputs, torch.tensor(labels))\n",
        "            #loss = criterion(outputs, labels.float())\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            # Calculate the statistics\n",
        "\n",
        "            probabilities = F.softmax(outputs, dim=0)\n",
        "            corr = (torch.max(probabilities,dim=1)[1]).squeeze().long() != torch.max(torch.tensor(labels),dim=1)[1]\n",
        "            total_train_err += int(corr.sum())\n",
        "\n",
        "            total_train_loss += loss.item()\n",
        "            total_epoch += len(labels)\n",
        "        train_err[epoch] = float(total_train_err) / total_epoch\n",
        "        train_loss[epoch] = float(total_train_loss) / (i+1)\n",
        "        val_err[epoch], val_loss[epoch] = evaluate(net, val_loader, criterion)\n",
        "        print((\"Epoch {}: Train err: {}, Train loss: {} |\"+\n",
        "               \"Validation err: {}, Validation loss: {}\").format(\n",
        "                   epoch + 1,\n",
        "                   train_err[epoch],\n",
        "                   train_loss[epoch],\n",
        "                   val_err[epoch],\n",
        "                   val_loss[epoch]))\n",
        "        # Save the current model (checkpoint) to a file\n",
        "        model_path = get_model_name(net.name, batch_size, learning_rate, epoch)\n",
        "        torch.save(net.state_dict(), model_path)\n",
        "    print('Finished Training')\n",
        "    end_time = time.time()\n",
        "    elapsed_time = end_time - start_time\n",
        "    print(\"Total time elapsed: {:.2f} seconds\".format(elapsed_time))\n",
        "    # Write the train/test loss/err into CSV file for plotting later\n",
        "    epochs = np.arange(1, num_epochs + 1)\n",
        "    np.savetxt(\"{}_train_err.csv\".format(model_path), train_err)\n",
        "    np.savetxt(\"{}_train_loss.csv\".format(model_path), train_loss)\n",
        "    np.savetxt(\"{}_val_err.csv\".format(model_path), val_err)\n",
        "    np.savetxt(\"{}_val_loss.csv\".format(model_path), val_loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "tLNaNZYtY1Li"
      },
      "outputs": [],
      "source": [
        "def get_model_name(name, batch_size, learning_rate, epoch):\n",
        "    \"\"\" Generate a name for the model consisting of all the hyperparameter values\n",
        "\n",
        "    Args:\n",
        "        config: Configuration object containing the hyperparameters\n",
        "    Returns:\n",
        "        path: A string with the hyperparameter name and value concatenated\n",
        "    \"\"\"\n",
        "    path = \"model_{0}_bs{1}_lr{2}_epoch{3}\".format(name,\n",
        "                                                   batch_size,\n",
        "                                                   learning_rate,\n",
        "                                                   epoch)\n",
        "    return path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "7EfgowS7Y4Qx"
      },
      "outputs": [],
      "source": [
        "def plot_training_curve(path):\n",
        "    \"\"\" Plots the training curve for a model run, given the csv files\n",
        "    containing the train/validation error/loss.\n",
        "\n",
        "    Args:\n",
        "        path: The base path of the csv files produced during training\n",
        "    \"\"\"\n",
        "    import matplotlib.pyplot as plt\n",
        "    train_err = np.loadtxt(\"{}_train_err.csv\".format(path))\n",
        "    val_err = np.loadtxt(\"{}_val_err.csv\".format(path))\n",
        "    train_loss = np.loadtxt(\"{}_train_loss.csv\".format(path))\n",
        "    val_loss = np.loadtxt(\"{}_val_loss.csv\".format(path))\n",
        "    plt.title(\"Train vs Validation Error\")\n",
        "    n = len(train_err) # number of epochs\n",
        "    plt.plot(range(1,n+1), train_err, label=\"Train\")\n",
        "    plt.plot(range(1,n+1), val_err, label=\"Validation\")\n",
        "    plt.xlabel(\"Epoch\")\n",
        "    plt.ylabel(\"Error\")\n",
        "    plt.legend(loc='best')\n",
        "    plt.show()\n",
        "    plt.title(\"Train vs Validation Loss\")\n",
        "    plt.plot(range(1,n+1), train_loss, label=\"Train\")\n",
        "    plt.plot(range(1,n+1), val_loss, label=\"Validation\")\n",
        "    plt.xlabel(\"Epoch\")\n",
        "    plt.ylabel(\"Loss\")\n",
        "    plt.legend(loc='best')\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 211
        },
        "id": "ceOdekq5ZIsq",
        "outputId": "33e1fb10-fa2a-4677-acb0-28368ff20a7d"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "ZeroDivisionError",
          "evalue": "division by zero",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mZeroDivisionError\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-12-7dfef87d6aec>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;36m1\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mlarge_net\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLargeNet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#reset weights\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mtrain_net\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlarge_net\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnum_epochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m20\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m32\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlearning_rate\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.005\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;31m#large_model_path = get_model_name(\"large\",batch_size=64,learning_rate=0.01,epoch=19) #best so far (not anymore)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;31m#large_model_path = get_model_name(\"large\",batch_size=64,learning_rate=0.008,epoch=19) #random crop first go\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mZeroDivisionError\u001b[0m: division by zero"
          ]
        }
      ],
      "source": [
        "large_net = LargeNet() #reset weights\n",
        "train_net(large_net,num_epochs=20,batch_size=32, learning_rate=0.005)\n",
        "#large_model_path = get_model_name(\"large\",batch_size=64,learning_rate=0.01,epoch=19) #best so far (not anymore)\n",
        "#large_model_path = get_model_name(\"large\",batch_size=64,learning_rate=0.008,epoch=19) #random crop first go\n",
        "large_model_path = get_model_name(\"large\",batch_size=32,learning_rate=0.005,epoch=19)\n",
        "#large_model_path = get_model_name(\"large\",batch_size=64,learning_rate=0.01,epoch=19)\n",
        "plot_training_curve(large_model_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TLrYLHNnZJKi"
      },
      "outputs": [],
      "source": [
        "%debug"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class_counts = {'Plastic':0,'Metal':0,'Glass':0,'Paper':0}\n",
        "for i, label in enumerate(main_data_y):\n",
        "  #if i == 140:\n",
        "    #break\n",
        "  class_counts[label]+=1\n",
        "\n",
        "print(class_counts)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BcBuXVxEROhR",
        "outputId": "ed56beff-14db-4cae-b140-4d5974af0ed2"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'Plastic': 561, 'Metal': 480, 'Glass': 371, 'Paper': 437}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_loader, val_loader, c = get_data_loader(1)"
      ],
      "metadata": {
        "id": "qxzmTw4PUisq"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class_counts = {'Plastic':0,'Metal':0,'Glass':0,'Paper':0}\n",
        "for i, data in enumerate(val_loader, 0):\n",
        "  inputs, labels = data\n",
        "  for label in labels:\n",
        "    class_counts[label]+=1\n",
        "print(class_counts)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jw3ITEBRRQwa",
        "outputId": "35f17375-1c07-4338-8b35-459799bd6d4f"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'Plastic': 119, 'Metal': 89, 'Glass': 74, 'Paper': 88}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class_counts = {'Plastic':0,'Metal':0,'Glass':0,'Paper':0}\n",
        "for i, data in enumerate(train_loader, 0):\n",
        "  inputs, labels = data\n",
        "  for label in labels:\n",
        "    class_counts[label]+=1\n",
        "print(class_counts)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M4u8SeLPUWkr",
        "outputId": "ed4d9148-2e36-4f85-d2c8-2177ba26a728"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'Plastic': 442, 'Metal': 391, 'Glass': 297, 'Paper': 349}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class ann(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(ann,self).__init__()\n",
        "    self.layer1 = nn.Linear(3*128*128,500)\n",
        "    self.layer2 = nn.Linear(500,80)\n",
        "    self.layer3 = nn.Linear(80,4)\n",
        "def forward(self,img):\n",
        "  flattened = img.view(-1, 3*128*128)\n",
        "  activation1 = F.relu(self.layer1(flattened))\n",
        "  activation2 = F.relu(self.layer2(activation1))\n",
        "  output = self.layer3(activation2)\n",
        "  return output"
      ],
      "metadata": {
        "id": "1XAd6nQOZPEx"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "large_net = LargeNet() #reset weights\n",
        "for param in large_net.parameters():\n",
        "    print(param.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iDKQ8wq7YnN9",
        "outputId": "4002fd4c-36ef-47b2-f004-dd0d0b0cee7e"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([5, 3, 5, 5])\n",
            "torch.Size([5])\n",
            "torch.Size([10, 5, 3, 3])\n",
            "torch.Size([10])\n",
            "torch.Size([32, 9000])\n",
            "torch.Size([32])\n",
            "torch.Size([4, 32])\n",
            "torch.Size([4])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ann_net = ann() #reset weights\n",
        "for param in large_net.parameters():\n",
        "    print(param.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V4UVRSpuZLV3",
        "outputId": "35c42b44-37f6-4932-8f2a-441da183fd8c"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([5, 3, 5, 5])\n",
            "torch.Size([5])\n",
            "torch.Size([10, 5, 3, 3])\n",
            "torch.Size([10])\n",
            "torch.Size([32, 9000])\n",
            "torch.Size([32])\n",
            "torch.Size([4, 32])\n",
            "torch.Size([4])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_parameters = filter(lambda p: p.requires_grad, large_net.parameters())\n",
        "params = sum([np.prod(p.size()) for p in model_parameters])\n",
        "params"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4wCN6ejmZG4r",
        "outputId": "99078393-c72f-4096-fcb8-3eeacaa1f09e"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "289004"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_parameters = filter(lambda p: p.requires_grad, ann_net.parameters())\n",
        "params = sum([np.prod(p.size()) for p in model_parameters])\n",
        "params"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8y6Gsg0zZIEK",
        "outputId": "dd201cb3-cdfc-4c27-8eae-33aa8b100878"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "24616904"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}