crashedice/signify
1
1# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_siamese.ipynb.2 3# %% auto 04__all__ = ['def_device', 'compare', 'with_cbs', 'run_cbs', 'to_cpu', 'to_device', 'CancelFitException', 'CancelBatchException',5 'CancelEpochException', 'Callback', 'MetricsCB', 'TrainCB', 'TrainContrastiveCB', 'TrainBceCB', 'DeviceCB',6 'ProgressCB', 'SaveModelCallback', 'LRFinderCB', 'HooksCallback', 'append_stats', 'ActivationStats',7 'LoadModelCallback', 'TwoDLVCallback', 'BatchTransformCB', 'BaseSchedCB', 'BatchSchedCB', 'rand_erase',8 'RandErase', 'SquareReflectPad', 'Learner', 'show_image', 'subplots', 'get_grid', 'show_images',9 'reshape_alternating']10 11# %% ../nbs/00_siamese.ipynb 412import pathlib, os, shutil, sys, cv2, torch, random, glob13 14import pandas as pd15import numpy as np16import matplotlib.pyplot as plt17import math18import statistics19 20from PIL import Image21from tqdm import tqdm22from pathlib import Path23from itertools import zip_longest24from copy import copy25from operator import attrgetter26from functools import partial27from collections.abc import Mapping28 29import torch.nn as nn30import torch.nn.functional as F31import torchvision.models as models32import torch.optim as optim33import torchvision.transforms as transforms34import torchvision35from torch.utils.data import DataLoader, Dataset36from torcheval.metrics import MulticlassAccuracy,Mean, BinaryAccuracy37from torch.nn import init38from fastprogress import progress_bar,master_bar39from torch.optim.lr_scheduler import ExponentialLR40 41# %% ../nbs/00_siamese.ipynb 842def compare(pic1, pic2): 43 return random.random()44 45# %% ../nbs/00_siamese.ipynb 1646class with_cbs:47 def __init__(self, nm): self.nm = nm48 def __call__(self, f):49 def _f(o, *args, **kwargs):50 try:51 o.callback(f'before_{self.nm}')52 f(o, *args, **kwargs)53 o.callback(f'after_{self.nm}')54 except globals()[f'Cancel{self.nm.title()}Exception']: pass55 finally: o.callback(f'cleanup_{self.nm}')56 return _f57 58 59def run_cbs(cbs, method_nm, learn=None):60 for cb in sorted(cbs, key=attrgetter('order')):61 method = getattr(cb, method_nm, None)62 if method is not None: method(learn)63 64# %% ../nbs/00_siamese.ipynb 1765def to_cpu(x):66 if isinstance(x, Mapping): return {k:to_cpu(v) for k,v in x.items()}67 if isinstance(x, list): return [to_cpu(o) for o in x]68 if isinstance(x, tuple): return tuple(to_cpu(list(x)))69 res = x.detach().cpu()70 return res.float() if res.dtype==torch.float16 else res71 72def_device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'73 74def to_device(x, device=def_device):75 if isinstance(x, list): return [to_device(o) for o in x]76 if isinstance(x, tuple): return tuple(to_device(list(x)))77 if isinstance(x, torch.Tensor): return x.to(device)78 79# %% ../nbs/00_siamese.ipynb 1880class CancelFitException(Exception): pass81class CancelBatchException(Exception): pass82class CancelEpochException(Exception): pass83 84# %% ../nbs/00_siamese.ipynb 1985class Callback(): order = 086 87class MetricsCB(Callback):88 def __init__(self, wandb, *ms, **metrics):89 for o in ms: metrics[type(o).__name__] = o90 self.metrics = metrics91 self.all_metrics = copy(metrics)92 self.all_metrics['loss'] = self.loss = Mean()93 self.wandb = wandb94 95 def _log(self, d): print(d)96 def before_fit(self, learn): learn.metrics = self97 def before_epoch(self, learn): [o.reset() for o in self.all_metrics.values()]98 99 def after_epoch(self, learn):100 log = {k:f'{v.compute():.3f}' for k,v in self.all_metrics.items()}101 log['epoch'] = learn.epoch102 log['train'] = 'train' if learn.model.training else 'eval'103 self._log(log)104 105 log_wandb = {f'{k}_{log["train"]}':round(float(v.compute()), 3) for k,v in self.all_metrics.items()}106 self.wandb.log(log_wandb)107 108 109 def after_batch(self, learn):110 x,x2, y = to_cpu(learn.batch)111 for m in self.metrics.values():112 m.update(to_cpu(learn.preds), y)113 114 loss = to_cpu(learn.loss)115 self.loss.update(loss, weight=len(x))116 self.wandb.log({"batch loss": loss.item()})117 118class TrainCB(Callback):119 def __init__(self, n_inp=1): self.n_inp = n_inp120 def predict(self, learn): 121 learn.x, learn.y = learn.batch[:self.n_inp], learn.batch[self.n_inp]122 learn.preds = learn.model(*learn.x)123 def get_loss(self, learn): 124 learn.loss = learn.loss_func(learn.preds, *learn.batch[self.n_inp:])125 def backward(self, learn): learn.loss.backward()126 def step(self, learn): learn.opt.step()127 def zero_grad(self, learn): learn.opt.zero_grad()128 129class TrainContrastiveCB(Callback):130 def __init__(self, n_inp=1): 131 self.n_inp = n_inp132 self.loss_func133 def predict(self, learn): 134 learn.preds = learn.model(*learn.batch[:self.n_inp])135 def get_loss(self, learn): 136 pred1, pred2 = learn.preds137 label = learn.batch[self.n_inp]138 learn.loss = learn.loss_func(pred1, pred2, label)139 def backward(self, learn): learn.loss.backward()140 def step(self, learn): learn.opt.step()141 def zero_grad(self, learn): learn.opt.zero_grad()142 143class TrainBceCB(Callback):144 def __init__(self, n_inp=1): 145 self.n_inp = n_inp146 self.loss_func = torch.nn.BCELoss()147 def predict(self, learn): 148 learn.x, learn.y = learn.batch[:self.n_inp], learn.batch[self.n_inp]149 learn.preds = learn.model(*learn.x)150 def get_loss(self, learn): 151 label = learn.batch[self.n_inp]152 learn.loss = self.loss_func(learn.preds, label)153 def backward(self, learn): learn.loss.backward()154 def step(self, learn): learn.opt.step()155 def zero_grad(self, learn): learn.opt.zero_grad()156 157class DeviceCB(Callback):158 order = 1159 def __init__(self, device): fc.store_attr()160 def before_fit(self, learn):161 if hasattr(learn.model, 'to'): learn.model.to(self.device)162 def before_batch(self, learn): 163 learn.batch = to_device(learn.batch, device=self.device)164 165class ProgressCB(Callback):166 order = MetricsCB.order+1167 def __init__(self, plot=False): self.plot = plot168 def before_fit(self, learn):169 learn.epochs = self.mbar = master_bar(learn.epochs)170 self.first = True171 if hasattr(learn, 'metrics'): learn.metrics._log = self._log172 self.losses = []173 self.val_losses = []174 175 176 def _log(self, d):177 if self.first:178 self.mbar.write(list(d), table=True)179 self.first = False180 self.mbar.write(list(d.values()), table=True)181 182 def before_epoch(self, learn): learn.dl = progress_bar(learn.dl, leave=False, parent=self.mbar)183 def after_batch(self, learn):184 learn.dl.comment = f'{learn.loss:.3f}'185 if self.plot and hasattr(learn, 'metrics') and learn.training:186 self.losses.append(learn.loss.item())187 if self.val_losses: self.mbar.update_graph([[fc.L.range(self.losses), self.losses],[fc.L.range(learn.epoch).map(lambda x: (x+1)*len(learn.dlt)), self.val_losses]])188 189 def after_epoch(self, learn): 190 if not learn.training:191 if self.plot and hasattr(learn, 'metrics'): 192 self.val_losses.append(learn.metrics.all_metrics['loss'].compute())193 self.mbar.update_graph([[fc.L.range(self.losses), self.losses],[fc.L.range(learn.epoch+1).map(lambda x: (x+1)*len(learn.dlt)), self.val_losses]])194 195 196class SaveModelCallback(Callback):197 "A `TrackerCallback` that saves the model's best during training and loads it at the end."198 order = ProgressCB.order + 1199 def __init__(self):200 201 try:202 old_loss = torch.load("model.pth")["loss"]203 except:204 old_loss = 1000205 206 self.valid_losses = [old_loss]207 self.valid_losses_batch = []208 209 def after_batch(self, learn):210 if not learn.training:211 self.valid_losses_batch.append(learn.loss.item())212 213 def after_epoch(self,learn):214 215 if not learn.training:216 217 current_valid_loss = statistics.mean(self.valid_losses_batch)218 prev_best = min(self.valid_losses)219 if current_valid_loss < prev_best:220 print(f"saving model in epoch {learn.epoch} with loss {current_valid_loss} (prev: {prev_best})")221 torch.save({222 'epoch': learn.epoch,223 'model_state_dict': learn.model.state_dict(),224 'optimizer_state_dict': learn.opt.state_dict(),225 'loss':current_valid_loss,226 }, "model.pth")227 self.valid_losses.append(current_valid_loss)228 self.valid_losses_batch = []229 230 231class LRFinderCB(Callback):232 order = 1233 def __init__(self, gamma=1.3, max_mult=3): fc.store_attr()234 235 def before_fit(self, learn):236 self.sched = ExponentialLR(learn.opt, self.gamma)237 self.lrs,self.losses = [],[]238 self.min = math.inf239 240 def after_batch(self, learn):241 if not learn.training: raise CancelEpochException()242 self.lrs.append(learn.opt.param_groups[0]['lr'])243 loss = to_cpu(learn.loss)244 self.losses.append(loss)245 if loss < self.min: self.min = loss246 if math.isnan(loss) or (loss > self.min*self.max_mult):247 raise CancelFitException()248 self.sched.step()249 250 def cleanup_fit(self, learn):251 plt.plot(self.lrs, self.losses)252 plt.xscale('log')253 254#| export255class HooksCallback(Callback):256 def __init__(self, hookfunc, mod_filter=fc.noop, on_train=True, on_valid=False, mods=None):257 fc.store_attr()258 super().__init__()259 260 def before_fit(self, learn):261 if self.mods: mods=self.mods262 else: mods = fc.filter_ex(learn.model.modules(), self.mod_filter)263 self.hooks = Hooks(mods, partial(self._hookfunc, learn))264 265 def _hookfunc(self, learn, *args, **kwargs):266 if (self.on_train and learn.training) or (self.on_valid and not learn.training): self.hookfunc(*args, **kwargs)267 268 def after_fit(self, learn): self.hooks.remove()269 def __iter__(self): return iter(self.hooks)270 def __len__(self): return len(self.hooks)271 272#| export273def append_stats(hook, mod, inp, outp):274 if not hasattr(hook,'stats'): hook.stats = ([],[],[])275 acts = to_cpu(outp)276 hook.stats[0].append(acts.mean())277 hook.stats[1].append(acts.std())278 hook.stats[2].append(acts.abs().histc(40,0,10))279 280#|export281class ActivationStats(HooksCallback):282 def __init__(self, mod_filter=fc.noop): super().__init__(append_stats, mod_filter)283 284 def color_dim(self, figsize=(11,5)):285 fig,axes = get_grid(len(self), figsize=figsize)286 for ax,h in zip(axes.flat, self):287 show_image(get_hist(h), ax, origin='lower')288 289 def dead_chart(self, figsize=(11,5)):290 fig,axes = get_grid(len(self), figsize=figsize)291 for ax,h in zip(axes.flatten(), self):292 ax.plot(get_min(h))293 ax.set_ylim(0,1)294 295 def plot_stats(self, figsize=(10,4)):296 fig,axs = plt.subplots(1,2, figsize=figsize)297 for h in self:298 for i in 0,1: axs[i].plot(h.stats[i])299 axs[0].set_title('Means')300 axs[1].set_title('Stdevs')301 plt.legend(fc.L.range(self))302 303 304 305class LoadModelCallback(Callback):306 order = 0307 def __init__(self, path):308 self.path = path309 310 def before_fit(self, learn):311 learn.model.load_state_dict(torch.load(self.path)["model_state_dict"])312 313class TwoDLVCallback(Callback):314 order = 0315 def __init__(self, dlv):316 self.dlv = dlv317 318 def after_epoch(self, learn):319 print("2nd valid")320 storedlearner = deepcopy(learn) #.copy()321 storedlearner.dlv = self.dlv322 torch.no_grad()(storedlearner._one_epoch)()323 324 325#| export326class BatchTransformCB(Callback):327 def __init__(self, tfm, on_train=True, on_val=True): fc.store_attr()328 329 def before_batch(self, learn):330 if (self.on_train and learn.training) or (self.on_val and not learn.training):331 learn.batch = self.tfm(learn.batch)332 333 334# %% ../nbs/00_siamese.ipynb 21335class BaseSchedCB(Callback):336 def __init__(self, sched): self.sched = sched337 def before_fit(self, learn): self.schedo = self.sched(learn.opt)338 def _step(self, learn):339 if learn.training: self.schedo.step()340#|export341class BatchSchedCB(BaseSchedCB):342 def after_batch(self, learn): self._step(learn)343 344# %% ../nbs/00_siamese.ipynb 23345def _rand_erase1(x, pct, xm, xs, mn, mx):346 szx = int(pct*x.shape[-2])347 szy = int(pct*x.shape[-1])348 stx = int(random.random()*(1-pct)*x.shape[-2])349 sty = int(random.random()*(1-pct)*x.shape[-1])350 init.normal_(x[:,:,stx:stx+szx,sty:sty+szy], mean=xm, std=xs)351 x.clamp_(mn, mx)352 353#|export354def rand_erase(x, pct=0.2, max_num = 4):355 xm,xs,mn,mx = x.mean(),x.std(),x.min(),x.max()356 num = random.randint(0, max_num)357 for i in range(num): _rand_erase1(x, pct, xm, xs, mn, mx)358 return x359 360class RandErase(nn.Module):361 def __init__(self, pct=0.2, max_num=4):362 super().__init__()363 self.pct,self.max_num = pct,max_num364 def forward(self, x): return rand_erase(x, self.pct, self.max_num)365class SquareReflectPad:366 def __call__(self, image):367 image = image.squeeze()368 s = image.size()369 max_wh = np.min([s[-1], s[-2] * 2.9])370 hp = np.max(int((max_wh - s[-1]) / 2), 0)371 vp = np.max(int((max_wh - s[-2]) / 2), 0)372 padding = (hp, vp, hp, vp)373 new_img = torchvision.transforms.functional.pad(image, padding, padding_mode='reflect')374 new_img = new_img.unsqueeze(1).expand(new_img.shape[0],3, new_img.shape[1]).permute(1,0,2)375 return new_img376 377 378# %% ../nbs/00_siamese.ipynb 25379def _flops(x, h, w):380 if x.dim()<3: return x.numel()381 if x.dim()==4: return x.numel()*h*w382 383class Learner():384 def __init__(self, model, dlt, dlv, lr=0.1, cbs=None, opt_func=optim.SGD):385 cbs = fc.L(cbs)386 fc.store_attr()387 388 @with_cbs('batch')389 def _one_batch(self):390 self.predict()391 self.callback('after_predict')392 self.get_loss()393 self.callback('after_loss')394 if self.training:395 self.backward()396 self.callback('after_backward')397 self.step()398 self.callback('after_step')399 self.zero_grad()400 401 @with_cbs('epoch')402 def _one_epoch(self):403 for self.iter ,self.batch in enumerate(self.dl): 404 self._one_batch()405 if self.iter > 100:406 break407 408 def one_epoch(self, training):409 self.model.train(training)410 self.dl = self.dlt if training else self.dlv411 self._one_epoch()412 413 @with_cbs('fit')414 def _fit(self, train, valid):415 for self.epoch in self.epochs:416 if train: self.one_epoch(True)417 if valid: torch.no_grad()(self.one_epoch)(False)418 419 def fit(self, n_epochs=1, train=True, valid=True, cbs=None, lr=None):420 cbs = fc.L(cbs)421 # `add_cb` and `rm_cb` were added in lesson 18422 for cb in cbs: self.cbs.append(cb)423 try:424 self.n_epochs = n_epochs425 self.epochs = range(n_epochs)426 if lr is None: lr = self.lr427 if self.opt_func: self.opt = self.opt_func(self.model.parameters(), lr)428 self._fit(train, True)429 finally:430 for cb in cbs: self.cbs.remove(cb)431 432 def __getattr__(self, name):433 if name in ('predict','get_loss','backward','step','zero_grad'): return partial(self.callback, name)434 raise AttributeError(name)435 436 def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)437 438 @property439 def training(self): return self.model.training440 441 def lr_find(self, gamma=1.3, max_mult=3, start_lr=1e-5, max_epochs=10):442 self.fit(max_epochs, lr=start_lr, cbs=LRFinderCB(gamma=gamma, max_mult=max_mult))443 444 def summary(self):445 res = '|Module|Input|Output|Num params|MFLOPS|\n|--|--|--|--|--|\n'446 totp,totf = 0,0447 def _f(hook, mod, inp, outp):448 nonlocal res,totp,totf449 nparms = sum(o.numel() for o in mod.parameters())450 totp += nparms451 *_,h,w = outp.shape452 flops = sum(_flops(o, h, w) for o in mod.parameters())/1e6453 totf += flops454 res += f'|{type(mod).__name__}|{tuple(inp[0].shape)}|{tuple(outp.shape)}|{nparms}|{flops:.1f}|\n'455 with Hooks(self.model, _f) as hooks: self.fit(1, lr=1, cbs=SingleBatchCB())456 print(f"Tot params: {totp}; MFLOPS: {totf:.1f}")457 if fc.IN_NOTEBOOK:458 from IPython.display import Markdown459 return Markdown(res)460 else: print(res)461 462 463 464# %% ../nbs/00_siamese.ipynb 40465@fc.delegates(plt.Axes.imshow)466def show_image(im, ax=None, figsize=None, title=None, noframe=True, **kwargs):467 "Show a PIL or PyTorch image on `ax`."468 if fc.hasattrs(im, ('cpu','permute','detach')):469 im = im.detach().cpu()470 if len(im.shape)==3 and im.shape[0]<5: im=im.permute(1,2,0)471 elif not isinstance(im,np.ndarray): im=np.array(im)472 if im.shape[-1]==1: im=im[...,0]473 if ax is None: _,ax = plt.subplots(figsize=figsize)474 ax.imshow(im, **kwargs, cmap='gray')475 if title is not None: ax.set_title(title, color="red")476 ax.set_xticks([]) 477 ax.set_yticks([]) 478 if noframe: ax.axis('off')479 480 return ax481 482@fc.delegates(plt.subplots, keep=True)483def subplots(484 nrows:int=1, # Number of rows in returned axes grid485 ncols:int=1, # Number of columns in returned axes grid486 figsize:tuple=None, # Width, height in inches of the returned figure487 imsize:int=3, # Size (in inches) of images that will be displayed in the returned figure488 suptitle:str=None, # Title to be set to returned figure489 **kwargs490): # fig and axs491 "A figure and set of subplots to display images of `imsize` inches"492 if figsize is None: figsize=(ncols*imsize, nrows*imsize)493 fig,ax = plt.subplots(nrows, ncols, figsize=figsize, **kwargs)494 if suptitle is not None: fig.suptitle(suptitle)495 if nrows*ncols==1: ax = np.array([ax])496 497 return fig,ax498 499@fc.delegates(subplots)500def get_grid(501 n:int, # Number of axes502 nrows:int=None, # Number of rows, defaulting to `int(math.sqrt(n))`503 ncols:int=None, # Number of columns, defaulting to `ceil(n/rows)`504 title:str=None, # If passed, title set to the figure505 weight:str='bold', # Title font weight506 size:int=14, # Title font size507 **kwargs,508): # fig and axs509 "Return a grid of `n` axes, `rows` by `cols`"510 if nrows: ncols = ncols or int(np.floor(n/nrows))511 elif ncols: nrows = nrows or int(np.ceil(n/ncols))512 else:513 nrows = int(math.sqrt(n))514 ncols = int(np.floor(n/nrows))515 fig,axs = subplots(nrows, ncols, **kwargs)516 for i in range(n, nrows*ncols): axs.flat[i].set_axis_off()517 if title is not None: fig.suptitle(title, weight=weight, size=size)518 return fig,axs519 520@fc.delegates(subplots)521def show_images(ims:list, # Images to show522 nrows:int|None=None, # Number of rows in grid523 ncols:int|None=None, # Number of columns in grid (auto-calculated if None)524 titles:list|None=None, # Optional list of titles for each image525 **kwargs):526 "Show all images `ims` as subplots with `rows` using `titles`"527 axs = get_grid(len(ims), nrows, ncols, **kwargs)[1].flat528 for im,t,ax in zip_longest(ims, titles or [], axs): show_image(im, ax=ax, title=t)529 530def reshape_alternating(tens1, tens2):531 new = torch.stack((tens1, tens2), dim=0)532 return torch.transpose(new,0,1).flatten(start_dim=0, end_dim=1)533