Merge pull request #16 from igorastashov/hw_3

Hw 3

.gitignore

@@ -7,6 +7,8 @@ __pycache__/
 
 # Folder
 **/PokemonData
+/runs
 
-# Weights
-/weights/model.pt
+# Weights and optimizer
+weights/model.pt
+weights/optimizer.pt

README.md

@@ -21,10 +21,10 @@ from [kaggle](https://www.kaggle.com/datasets/lantian773030/pokemonclassificatio
 
 ## (2) Quick start
 
-### Download model weights
+### Download model and optimizer
 
 ```
-# Download model weights
+# Download model and optimizer
 cd weights
 bash download_weights.sh
 cd ../..
@@ -59,17 +59,33 @@ All of them with the Pokémon in the center.
 Most (if not all) of the images have relatively high quality (correct labels, centered).
 The images don't have extremely high resolutions so it's perfect for some light deep learning.
 
-**If the script doesn't work, an alternative will be to download the zip files manually
+If the script doesn't work, an alternative will be to download the zip files manually
 from the [link](https://www.kaggle.com/datasets/lantian773030/pokemonclassification/download?datasetVersionNumber=1).
-One can place the dataset zip files in `data`, respectively, and then unzip the zip file to set everything up.**
+One can place the dataset zip files in `data`, respectively, and then unzip the zip file to set everything up.
+
+**PAY ATTENTION**
+
+**This repository runs Data Version Control (DVC) for training and validation data.
+Pre-configured Google Drive remote storage stores raw input data.**
+
+```console
+$ dvc remote list
+my_remote gdrive://1RXz3Mv7OxVveHtQ7c1ZtGgazDh6bPFJz
+```
+
+You can run `dvc pull` to download the data:
+
+```console
+$ dvc pull
+```
 
 ## (4) Train and Evaluation model
 
 Example script to train and evaluate model.
 
 ```
-# Train MobileNet_V2
-python train.py
+# Train ConvNet
+python main.py
 ```
 
 ## (A) Acknowledgments

datasets/dataset.py → ds/dataset.py

@@ -1,5 +1,7 @@
 import glob
 import os
+from pathlib import Path
+from typing import Any
 
 import torchvision.transforms as T
 from PIL import Image
@@ -8,16 +10,47 @@
 from tqdm import tqdm
 
 
-bad_images = glob.glob("data/PokemonData/*/*.svg")
-for bad_image in bad_images:
-    os.remove(bad_image)
+def remove_bed_images(data_dir: Path):
+    bad_images = glob.glob(f"{data_dir}/*/*.svg")
+    for bad_image in bad_images:
+        os.remove(bad_image)
 
 
-class PokemonDataset(Dataset):
+TRAIN_NORMALIZED_MEAN = [0.485, 0.456, 0.406]
+TRAIN_NORMALIZED_STDEV = [0.229, 0.224, 0.225]
+NORMALIZE = T.Normalize(mean=TRAIN_NORMALIZED_MEAN, std=TRAIN_NORMALIZED_STDEV)
+
+
+def prepare_train_data():
+    """Add Augmentations"""
+    train_transform = T.Compose(
+        [
+            T.RandomResizedCrop(size=224, scale=(0.7, 1.0)),
+            T.RandomHorizontalFlip(),
+            T.ToTensor(),
+            NORMALIZE,
+        ]
+    )
+    return train_transform
+
+
+def prepare_test_data():
+    test_transform = T.Compose(
+        [
+            T.Resize(size=256),
+            T.CenterCrop(size=224),
+            T.ToTensor(),
+            NORMALIZE,
+        ]
+    )
+    return test_transform
+
+
+class PokemonDataset(Dataset[Any]):
     SPLIT_RANDOM_SEED = 42
     TEST_SIZE = 0.25
 
-    def __init__(self, root, train=True, load_to_ram=True, transform=None):
+    def __init__(self, root, train=True, load_to_ram=False, transform=None):
         super().__init__()
         self.root = root
         self.train = train
@@ -75,35 +108,40 @@ def __getitem__(self, item):
         return image, label
 
 
-def create_dataloader():
-    normalize = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
-    test_transform = T.Compose(
-        [
-            T.Resize(256),
-            T.CenterCrop(224),
-            T.ToTensor(),
-            normalize,
-        ]
-    )
-
+def create_dataloader(
+    root: Path,
+    batch_size: int,
+    load_to_ram: bool = False,
+    pin_memory: bool = True,
+    num_workers: int = 2,
+) -> tuple[DataLoader[Any], DataLoader[Any]]:
     train_dataset = PokemonDataset(
-        root="data/PokemonData",
+        root=root,
         train=True,
-        load_to_ram=False,
-        transform=test_transform,
+        load_to_ram=load_to_ram,
+        transform=prepare_train_data(),
     )
+
     test_dataset = PokemonDataset(
-        root="data/PokemonData",
+        root=root,
         train=False,
-        load_to_ram=False,
-        transform=test_transform,
+        load_to_ram=load_to_ram,
+        transform=prepare_test_data(),
     )
 
     train_loader = DataLoader(
-        train_dataset, batch_size=32, shuffle=True, pin_memory=True, num_workers=4
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        pin_memory=pin_memory,
+        num_workers=num_workers,
     )
     test_loader = DataLoader(
-        test_dataset, batch_size=32, shuffle=False, pin_memory=True, num_workers=4
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=pin_memory,
+        num_workers=num_workers,
     )
-    return train_dataset, test_dataset, train_loader, test_loader
+
+    return train_loader, test_loader

ds/models.py

@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+
+
+class ConvNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.network = nn.Sequential(
+            nn.Conv2d(
+                in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1
+            ),  # 16 x 224 x 224
+            nn.BatchNorm2d(16),
+            nn.ReLU(),
+            nn.Conv2d(
+                in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1
+            ),  # 32 x 224 x 224
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=28),  # 32 x 8 x 8
+            nn.Flatten(),  # 32*8*8 = 2048
+            nn.Linear(2048, 150),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)

ds/runner.py

@@ -0,0 +1,101 @@
+import torch
+import torch.nn as nn
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+from tqdm.notebook import tqdm
+
+from ds.tracking import plot_losses
+
+
+def training_epoch(
+    model: nn.Module,
+    optimizer: Optimizer,
+    criterion,
+    train_loader,
+    device: torch.device,
+    tqdm_desc,
+):
+    train_loss, train_accuracy = 0.0, 0.0
+    model.train()
+
+    for _, (images, labels) in enumerate(tqdm(train_loader, desc=tqdm_desc), 1):
+        images = images.to(device)  # images: batch_size x num_channels x height x width
+        labels = labels.to(device)  # labels: batch_size
+
+        optimizer.zero_grad()
+        logits = model(images)  # logits: batch_size x num_classes
+        loss = criterion(logits, labels)
+        loss.backward()
+        optimizer.step()
+
+        train_loss += loss.item() * images.shape[0]
+        accuracy = (logits.argmax(dim=1) == labels).sum().item()
+        train_accuracy += accuracy
+
+    train_loss /= len(train_loader.dataset)
+    train_accuracy /= len(train_loader.dataset)
+    return train_loss, train_accuracy
+
+
+@torch.no_grad()
+def validation_epoch(
+    model: nn.Module, criterion, test_loader: DataLoader, device: torch.device, tqdm_desc
+):
+    test_loss, test_accuracy = 0.0, 0.0
+    model.eval()
+    for images, labels in tqdm(test_loader, desc=tqdm_desc):
+        images = images.to(device)  # images: batch_size x num_channels x height x width
+        labels = labels.to(device)  # labels: batch_size
+        logits = model(images)  # logits: batch_size x num_classes
+        loss = criterion(logits, labels)
+
+        test_loss += loss.item() * images.shape[0]
+        test_accuracy += (logits.argmax(dim=1) == labels).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+    test_accuracy /= len(test_loader.dataset)
+    return test_loss, test_accuracy
+
+
+def train(
+    model: object,
+    optimizer: object,
+    scheduler: object,
+    criterion: object,
+    train_loader: object,
+    test_loader: object,
+    num_epochs: object,
+    device: object,
+    title: object,
+) -> object:
+    train_losses, train_accuracies = [], []
+    test_losses, test_accuracies = [], []
+
+    for epoch in range(1, num_epochs + 1):
+        train_loss, train_accuracy = training_epoch(
+            model,
+            optimizer,
+            criterion,
+            train_loader,
+            device,
+            tqdm_desc=f'Training {epoch}/{num_epochs}',
+        )
+        test_loss, test_accuracy = validation_epoch(
+            model,
+            criterion,
+            test_loader,
+            device,
+            tqdm_desc=f'Validating {epoch}/{num_epochs}',
+        )
+
+        if scheduler is not None:
+            scheduler.step()
+
+        train_losses += [train_loss]
+        train_accuracies += [train_accuracy]
+        test_losses += [test_loss]
+        test_accuracies += [test_accuracy]
+
+    plot_losses(train_losses, test_losses, train_accuracies, test_accuracies, title)
+
+    return train_losses, test_losses, train_accuracies, test_accuracies

ds/tracking.py

@@ -0,0 +1,36 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+from IPython.display import clear_output
+
+
+sns.set_style('darkgrid')
+plt.rcParams.update({'font.size': 15})
+
+LOG_PATH = "runs"
+
+
+def plot_losses(
+    train_losses: list[float],
+    test_losses: list[float],
+    train_accuracies: list[float],
+    test_accuracies: list[float],
+    title: str,
+):
+    clear_output()
+    fig, axs = plt.subplots(1, 2, figsize=(13, 4))
+    axs[0].plot(range(1, len(train_losses) + 1), train_losses, label='train')
+    axs[0].plot(range(1, len(test_losses) + 1), test_losses, label='test')
+    axs[0].set_ylabel('loss')
+    axs[0].set_title(title + ' loss')
+
+    axs[1].plot(range(1, len(train_accuracies) + 1), train_accuracies, label='train')
+    axs[1].plot(range(1, len(test_accuracies) + 1), test_accuracies, label='test')
+    axs[1].set_ylabel('accuracy')
+    axs[1].set_title(title + ' accuracy')
+
+    for ax in axs:
+        ax.set_xlabel('epoch')
+        ax.legend()
+
+    plt.savefig(f'{LOG_PATH}/{title}.png')
+    plt.show()