<jupyter_start><jupyter_text>Buliding Flexible Code BaseWe will build a flexible code base that increases the ease of development and code-reuse. I will take bits and pieces of FASTAI library and use this as a base for our experiments.  Let's setup dummy dataset<jupyter_code>from sklearn import datasets

noise = 0.1
X_train, y_train = datasets.make_moons(n_samples=2000, noise=noise)
y_train
#create dummy dataset
import torch
from torch.utils.data import Dataset,DataLoader
from torch.nn import Sequential
import torch.nn as nn
from tqdm import tqdm
class DummyDset(Dataset):
    def __init__(self,sz,noise=0.1):
        self.x,self.y = datasets.make_moons(n_samples=sz,noise=noise)
        self.x = self.x.astype('float32')
    def __getitem__(self,ind):
        return self.x[ind],self.y[ind]
    def __len__(self):
        return len(self.x)
        
        
class Dloaders:
    def __init__(self,*dls):
        self.train,self.valid = dls
sampl = DummyDset(100)
for x,y in sampl:
    print(f'{x},{x.shape}')
    print(f'{y},{y.shape}')
    break
trn_dset,vld_dset,tst_dset = DummyDset(2000),DummyDset(500),DummyDset(500)
trn_dl = DataLoader(trn_dset,shuffle=True,batch_size=50)
vld_dl = DataLoader(vld_dset,batch_size=50)
dls = Dloaders(trn_dl,vld_dl)
for x_bl,y_bl in trn_dl:
    print(x_bl.shape)
    print(y_bl.shape)
    break
dmodel = Sequential(nn.Linear(2,164),
          nn.LeakyReLU(),
          nn.Linear(164,32),
          nn.LeakyReLU(),
          nn.Linear(32,2),
          )
dmodel
loss_fn = nn.CrossEntropyLoss()
optimizer_fn = torch.optim.Adam
# define Python user-defined exceptions
class CancelFitException(Exception):
    pass
xb,yb = None,None
for xb,yb in trn_dl:
    break
    
out = dmodel(xb)<jupyter_output><empty_output><jupyter_text>Flexible Trainer Class<jupyter_code>class Trainer:
    def __init__(self,model,dls,loss_func,lr,opt_func,cbs,device='cpu'):
        self.model = model
        self.dl = dls
        self.loss_func = loss_func
        self.lr = lr
        self.opt_func = opt_func
        self.cbs =cbs
        self.device = device
        for cb in cbs:cb.trainer = self
        
    def one_batch(self,batch):
        self('before_batch')
        xb,yb = batch
        self.opt.zero_grad() 
        #storing y_preds onto the class for a potential use by a callback
        self.y_preds = self.model(xb)
        self.loss = self.loss_func(self.y_preds,yb)
        if self.model.training:
            #take backward step only if in training mode
            self.loss.backward()
            self.opt.step()
        self('after_batch')
        
    def one_epoch(self,is_train):
        self.model.training = is_train
        self('before_epoch')
        #choose the appropriate data-loader
        dl = self.dl.train if is_train else self.dl.valid
        for self.epochnum,self.batch in enumerate(dl):
            #import pdb;pdb.set_trace()
            self.one_batch(self.batch)
        self('after_epoch')
    
    def fit(self,nepochs):
        self('before_fit')
        self.opt = self.opt_func(self.model.parameters(),self.lr)
        self.nepochs = nepochs
        try:
            for self.epoch in tqdm(range(self.nepochs)):
                self.one_epoch(True)
                self.one_epoch(False)
        except CancelFitException:
            pass
        self('after_fit')
        
    def predict(self,dl):
        print('NOT IMPLEMENTED')
        pass
    
    def infer(self,x):
        pass
    
    def __call__(self,name):
        for cb in self.cbs:getattr(cb,name)()
        
        
class CallBack:
    def before_epoch(self):
        pass
    def after_epoch(self):
        pass
    def before_batch(self):
        pass
    def after_batch(self):
        pass
    def before_fit(self):
        pass
    def after_fit(self): 
        pass
    
    
    

class SetupTrainerCB(CallBack):
    def before_batch(self):
        xb,yb = to_device(self.trainer.device,self.trainer.batch)
        self.trainer.batch = xb,yb

    def before_fit(self): 
        if self.trainer.device != 'cpu':
            self.trainer.model.cuda()
class PrinterCB(CallBack):
    def before_epoch(self):
        self.ns,self.losses,self.accs = [],[],[]
        
    def after_epoch(self):
        n = sum(self.ns)
        epoch_type = 'Training' if self.trainer.model.training else 'Validation'
        acc = sum(self.accs)/n
        loss = sum(self.losses)/n
        print(f"{epoch_type} acc: {acc},loss:{loss}")
        
    def after_batch(self):
        xb,yb = self.trainer.batch
        acc = (self.trainer.y_preds.argmax(dim=-1) == yb).float().sum()
        ns = len(xb)
        l = (self.trainer.loss.item())*ns
        self.ns.append(ns)
        self.losses.append(l)
        self.accs.append(acc)
    
cbs = [PrinterCB()]
trainer = Trainer(dmodel,dls,loss_fn,0.0002,optimizer_fn,cbs,'cpu')
trainer.fit(10)<jupyter_output>100%|███████████████████████████████████████████| 10/10 [00:00<00:00, 41.63it/s]<jupyter_text>Using tools<jupyter_code>import os
import sys
import torch
import numpy as np
import torch.nn as nn
import torch.backends.cudnn as cudnn
import matplotlib.pyplot as plt
from torch.nn import functional as F

import pytorch_lightning as pl


from pathlib import Path
import matplotlib.pyplot as plt
#import pandas as pd
import numpy as np

cwd = os.getcwd()
module_path = "/".join(cwd.split('/')[0:-1])
if module_path not in sys.path:
    sys.path.append(module_path)
from src.models import vgg16


def accuracy(y_hat,y):
    #import pdb;pdb.set_trace()
    return torch.mean((y_hat == y).type(torch.float32))
#configuring pytorch-lightning with wandb
import pytorch_lightning as pl
str_model = {'dmodel':dmodel}

hparams = {'model':'dmodel','optimizer':{'lr':0.0002}}
class Model(pl.LightningModule):
    def __init__(self,hparams):
        super().__init__()
        self.save_hyperparameters(hparams)
        self.model = str_model[hparams['model']]
        #import pdb;pdb.set_trace()
    def forward(self,x):
        #? write for reading off any model specific hparams before the call(hydra)
        x = self.model(x)
        return x
    
    def training_step(self,batch,batch_idx):
        x,y = batch
        y_hat = self.model(x)
        #y_true = y.squeeze().type(torch.LongTensor)
        loss = F.cross_entropy(y_hat,y)
        
        acc = torch.mean((y_hat.argmax(dim=-1) == y).type(torch.float32))
        #import pdb;pdb.set_trace()
        #print(f'batch indx:{batch_idx},{loss.item()}')
        self.log("train_acc", acc, on_step=False, on_epoch=True)
        self.log("train_loss", loss,on_step=False, on_epoch=True)
        
        
        return loss
        
    
    def validation_step(self,batch,batch_idx):
        x,y = batch
        y_hat = self(x)
        
        #y_hat = torch.argmax(y_hat, dim=1)
        #import pdb;pdb.set_trace()
        acc = torch.mean((y_hat.argmax(dim=-1) == y).type(torch.float32))
        self.log('val_acc', acc,on_step=False, on_epoch=True)
        return acc

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams['optimizer']['lr'])
        return optimizer

        
        

pl.seed_everything(0)
#model
net = Model(hparams)

from pytorch_lightning.loggers import CSVLogger
logger = CSVLogger("logs", name="temp")
trainer = pl.Trainer(max_epochs=10,logger=logger)
trainer.fit(net,trn_dl,vld_dl)


import pandas as pd

log_dir = '/Users/vinay/Projects/uncertainty-estimates/nbs/logs/temp/version_4/'
met = pd.read_csv(log_dir+'metrics.csv')
met
class Model(pl.LightningModule):
    def __init__(self,hparams):
        super().__init__()
        self.save_hyperparameters(hparams)
        self.model = str_model[hparams['model']]
        #import pdb;pdb.set_trace()
    def forward(self,x):
        #? write for reading off any model specific hparams before the call(hydra)
        x = self.model(x)
        return x
    
    def training_step(self,batch,batch_idx):
        x,y = batch
        y_hat = self.model(x)
        #y_true = y.squeeze().type(torch.LongTensor)
        loss = F.cross_entropy(y_hat,y)
        
        acc = torch.mean((y_hat.argmax(dim=-1) == y).type(torch.float32))
        #import pdb;pdb.set_trace()
        #print(f'batch indx:{batch_idx},{loss.item()}')
        self.log("train_acc", acc, on_step=False, on_epoch=True)
        self.log("train_loss", loss)
        
        
        return loss
        
    
    def validation_step(self,batch,batch_idx):
        x,y = batch
        y_hat = self(x)
        
        #y_hat = torch.argmax(y_hat, dim=1)
        #import pdb;pdb.set_trace()
        acc = torch.mean((y_hat.argmax(dim=-1) == y).type(torch.float32))
        self.log('val_acc', acc,on_step=False, on_epoch=True)
        return acc

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams['optimizer']['lr'])
        return optimizer

        
        
## Integrating wandb


import wandb


from pytorch_lightning.loggers import WandbLogger

wandb.login()


wandb_logger = WandbLogger(project="temp-tools")
from pytorch_lightning.loggers import CSVLogger
csvlogger = CSVLogger("logs", name="temp")



pl.seed_everything(0)
#model
net = Model(hparams)




trainer = pl.Trainer(
    logger=[wandb_logger,csvlogger],    # W&B integration
    log_every_n_steps=50,   # set the logging frequency                
    max_epochs=10,           # number of epochs
    deterministic=True,     # keep it deterministic
    )
trainer.fit(net,trn_dl,vld_dl)



import pandas as pd
log_dir = '/Users/vinay/Projects/uncertainty-estimates/nbs/logs/temp/version_5/'
met = pd.read_csv(log_dir+'metrics.csv')
met<jupyter_output><empty_output><jupyter_text>wandb sweeps<jupyter_code>import wandb

sweep_config = {
  "name" : "my-sweep",
  "method" : "random",
  "parameters" : {
    "epochs" : {
      "values" : [10, 20, 50]
    },
    "learning_rate" :{
      "min": 0.0001,
      "max": 0.1
    }
  }
}

sweep_id = wandb.sweep(sweep_config)
wandb.config
'/'.join(os.getcwd().split('/')[:-1])
from torch import as_tensor,Tensor,ByteTensor,LongTensor,FloatTensor,HalfTensor,DoubleTensor
from types import GeneratorType

def is_listy(x):
    "`isinstance(x, (tuple,list,slice,GeneratorType))`"
    return isinstance(x, (tuple,list,slice,GeneratorType))



def apply(func, x, *args, **kwargs):
    "Apply `func` recursively to `x`, passing on args"
    if is_listy(x): return type(x)([apply(func, o, *args, **kwargs) for o in x])
    if isinstance(x,dict):  return {k: apply(func, v, *args, **kwargs) for k,v in x.items()}
    res = func(x, *args, **kwargs)
    return res 

def to_device(b, device=None, non_blocking=False):
    "Recursively put `b` on `device`."
    #if defaults.use_cuda==False: device='cpu'
    #elif device is None: device=default_device()
    def _inner(o):
        if isinstance(o,Tensor): return o.to(device, non_blocking=non_blocking)
#         if hasattr(o, "to_device"): return o.to_device(device)
        return o
    return apply(_inner, b)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12867082316797687, 'fraction_numerical': 0.01039125327550375, 'mean_word_length': 2.8726382085374387, 'pattern_counts': {'":': 2, '<': 17, '<?xml version=': 0, '>': 16, 'https://': 0, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '56750', 'license': 'MIT', 'n_tokens_mistral': 4030, 'n_tokens_neox': 3920, 'n_words': 760, 'path': 'nbs/tools.ipynb', 'repo_name': 'vin136/uncertainty-aware', 'size': '11067'}