<jupyter_start><jupyter_text>**Bootstrap Your Own Latent A New Approach to Self-Supervised Learning:** https://arxiv.org/pdf/2006.07733.pdf<jupyter_code># !pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
# !pip install -qqU fastai fastcore
# !pip install nbdev
import fastai, fastcore, torch
fastai.__version__ , fastcore.__version__, torch.__version__
from fastai.vision.all import *<jupyter_output><empty_output><jupyter_text>SizesResize -> RandomCrop320 -> 256 | 224 -> 192 | 160 -> 128<jupyter_code>resize = 320
size = 256<jupyter_output><empty_output><jupyter_text>1. Implementation Details (Section 3.2 from the paper)  1.1 Image AugmentationsSame as SimCLR with optional grayscale<jupyter_code>import kornia
def get_aug_pipe(size, stats=imagenet_stats, s=.6):
    "SimCLR augmentations"
    rrc = kornia.augmentation.RandomResizedCrop((size, size), scale=(0.2, 1.0), ratio=(3/4, 4/3))
    rhf = kornia.augmentation.RandomHorizontalFlip()
    rcj = kornia.augmentation.ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s)
    rgs = kornia.augmentation.RandomGrayscale(p=0.2)
    
    tfms = [rrc, rhf, rcj, rgs, Normalize.from_stats(*stats)]
    pipe = Pipeline(tfms)
    pipe.split_idx = 0
    return pipe<jupyter_output><empty_output><jupyter_text>1.2 Architecture<jupyter_code>def create_encoder(arch, n_in=3, pretrained=True, cut=None, concat_pool=True):
    "Create encoder from a given arch backbone"
    encoder = create_body(arch, n_in, pretrained, cut)
    pool = AdaptiveConcatPool2d() if concat_pool else nn.AdaptiveAvgPool2d(1)
    return nn.Sequential(*encoder, pool, Flatten())
class MLP(Module):
    "MLP module as described in paper"
    def __init__(self, dim, projection_size=256, hidden_size=2048):
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_size, projection_size)
        )

    def forward(self, x):
        return self.net(x)
class BYOLModel(Module):
    "Compute predictions of v1 and v2" 
    def __init__(self,encoder,projector,predictor):
        self.encoder,self.projector,self.predictor = encoder,projector,predictor    

    def forward(self,v1,v2):
        q1 = self.predictor(self.projector(self.encoder(v1)))
        q2 = self.predictor(self.projector(self.encoder(v2)))
        return (q1,q2)
def create_byol_model(arch=resnet50, hidden_size=4096, pretrained=True, projection_size=256, concat_pool=False):
    encoder = create_encoder(arch, pretrained=pretrained, concat_pool=concat_pool)
    with torch.no_grad(): 
        x = torch.randn((2,3,128,128))
        representation = encoder(x)

    projector = MLP(representation.size(1), projection_size, hidden_size=hidden_size)     
    predictor = MLP(projection_size, projection_size, hidden_size=hidden_size)
    apply_init(projector)
    apply_init(predictor)
    return BYOLModel(encoder, projector, predictor)<jupyter_output><empty_output><jupyter_text>1.3 BYOLCallback<jupyter_code>def _mse_loss(x, y):
    x = F.normalize(x, dim=-1, p=2)
    y = F.normalize(y, dim=-1, p=2)
    return 2 - 2 * (x * y).sum(dim=-1)

def symmetric_mse_loss(pred, *yb):
    (q1,q2),z1,z2 = pred,*yb
    return (_mse_loss(q1,z2) + _mse_loss(q2,z1)).mean()
x = torch.randn((64,256))
y = torch.randn((64,256))
test_close(symmetric_mse_loss((x,y),y,x), 0) # perfect
test_close(symmetric_mse_loss((x,y),x,y), 4, 1e-1) # random<jupyter_output><empty_output><jupyter_text>Useful Discussions and Supportive Material:- https://www.reddit.com/r/MachineLearning/comments/hju274/d_byol_bootstrap_your_own_latent_cheating/fwohtky/- https://untitled-ai.github.io/understanding-self-supervised-contrastive-learning.html<jupyter_code>import copy

class BYOLCallback(Callback):
    "Implementation of https://arxiv.org/pdf/2006.07733.pdf"
    def __init__(self, T=0.99, debug=True, size=224, **aug_kwargs):        
        self.T, self.debug = T, debug
        self.aug1 = get_aug_pipe(size, **aug_kwargs)
        self.aug2 = get_aug_pipe(size, **aug_kwargs)


    def before_fit(self):
        "Create target model"
        self.target_model = copy.deepcopy(self.learn.model).to(self.dls.device)        
        self.T_sched = SchedCos(self.T, 1) # used in paper
        # self.T_sched = SchedNo(self.T, 1) # used in open source implementation
  
        
    def before_batch(self):
        "Generate 2 views of the same image and calculate target projections for these views"
        if self.debug: print(f"self.x[0]: {self.x[0]}")
        
        v1,v2 = self.aug1(self.x), self.aug2(self.x.clone())
        self.learn.xb = (v1,v2)
        
        if self.debug:
            print(f"v1[0]: {v1[0]}\nv2[0]: {v2[0]}")
            self.show_one()
            assert not torch.equal(*self.learn.xb)

        with torch.no_grad():
            z1 = self.target_model.projector(self.target_model.encoder(v1))
            z2 = self.target_model.projector(self.target_model.encoder(v2))
            self.learn.yb = (z1,z2)


    def after_step(self):
        "Update target model and T"
        self.T = self.T_sched(self.pct_train)
        with torch.no_grad():
            for param_k, param_q in zip(self.target_model.parameters(), self.model.parameters()):
                param_k.data = param_k.data * self.T + param_q.data * (1. - self.T)
          

    def show_one(self):
        b1 = self.aug1.normalize.decode(to_detach(self.learn.xb[0]))
        b2 = self.aug1.normalize.decode(to_detach(self.learn.xb[1]))
        i = np.random.choice(len(b1))
        show_images([b1[i],b2[i]], nrows=1, ncols=2)

    def after_train(self):      
        if self.debug: self.show_one()

    def after_validate(self):   
        if self.debug: self.show_one()<jupyter_output><empty_output><jupyter_text>2. Pretext Training<jupyter_code>sqrmom=0.99
mom=0.95
beta=0.
eps=1e-4
opt_func = partial(ranger, mom=mom, sqr_mom=sqrmom, eps=eps, beta=beta)
bs=128
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=0.9)], 
            [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
dls = get_dls(resize, bs)
model = create_byol_model(arch=xresnet34, pretrained=False)
learn = Learner(dls, model, symmetric_mse_loss, opt_func=opt_func,
                cbs=[BYOLCallback(T=0.99, size=size, debug=False), TerminateOnNaNCallback()])
learn.to_fp16();
learn.lr_find()
lr=1e-3
wd=1e-2
epochs=100
learn.unfreeze()
learn.fit_flat_cos(epochs, lr, wd=wd, pct_start=0.5)
save_name = f'byol_iwang_sz{size}_epc{epochs}'
learn.save(save_name)
torch.save(learn.model.encoder.state_dict(), learn.path/learn.model_dir/f'{save_name}_encoder.pth')
learn.load(save_name);
lr=1e-4
wd=1e-2
epochs=100
learn.unfreeze()
learn.fit_flat_cos(epochs, lr, wd=wd, pct_start=0.5)
save_name = f'byol_iwang_sz{size}_epc200'
learn.save(save_name)
torch.save(learn.model.encoder.state_dict(), learn.path/learn.model_dir/f'{save_name}_encoder.pth')
lr=1e-4
wd=1e-2
epochs=30
learn.unfreeze()
learn.fit_flat_cos(epochs, lr, wd=wd, pct_start=0.5)
save_name = f'byol_iwang_sz{size}_epc230'
learn.save(save_name)
torch.save(learn.model.encoder.state_dict(), learn.path/learn.model_dir/f'{save_name}_encoder.pth')
lr=5e-5
wd=1e-2
epochs=30
learn.unfreeze()
learn.fit_flat_cos(epochs, lr, wd=wd, pct_start=0.5)
save_name = f'byol_iwang_sz{size}_epc260'
learn.save(save_name)
torch.save(learn.model.encoder.state_dict(), learn.path/learn.model_dir/f'{save_name}_encoder.pth')
learn.recorder.plot_loss()
save_name<jupyter_output><empty_output><jupyter_text>3. Downstream Task - Image Classification<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
def do_train(epochs=5, runs=5,  lr=2e-2, size=size, bs=bs, save_name=None):
    dls = get_dls(size, bs)
    for run in range(runs):
        print(f'Run: {run}')
        learn = cnn_learner(dls, xresnet34, opt_func=opt_func, normalize=False,
                metrics=[accuracy,top_k_accuracy], loss_func=LabelSmoothingCrossEntropy(),
                pretrained=False)
#         learn.to_fp16()
        
        if save_name is not None:
            state_dict = torch.load(learn.path/learn.model_dir/f'{save_name}_encoder.pth')
            learn.model[0].load_state_dict(state_dict)
            print("Model loaded...")
            
        learn.unfreeze()
        learn.fit_flat_cos(epochs, lr, wd=wd)<jupyter_output><empty_output><jupyter_text>ImageWang Leaderboard **sz-256****Contrastive Learning**- 5 epochs:  67.70%- 20 epochs: 70.03%- 80 epochs: 70.71%- 200 epochs: 71.78%**BYOL**- 5 epochs: 64.74%- 20 epochs: **71.01%**- 80 epochs: **72.58%**- 200 epochs: **72.13%**  5 epochs<jupyter_code># we are using old pretrained model with size 192 for transfer learning
# link: https://github.com/KeremTurgutlu/self_supervised/blob/252269827da41b41091cf0db533b65c0d1312f85/nbs/byol_iwang_192.ipynb
save_name = 'byol_iwang_sz192_epc230'
lr = 1e-2
wd=1e-2
bs=128
epochs = 5
runs = 5
do_train(epochs, runs, lr=lr, bs=bs, save_name=save_name)
np.mean([0.657165,0.637312,0.631967,0.646729,0.664291])<jupyter_output><empty_output><jupyter_text>20 epochs<jupyter_code>lr=2e-2
epochs = 20
runs = 3
do_train(epochs, runs, lr=lr, save_name=save_name)
np.mean([0.711631, 0.705269, 0.713413])<jupyter_output><empty_output><jupyter_text>80 epochs<jupyter_code>epochs = 80
runs = 1
do_train(epochs, runs, save_name=save_name)<jupyter_output>Run: 0
Model loaded...<jupyter_text>200 epochs<jupyter_code>epochs = 200
runs = 1
do_train(epochs, runs, save_name=save_name)<jupyter_output>Run: 0
Model loaded...
{'fraction_non_alphanumeric': 0.1371005973262539, 'fraction_numerical': 0.044846875888878354, 'mean_word_length': 4.150390625, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 51, 'https://': 6, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '6013', 'license': 'Apache-2.0', 'n_tokens_mistral': 4345, 'n_tokens_neox': 4051, 'n_words': 662, 'path': 'nbs/examples/byol_iwang_256.ipynb', 'repo_name': 'adbmd/self_supervised', 'size': '10547'}