# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_learner.ipynb.

# %% auto 0
__all__ = ['run_cbs', 'CancelFitException', 'CancelBatchException', 'CancelEpochException', 'Learner']

# %% ../nbs/05_learner.ipynb 3
import torch
from contextlib import contextmanager
from operator import attrgetter
from torch import optim
import torch.nn.functional as F
from functools import partial

from .callbacks import Callback
from .model import get_seg_model, get_model_name

# %% ../nbs/05_learner.ipynb 4
def run_cbs(cbs, method_nm, learn=None):
    for cb in sorted(cbs, key=attrgetter('order')):
        method = getattr(cb, method_nm, None)
        
        if method is not None: 
            method(learn)

# %% ../nbs/05_learner.ipynb 7
class CancelFitException(Exception): pass
class CancelBatchException(Exception): pass
class CancelEpochException(Exception): pass

# %% ../nbs/05_learner.ipynb 8
class Learner():
    def __init__(self, model, dls, cbs, loss_func=F.cross_entropy, lr=0.1, opt_func=optim.SGD):
        self.model = model
        self.dls = dls
        self.loss_func = loss_func
        self.opt_func = opt_func
        self.cbs = cbs
        self.n_epochs = 1
        self.epochs = None
        self.opt = None
        self.dl = None
        self.lr = lr
       
    def train_epoch(self):
        self.model.train(True)
        self.dl = self.dls.train
        self.callback('before_epoch')
        for self.iter,self.batch in enumerate(self.dl):
            self.callback('before_batch')
            self.preds = self.model(self.batch[0])
            self.loss = self.loss_func(self.preds, self.batch[1])
            self.loss.backward()
            self.opt.step()
            self.opt.zero_grad()       
            self.callback('after_batch')
        self.callback('after_epoch')       
        
    def eval_epoch(self):
        self.model.train(False)
        self.dl = self.dls.valid
        with torch.no_grad():
            self.callback('before_epoch')
            for self.iter,self.batch in enumerate(self.dl):
                self.callback('before_batch')
                self.preds = self.model(self.batch[0])
                self.loss = self.loss_func(self.preds, self.batch[1])
                self.callback('after_batch')
            self.callback('after_epoch')   
        
    def fit(self, n_epochs):
        self.n_epochs = n_epochs
        self.epochs = range(n_epochs)
        self.opt = self.opt_func(self.model.parameters(), self.lr)
        self.callback('before_fit')
        for self.epoch in self.epochs:
            self.callback('before_full_epoch')
            self.train_epoch()
            torch.cuda.empty_cache()
            self.eval_epoch()
            torch.cuda.empty_cache()
            self.callback('after_full_epoch')
        self.callback('after_fit')

    def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)
