# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/08_hooks.ipynb.

# %% auto 0
__all__ = ['ToyCNN', 'Hook']

# %% ../nbs/08_hooks.ipynb 3
# from copy import copy
# import math
# import fastcore.all as fc
# from fastprogress import progress_bar, master_bar
# from operator import attrgetter
# from collections.abc import Mapping

# import matplotlib.pyplot as plt

# import torch
# from torchvision import transforms as T
# from torcheval.metrics import Mean

# %% ../nbs/08_hooks.ipynb 4
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split

from .loading import load_ossl
from .models import conv
from .callbacks import to_cpu
from .preprocessing import Log1p, SNV
from .dataloaders import SpectralDataset, get_dls

import torch
from torch import nn, tensor

# %% ../nbs/08_hooks.ipynb 7
from functools import partial
from torch.optim import lr_scheduler

from torcheval.metrics import R2Score
from .callbacks import (MetricsCB, BatchSchedCB, BatchTransformCB,
                            DeviceCB, TrainCB, ProgressCB)
from .transforms import GADFTfm, _resizeTfm, StatsTfm
from .learner import Learner

# %% ../nbs/08_hooks.ipynb 8
class ToyCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            conv(1, 8, ks=5),        # 8, 851
            conv(8, 16),             # 16, 426
            conv(16, 32),            # 32, 213
            conv(32, 64),            # 64, 107
            nn.Flatten(),
            nn.Linear(64*107, 1))
        
    def forward(self, x):
        return self.layers(x)

# %% ../nbs/08_hooks.ipynb 11
class Hook():
    def __init__(self, m, f): self.hook = m.register_forward_hook(partial(f, self))
    def remove(self): self.hook.remove()
    def __del__(self): self.remove()
