application_vnd.jupyter.stderr_output_2.png text_html_output_1.png text_plain_output_1.png
from PIL import Image

from pathlib import Path

from pytorch_lightning.callbacks import ModelCheckpoint

from pytorch_lightning.loggers import WandbLogger

from pytorch_lightning.trainer import Trainer

from torch import nn

from torch.utils.data import Dataset, DataLoader

from torchmetrics import Accuracy

from torchvision import transforms

import os

import pytorch_lightning as pl

import torch

import torchvision

import wandb

img_path = Path('../input/celeba-dataset/img_align_celeba/img_align_celeba')
csv_path = Path('../input/celeba-dataset/list_attr_celeba.csv')

class LitModel(pl.LightningModule):
    def __init__(self, n_classes, download_pretrained=True, wandb_logger=None, **kwargs):
        super().__init__()
        self.model = torchvision.models.resnet18(pretrained=True)
        self.model.fc = nn.Linear(512, n_classes)
        self.train_accuracy = Accuracy()
        self.val_accuracy = Accuracy()
        self.test_accuracy = Accuracy()
        self.loss = nn.CrossEntropyLoss()
        self.wandb_logger = wandb_logger
        self.opt_params = {'lr': 0.001}
    def forward(self, x):
        return self.model(x)
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat_logits = self.forward(x)
        loss = self.loss(y_hat_logits, y)
        train_acc = self.train_accuracy(y_hat_logits, y)
        self.log('train_acc', train_acc, prog_bar=True)
        self.log('train_loss', loss)
        return {'loss': loss, 'progress_bar': {'train_acc': train_acc}}
    def training_epoch_end(self, outputs):
        avg_train_loss = torch.stack([x['loss'] for x in outputs]).mean()
        self.log('train_avg_loss', avg_train_loss, on_epoch=True)
        self.log('train_acc_epoch', self.train_accuracy, on_epoch=True)
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat_logits = self.forward(x)
        self.val_accuracy(y_hat_logits, y)
        val_loss = self.loss(y_hat_logits, y)
        return {'val_loss': val_loss, 'out_logits': y_hat_logits}
    def on_validation_batch_end(self, outputs, batch, batch_idx, dataloader_idx) -> None:
        if not self.wandb_logger:
            return
        if batch_idx == 0:
            outputs = outputs['out_logits']
            n = 20
            x, y = batch
            images = [img for img in x[:n]]
            assert len(outputs.shape) == 2
            assert outputs.shape[-1] == 2
            preds = outputs[:n].softmax(dim=-1)
            pred_prob, pred_class = preds.max(dim=-1)
            columns = ['image', 'pred_class', 'true_class', 'pred_prob']
            label_to_text = {0: 'Female', 1: 'Male'}
            td = []
            for image, y_pred_class, y_real, y_pred_prob in zip(x[:n], pred_class, y[:n], pred_prob):
                pred_class = label_to_text[y_pred_class.item()]
                true_class = label_to_text[y_real.item()]
                td.append([wandb.Image(image), pred_class, true_class, y_pred_prob])
            self.wandb_logger.log_table(key='samples', columns=columns, data=td)
    def validation_epoch_end(self, outputs):
        avg_val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        self.log('val_acc_epoch', self.val_accuracy, on_epoch=True, prog_bar=True)
        self.log('val_loss_epoch', avg_val_loss, on_epoch=True, prog_bar=True)
        return {'val_loss': avg_val_loss}
    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat_logits = self.forward(x)
        self.test_accuracy(y_hat_logits, y)
        test_loss = self.loss(y_hat_logits, y)
        self.log('test_loss', test_loss, on_epoch=True)
        self.log('test_acc', self.test_accuracy, on_epoch=True)
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), **self.opt_params)
        return optimizer

wandb_logger = WandbLogger(project='gender-detection-vit')

model = LitModel(2, wandb_logger=wandb_logger)

class CelebADataset(Dataset):
    def __init__(self, image_attr_ordered_map, base_image_dir: str, transforms):
        self.image_attr_ordered_map = image_attr_ordered_map
        self.base_image_dir = base_image_dir
        self.transforms = transforms
    def __len__(self):
        return len(self.image_attr_ordered_map)
    def __getitem__(self, index):
        img_name, labels = self.image_attr_ordered_map[index]
        image_path = os.path.join(self.base_image_dir, img_name)
        image = Image.open(image_path).convert('RGB')
        tensor_x = self.transforms(image)
        return (tensor_x, labels)

single_image_path = img_path / '000002.jpg'
orig_image = Image.open(single_image_path).convert('RGB')
orig_image = transforms.ToTensor()(orig_image)
orig_image.shape

image_transforms = {'train': transforms.Compose([transforms.RandomHorizontalFlip(), transforms.ToTensor()]), 'val': transforms.Compose([transforms.Resize(size=(218, 178)), transforms.ToTensor()])}

data_mod = CelebADataModule(csv_path=csv_path, image_dir=img_path, train_transforms=image_transforms['train'], val_transforms=image_transforms['val'], batch_size=32, num_workers=3)

checkpoint_callback = ModelCheckpoint(monitor='val_loss_epoch')

trainer = Trainer(gpus=1, max_epochs=6, logger=wandb_logger, callbacks=[checkpoint_callback])

trainer.fit(model, data_mod)

trainer.validate(model, data_mod)

best_model_path = trainer.checkpoint_callback.best_model_path
print(f'Best model {best_model_path}')
model = LitModel.load_from_checkpoint(best_model_path, n_classes=2)