Final commit

cnn.py

-import pytorch_lightning as pl
-import torch
-from torch import nn, optim
-import torch.nn.functional as F
-import torchmetrics
-
-
-def conv2d_size_out(size, kernel_size=5, stride=1):
-    """
-    Determine the output size after applying convolution operation.
-    """
-    return (size - (kernel_size - 1) - 1) // stride + 1
-
-
-class CNN(pl.LightningModule):
-    def __init__(self, img_size, conv_filters, conv_kernels, conv_strides):
-        super().__init__()
-
-        self.fc_input_size_h = img_size
-        self.fc_input_size_w = img_size
-
-        # Conv layers
-        self.conv_layers = []
-        for i in range(len(conv_filters) - 1):
-            self.conv_layers.append(nn.Conv2d(in_channels=conv_filters[i],
-                                              out_channels=conv_filters[i + 1],
-                                              kernel_size=conv_kernels[i],
-                                              stride=conv_strides[i]))
-            self.conv_layers.append(nn.BatchNorm2d(conv_filters[i + 1]))
-            self.conv_layers.append(nn.ReLU())
-            self.fc_input_size_h = conv2d_size_out(self.fc_input_size_h, conv_kernels[i], 2)
-            self.fc_input_size_w = conv2d_size_out(self.fc_input_size_w, conv_kernels[i], 2)
-
-        self.conv_layers = nn.Sequential(*self.conv_layers)
-
-        # Fully connected layers
-        fc_layers = [self.fc_input_size_h * self.fc_input_size_w * conv_filters[-1], 100]
-        self.fc_layers = []
-
-        for i in range(len(fc_layers) - 1):
-            self.fc_layers.append(nn.Linear(fc_layers[i], fc_layers[i + 1]))
-            self.fc_layers.append(nn.ReLU() if i != len(fc_layers) - 2 else nn.LogSoftmax(dim=0))
-
-        self.fc_layers = nn.Sequential(nn.Flatten(), *self.fc_layers)
-
-    def forward(self, x):
-        x = self.conv_layers(x)
-        x = self.fc_layers(x)
-        return x
-
-    def configure_optimizers(self):
-        return optim.Adam(self.parameters(), lr=1e-3)
-
-    def common_step(self, batch, batch_id):
-        images, targets = batch
-        images = images.to(self.device).permute(0, 3, 1, 2).float()
-        target = targets.to(self.device).long()
-        output = self(images)
-        loss = F.nll_loss(output, target)
-        f1 = torchmetrics.functional.f1(output, target, average='macro', num_classes=100)
-        return loss, f1
-
-    def training_step(self, batch, batch_id):
-        loss, f1 = self.common_step(batch, batch_id)
-        self.log('train/loss', loss)
-        self.log('train/f1', f1)
-        return loss
-
-    @torch.no_grad()
-    def validation_step(self, batch, batch_id):
-        loss, f1 = self.common_step(batch, batch_id)
-        self.log('val/loss', loss)
-        self.log('val/f1', f1)
-        return loss

train.py

-from typing import Tuple
-from pathlib import Path
-import logging
-
-import pandas as pd
-from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.loggers import WandbLogger
-from skimage.io import imread
-import albumentations as A
-import torch
-from torch.utils.data import Dataset, DataLoader
-import pytorch_lightning as pl
-
-from cnn import CNN
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-
-
-class PetDataset(Dataset):
-    def __init__(self, image_paths, df, transforms) -> None:
-        self.path_names = image_paths
-        self.targets = None
-        if 'Pawpularity' in df.columns:
-            self.targets = df[df.index.isin([p.stem for p in self.path_names])]['Pawpularity'].values
-        self.transforms = transforms
-
-    def __len__(self) -> int:
-        return len(self.path_names)
-
-    def __getitem__(self, index: int) -> Tuple:
-        image = imread(self.path_names[index])
-        image = self.transforms(image=image)['image']
-        if self.targets is None:
-            return image
-
-        return image, self.targets[index] - 1
-
-
-def train():
-    LOAD_PATH = Path('data')
-
-    logger = WandbLogger('petfinder_cnn', project='petfinder', log_model='all')
-
-    train_df = pd.read_csv(Path(LOAD_PATH, 'train.csv'), index_col='Id')
-    test_df = pd.read_csv(Path(LOAD_PATH, 'test.csv'), index_col='Id')
-
-    train_image_paths = list(Path(LOAD_PATH, 'train').glob('*'))
-    test_image_paths = list(Path(LOAD_PATH, 'test').glob('*'))
-
-    img_size = 64
-
-    conv_filters = [3, 8, 16, 32, 64]
-    conv_kernels = [5, 5, 5, 5]
-    conv_strides = [2, 2, 2, 2]
-
-    train_transforms = A.Compose([
-        A.Resize(img_size, img_size)
-    ])
-
-    val_transforms = A.Compose([
-        A.Resize(img_size, img_size)
-    ])
-
-    train_dataset = PetDataset(train_image_paths, train_df, train_transforms)
-    test_dataset = PetDataset(test_image_paths, test_df, val_transforms)
-
-    val_size = int(len(train_dataset) * 0.8)
-    train_set, val_set = torch.utils.data.random_split(train_dataset, lengths=[len(train_dataset) - val_size, val_size])
-
-    batch_size = 128
-    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)
-    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=0)
-    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
-
-    model = CNN(img_size, conv_filters, conv_kernels, conv_strides)
-
-    callback = ModelCheckpoint(
-        monitor='val/f1',
-        mode='max',
-        save_top_k=5,
-        filename=CNN.__name__
-    )
-
-    trainer = pl.Trainer(
-        gpus=1,
-        max_epochs=150,
-        logger=logger,
-        callbacks=[callback]
-    )
-
-    trainer.fit(model, train_loader, val_loader)
-    # prediction = trainer.predict(model, test_loader)
-    # torch.save(prediction, 'sub.pt')
-    logger.finalize('ok')
-
-
-if __name__ == '__main__':
-    train()