text_plain_output_1.png image_output_1.png
from torch import nn

from torchvision.datasets import ImageFolder

from torchvision.models import resnet18

from torchvision.transforms import Compose, Normalize, Resize, ToTensor

from tqdm.auto import tqdm

import numpy as np

import sys

import torch

from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, Normalize, Resize, ToTensor
dataset = ImageFolder('/kaggle/input/the-office-characters/Office Dataset/', transform=Compose([Resize((224, 224)), ToTensor(), Normalize((0.5, 0.5, 0.5), (1, 1, 1))]))
train_set, test_set = torch.utils.data.random_split(dataset, [int(0.8 * len(dataset)), len(dataset) - int(0.8 * len(dataset))])

train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=256, shuffle=True)
test_dataloader = torch.utils.data.DataLoader(test_set, batch_size=256, shuffle=False)

from torchvision.models import resnet18
model = resnet18(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
model.fc = nn.Linear(512, 6)

def train_epoch(model, data_loader, optimizer, criterion, return_losses=False, device='cuda:0'):
    model = model.to(device).train()
    total_loss = 0
    num_batches = 0
    all_losses = []
    total_predictions = np.array([])
    total_labels = np.array([])
    with tqdm(total=len(data_loader), file=sys.stdout) as prbar:
        for images, labels in data_loader:
            images = images.to(device)
            labels = labels.to(device)
            predicted = model(images)
            loss = criterion(predicted, labels)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            accuracy = (predicted.argmax(1) == labels).float().mean()
            prbar.set_description(f'Loss: {round(loss.item(), 4)} Accuracy: {round(accuracy.item() * 100, 4)}')
            prbar.update(1)
            total_loss += loss.item()
            total_predictions = np.append(total_predictions, predicted.argmax(1).cpu().detach().numpy())
            total_labels = np.append(total_labels, labels.cpu().detach().numpy())
            num_batches += 1
            all_losses.append(loss.detach().item())
    metrics = {'loss': total_loss / num_batches}
    metrics.update({'accuracy': (total_predictions == total_labels).mean()})
    if return_losses:
        return (metrics, all_losses)
    else:
        return metrics

def validate(model, data_loader, criterion, device='cuda:0'):
    model = model.eval()
    total_loss = 0
    num_batches = 0
    total_predictions = np.array([])
    total_labels = np.array([])
    with tqdm(total=len(data_loader), file=sys.stdout) as prbar:
        for images, labels in data_loader:
            images = images.to(device)
            labels = labels.to(device)
            predicted = model(images)
            loss = criterion(predicted, labels)
            accuracy = (predicted.argmax(1) == labels).float().mean()
            prbar.set_description(f'Loss: {round(loss.item(), 4)} Accuracy: {round(accuracy.item() * 100, 4)}')
            prbar.update(1)
            total_loss += loss.item()
            total_predictions = np.append(total_predictions, predicted.argmax(1).cpu().detach().numpy())
            total_labels = np.append(total_labels, labels.cpu().detach().numpy())
            num_batches += 1
    metrics = {'loss': total_loss / num_batches}
    metrics.update({'accuracy': (total_predictions == total_labels).mean()})
    return metrics

def fit(model, epochs, train_data_loader, validation_data_loader, optimizer, criterion, device='cuda:0'):
    all_train_losses = []
    epoch_train_losses = []
    epoch_eval_losses = []
    for epoch in range(epochs):
        train_metrics, one_epoch_train_losses = train_epoch(model=model, data_loader=train_data_loader, optimizer=optimizer, return_losses=True, criterion=criterion, device=device)
        all_train_losses.extend(one_epoch_train_losses)
        epoch_train_losses.append(train_metrics['loss'])
        with torch.no_grad():
            validation_metrics = validate(model=model, data_loader=validation_data_loader, criterion=criterion)
        epoch_eval_losses.append(validation_metrics['loss'])

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.fc.parameters(), 0.0001)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

fit(model, 30, train_dataloader, test_dataloader, optimizer, criterion, device=device)