# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] papermill={"duration": 0.02177, "end_time": "2022-03-21T00:29:43.359044", "exception": false, "start_time": "2022-03-21T00:29:43.337274", "status": "completed"} tags=[]
#
#
# # Example: Train SimCLR on CIFAR10
#
# In this tutorial, we will train a SimCLR model using lightly. The model,
# augmentations and training procedure is from 
# `A Simple Framework for Contrastive Learning of Visual Representations <https://arxiv.org/abs/2002.05709>`_.
#
# The paper explores a rather simple training procedure for contrastive learning.
#
# -

# !nvidia-smi

# + [markdown] papermill={"duration": 0.020242, "end_time": "2022-03-21T00:29:43.401855", "exception": false, "start_time": "2022-03-21T00:29:43.381613", "status": "completed"} tags=[]
# ## Imports
#
# Import the Python frameworks we need for this tutorial.
#
#

# + papermill={"duration": 6.108637, "end_time": "2022-03-21T00:30:02.386483", "exception": false, "start_time": "2022-03-21T00:29:56.277846", "status": "completed"} tags=[]
import torch
import torch.nn as nn
import torchvision
import pytorch_lightning as pl
import lightly

from lightly.data import LightlyDataset
from torchvision.datasets import CIFAR10
from torch.utils.data import Subset
import sys


sys.path.append('../src')
from utils import get_classes, generate_embeddings, custom_collate_fn
from my_resnet import resnet20
# %matplotlib inline

# + [markdown] papermill={"duration": 0.02944, "end_time": "2022-03-21T00:30:02.445604", "exception": false, "start_time": "2022-03-21T00:30:02.416164", "status": "completed"} tags=[]
# ## Configuration
#
# We set some configuration parameters for our experiment.
# Feel free to change them and analyze the effect.
#
#

# + papermill={"duration": 0.044385, "end_time": "2022-03-21T00:30:02.519766", "exception": false, "start_time": "2022-03-21T00:30:02.475381", "status": "completed"} tags=[]
num_workers = 2
batch_size = 128
seed = 1

max_epochs = 150
input_size = 32 # image height, assume its always square

# Let's set the seed for our experiments
data_path = "../data/cifar10"
NORMAL_CLASS = 'dog'


pl.seed_everything(seed)

# + [markdown] papermill={"duration": 0.029471, "end_time": "2022-03-21T00:30:02.644733", "exception": false, "start_time": "2022-03-21T00:30:02.615262", "status": "completed"} tags=[]
# ## Setup data augmentations and loaders
#

# +
cifar10_train = CIFAR10(data_path, download=True, train=True)
cifar10_test = CIFAR10(data_path, download=False, train=False)

classes_ids_train = get_classes(cifar10_train) # long!
classes_ids_test = get_classes(cifar10_test)
# -

# `collate_fn` Additional augmentations such as vertical flip or random rotation (90 degrees). 
# By adding these augmentations we learn our model invariance regarding the 
# orientation of the images. E.g. we don't care if a shirt is upside down
# but more about the strcture which make it.

# + jupyter={"outputs_hidden": false} papermill={"duration": 2.54491, "end_time": "2022-03-21T00:30:05.219039", "exception": false, "start_time": "2022-03-21T00:30:02.674129", "status": "completed"} tags=[]
dataset_train = LightlyDataset.from_torch_dataset(Subset(cifar10_train, classes_ids_train[NORMAL_CLASS]))
dataset_test = LightlyDataset.from_torch_dataset(Subset(cifar10_train, classes_ids_train[NORMAL_CLASS]))

dataloader_train = torch.utils.data.DataLoader(dataset_train,
    batch_size=batch_size, shuffle=True, collate_fn=custom_collate_fn, drop_last=True, num_workers=num_workers)

dataloader_test = torch.utils.data.DataLoader(dataset_test,
                                              batch_size=batch_size, shuffle=False, drop_last=False, collate_fn=custom_collate_fn, num_workers=num_workers)

# + [markdown] papermill={"duration": 0.02912, "end_time": "2022-03-21T00:30:05.362705", "exception": false, "start_time": "2022-03-21T00:30:05.333585", "status": "completed"} tags=[]
# # 1. Create the SimCLR Model
# Now we create the SimCLR model. We implement it as a PyTorch Lightning Module
# and use custom ResNet-20 backbone provided by <NAME>. Lightly provides implementations
# of the SimCLR projection head and loss function in the `SimCLRProjectionHead`
# and `NTXentLoss` classes. We can simply import them and combine the building
# blocks in the module. We will import constructed model from our `src`.
#
#
# -

from simclr_model import SimCLRModel

# + [markdown] papermill={"duration": 0.030096, "end_time": "2022-03-21T00:30:05.493609", "exception": false, "start_time": "2022-03-21T00:30:05.463513", "status": "completed"} tags=[]
# We first check if a GPU is available and then train the module
# using the PyTorch Lightning Trainer.
#
#

# + jupyter={"outputs_hidden": false} papermill={"duration": 1559.208489, "end_time": "2022-03-21T00:56:04.731487", "exception": false, "start_time": "2022-03-21T00:30:05.522998", "status": "completed"} tags=[]
gpus = 1 if torch.cuda.is_available() else 0
device = "cuda" if torch.cuda.is_available() else "cpu"

resnet_backbone = resnet20(num_classes=1)
model = SimCLRModel(resnet_backbone, img_size = input_size)
trainer = pl.Trainer(
    max_epochs=50, gpus=gpus, progress_bar_refresh_rate=10
)
trainer.fit(model, dataloader_train)
# -

torch.save(model.state_dict(), '../weights/weights_simclr')

# # 2. Embeddings extraction
# Next we create a helper function to generate embeddings
# from our test images using the model we just trained.
# Note that only the backbone is needed to generate embeddings,
# the projection head is only required for the training.
# Make sure to put the model into eval mode for this part!
#
#

model.to(device)
model.eval()
embeddings_test, filenames_test = generate_embeddings(model, dataloader_test, device)
embeddings_train, filenames_train = generate_embeddings(model, dataloader_train, device)

print(f'Shape of TEST embeddings {embeddings_test.shape}')
print(f'Shape of TRAIN embeddings {embeddings_train.shape}')

# # 3. Calculate Hausdorff distance between poin clouds

# +
from scipy.spatial.distance import directed_hausdorff

hausdorff_dist = directed_hausdorff(embeddings_train, embeddings_test)[0]

print(f'Hausdorff Dist: {hausdorff_dist:.3f}')

StarcoderdataJupyterScriptsDedupFiltered