# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: PyCharm (celeb_a)
#     language: python
#     name: pycharm-b8c63a84
# ---

# +
from __future__ import print_function, division

import numpy as np
import matplotlib.pyplot as plt
# %matplotlib inline

import torch
import torch.nn as nn
import torch.nn.functional as F

import torchvision
from torchvision import models, transforms, datasets

import torch.optim as optim
from torch.optim import lr_scheduler

from torch.utils.data import DataLoader
from torch.utils.data import sampler

import time
import os
import copy


# hyper parameters
num_epochs = 5
num_classes = 2
batch_size = 8
learning_rate = 0.001
model_name = "resnet"

# +
import PIL.Image as Image

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((224, 224)),
    transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)
])


# +
# get dataset

train_dataset = datasets.CelebA(root='./data',
                                split='train',
                                target_type='attr',
                                transform=transform,
                                download=False)
val_dataset = datasets.CelebA(root='./data',
                                split='valid',
                                target_type='attr',
                                transform=transform,
                                download=False)
test_dataset = datasets.CelebA(root='./data',
                                split='test',
                                target_type='attr',
                                transform=transform,
                                download=False)

# data loader
train_loader = DataLoader(dataset=train_dataset,
                            batch_size=batch_size,
                            shuffle=True)
val_loader = DataLoader(dataset=val_dataset,
                            batch_size=batch_size,
                            shuffle=False)
test_loader = DataLoader(dataset=test_dataset,
                            batch_size=batch_size,
                            shuffle=False)

# -

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

"""unzip the dataset"""
# import zipfile
#
# with zipfile.ZipFile("img_align_celeba.zip","r") as zip_ref:
#     zip_ref.extractall("data/")

# + pycharm={"name": "#%%\n"}
def show(img,renorm=False,nrow=8,interpolation='bicubic'):
    if renorm:
        img = img/2 + 0.5
    img_grid = torchvision.utils.make_grid(img,nrow=nrow).numpy()
    plt.figure()
    plt.imshow(np.transpose(img_grid, (1,2,0)), interpolation=interpolation)
    plt.axis('off')
    plt.show()


# -

"""Visualize a few images"""
batch, classes = next(iter(train_loader))
show(batch, renorm = True, nrow=4)

batch.shape

# + pycharm={"name": "#%%\n"}
classes.shape

# + pycharm={"name": "#%%\n"}
model = models.resnet18(pretrained=True)
# freeze all params in the layers except for the final fc
for param in model.parameters():
    param.requires_grad = False

num_ftrs = model.fc.in_features
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
model.fc = nn.Linear(num_ftrs, 2)

model = model.to(device)

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

# + pycharm={"name": "#%%\n"}
for name, params in model.named_parameters():
      print(name, params.size())

# + pycharm={"name": "#%%\n"}
train_step = len(train_loader) # total number of `i`

for epoch in range(num_epochs): # repeat the entire training `num_epochs` times
    # for each training sample
    for i, (images, labels) in enumerate(train_loader):
        # move to gpu if available
        labels = labels[:, 2]
        images = images.to(device)
        labels = labels.to(device)
        # forward pass
        outputs = model(images)
        loss = criterion(outputs, labels) # still a tensor so we need to use .item() when printing
        # backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # print
        if (i+1) % 100 == 0:
            print('Epoch: [{}/{}], Step[{}/{}], Loss:{:.4f}' \
                  .format(epoch+1, num_epochs, i+1, train_step, loss.item()))

# + pycharm={"name": "#%%\n"}
# save model
torch.save(model.state_dict(), 'cnn.ckpt')

# load
# model = TheModelClass(*args, **kwargs)
# model.load_state_dict(torch.load(PATH))
# model.eval()

# + pycharm={"name": "#%%\n"}
model.eval() # BatchNorm uses moving mean/variance instead of mini-batch mean/variance

with torch.no_grad():
    # initialize the stats
    correct = 0
    total = 0
    # pass through testing data once
    for images, labels in val_loader:

        labels = labels[:, 2]
        # again move to device first
        images = images.to(device)
        labels = labels.to(device)
        # forward once
        outputs = model(images)
        # instead of calculating loss we will get predictions
        # it's essetially outputs just reformatting imo
        _, predicted = torch.max(outputs.data, 1)
        # accumulate stats
        total += labels.size(0) # yeah again, number of elements in the tensor
        correct += (labels == predicted).sum().item()

    # print
    print('Test accuracy on 10000 test images: {}%' \
            .format(100 * correct / total))

# + pycharm={"name": "#%%\n"}


StarcoderdataJupyterScriptsDedupFiltered