text_plain_output_3.png text_plain_output_2.png text_plain_output_1.png
from kaggle_secrets import UserSecretsClient

from kaggle_secrets import UserSecretsClient

from torch import nn

from torchmetrics import Accuracy

import pytorch_lightning as pl

import torch

import torchvision

import wandb

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

from kaggle_secrets import UserSecretsClient
secret_label = 'wandb'
wandb_key = UserSecretsClient().get_secret(secret_label)

from kaggle_secrets import UserSecretsClient
secret_label = 'wandb'
secret_value = UserSecretsClient().get_secret(secret_label)

wandb.login(key=secret_value)

wandb.finish()