<jupyter_start><jupyter_text><jupyter_code>#hide
#skip
! [ -e /content ] && pip install -Uqq self-supervised
! pip install transcv nbdev fastai -q --upgrade
# hide
# all_slow<jupyter_output><empty_output><jupyter_text>A tutorial for SimCLR > In this tutorial, we will pre-train a ViT model, using SimCLR, and then, we shall fine-tune the model for a downstream visual recognition task. For a detailed and thorough explanation, it is recommended to first go through the tutorial notebooks of [this library](https://keremturgutlu.github.io/self_supervised/), and then this notebook.  **Imports**<jupyter_code>from fastai.vision.all import *
from transcv.visrectrans import *
torch.backends.cudnn.benchmark = True
from self_supervised.augmentations import *
from self_supervised.layers import *
from self_supervised.vision.simclr import *<jupyter_output><empty_output><jupyter_text>**Datasets & Dataloaders for self-supervised pre-training**<jupyter_code>def get_dls(size, bs, workers=None):
    path = URLs.IMAGEWANG_160 if size <= 160 else URLs.IMAGEWANG
    source = untar_data(path)
    
    files = get_image_files(source)
    tfms = [[PILImage.create, ToTensor, RandomResizedCrop(size, min_scale=1.)], 
            [parent_label, Categorize()]]
    
    dsets = Datasets(files, tfms=tfms, splits=RandomSplitter(valid_pct=0.1)(files))
    
    batch_tfms = [IntToFloatTensor]
    dls = dsets.dataloaders(bs=bs, num_workers=workers, after_batch=batch_tfms)
    return dls
bs, resize, size = 8, 256, 224
dls = get_dls(resize, bs)<jupyter_output><empty_output><jupyter_text>**ViT model**Getting the ViT model, using the `VisRecTrans` class :<jupyter_code>vis_rec_ob = VisRecTrans('vit_small_patch16_224', 0, False)
model = vis_rec_ob.create_model()
vis_rec_ob.initialize(model)<jupyter_output>/usr/local/lib/python3.7/dist-packages/torch/nn/init.py:388: UserWarning: Initializing zero-element tensors is a no-op
  warnings.warn("Initializing zero-element tensors is a no-op")
/usr/local/lib/python3.7/dist-packages/torch/nn/init.py:426: UserWarning: Initializing zero-element tensors is a no-op
  warnings.warn("Initializing zero-element tensors is a no-op")<jupyter_text>We need a module for extracting the class token, from the feature representation computed by the ViT :<jupyter_code>class Encoder_transcv (Module) :
  def __init__ (self) :
    pass 
  def forward (self, x) : 
    x = x[:, 0]
    return x
encoder = nn.Sequential(model[0], model[1], model[2], Encoder_transcv())
def create_simclr_model(encoder, hidden_size=256, projection_size=128, bn=False, nlayers=2):
    "Create SimCLR model"
    n_in  = in_channels(encoder)
    with torch.no_grad(): representation = encoder(torch.randn((2,n_in,224,224)))
    projector = create_mlp_module(representation.size(1), hidden_size, projection_size, bn=bn, nlayers=nlayers)
    apply_init(projector)
    return nn.Sequential(encoder, projector)
model = create_simclr_model(encoder)
aug_pipelines = get_simclr_aug_pipelines(size, rotate=True, rotate_deg=10, jitter=True, bw=True, blur=False) 
cbs=[SimCLR(aug_pipelines)]
learn = Learner(dls, model, cbs=cbs)
b = dls.one_batch()
learn._split(b)
learn('before_batch')
learn.sim_clr.show(n=5);
learn.to_fp16();
lr,wd,epochs=1e-2,1e-2,100
learn.lr_find()<jupyter_output><empty_output><jupyter_text>Since we have a different model for SimCLR, we need a custom `embed_callback` callback :<jupyter_code>class TrainEmbedCallback (Callback) :
  def before_train (self) :
    self.learn.model[0][1].training = True
    self.learn.model[0][1].requires_grad_(True)
  def before_validation (self) :
    self.learn.model[0][1].training = False
    self.learn.model[0][1].requires_grad_(False)
learn.summary()<jupyter_output><empty_output><jupyter_text>**Self-supervised pre-training**Here, we pre-train the ViT model for just 1 epoch, since the training time for even 1 epoch is significantly high. It is recommended to pre-train the ViT model for 100 epochs (according to [this tutorial](https://keremturgutlu.github.io/self_supervised/01%20-%20training_simclr_iwang.html)).<jupyter_code>learn.unfreeze()
learn.fit_flat_cos(1, 1e-4, wd=wd, pct_start=0.5, cbs = [TrainEmbedCallback])<jupyter_output><empty_output><jupyter_text>Saving the parameters of the encoder :<jupyter_code>save_name = f'simclr_iwang_sz{size}_epc{epochs}'
learn.save(save_name)
encoder_path = learn.path/learn.model_dir/f'{save_name}_encoder.pth'
torch.save(learn.model[0].state_dict(), encoder_path)
bs, size<jupyter_output><empty_output><jupyter_text>**Datasets & Dataloaders for fine-tuning**<jupyter_code>def get_dls(size, bs, workers=None):
    path = URLs.IMAGEWANG_160 if size <= 160 else URLs.IMAGEWANG
    source = untar_data(path)
    files = get_image_files(source, folders=['train', 'val'])
    splits = GrandparentSplitter(valid_name='val')(files)
    
    item_aug = [RandomResizedCrop(size, min_scale=0.35), FlipItem(0.5)]
    tfms = [[PILImage.create, ToTensor, *item_aug], 
            [parent_label, Categorize()]]
    
    dsets = Datasets(files, tfms=tfms, splits=splits)
    
    batch_tfms = [IntToFloatTensor, Normalize.from_stats(*imagenet_stats)]
    dls = dsets.dataloaders(bs=bs, num_workers=workers, after_batch=batch_tfms)
    return dls<jupyter_output><empty_output><jupyter_text>**Getting the learner for fine-tuning**<jupyter_code>def split_func(m): return L(m[0], m[1]).map(params)

def create_learner(size=size, encoder_path="models/swav_iwang_sz128_epc100_encoder.pth"):
    
    dls = get_dls(size, bs=bs)
    pretrained_encoder = torch.load(encoder_path)
    vis_rec_ob = VisRecTrans('vit_small_patch16_224', 0, False)
    model = vis_rec_ob.create_model()
    vis_rec_ob.initialize(model)
    encoder = nn.Sequential(model[0], model[1], model[2], Encoder_transcv())
    encoder.load_state_dict(pretrained_encoder)
    nf = encoder(torch.randn(2,3,224,224)).size(-1)
    classifier = create_cls_module(nf, dls.c)
    model = nn.Sequential(encoder, classifier)
    learn = Learner(dls, model, splitter=split_func,
                metrics=top_k_accuracy, loss_func=LabelSmoothingCrossEntropy())
    return learn
learn = create_learner(size, encoder_path)
learn.unfreeze()
learn.lr_find()<jupyter_output>/usr/local/lib/python3.7/dist-packages/torch/nn/init.py:388: UserWarning: Initializing zero-element tensors is a no-op
  warnings.warn("Initializing zero-element tensors is a no-op")
/usr/local/lib/python3.7/dist-packages/torch/nn/init.py:426: UserWarning: Initializing zero-element tensors is a no-op
  warnings.warn("Initializing zero-element tensors is a no-op")<jupyter_text>**Fine-tuning**<jupyter_code>learn.fit_one_cycle(5, 1e-4, cbs = TrainEmbedCallback)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12253436939629408, 'fraction_numerical': 0.020621637776449493, 'mean_word_length': 5.788032454361055, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 44, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '127699', 'license': 'Apache-2.0', 'n_tokens_mistral': 2506, 'n_tokens_neox': 2391, 'n_words': 466, 'path': 'A tutorial for SimCLR.ipynb', 'repo_name': 'Ys-Prakash/transcv', 'size': '6692'}