nathanhubens / fasterai Goto Github PK

# Launch training
teacher.fine_tune(epochs)

# Distill
m = globals()[config["model"]] # resnet18()
student = Learner(data_loaders, m(num_classes=int(teacher.n_out)),
                  metrics=[accuracy])
kd_cb = KnowledgeDistillationCallback(teacher.model, SoftTarget)
student.fit_one_cycle(epochs, 1e-4, cbs=kd_cb)

# Sparsify
sp_cb = SparsifyCallback(sparsity=50, granularity='filter',
                         context='global', criteria=large_final,
                         schedule=cos)
student.fit(5, 1e-5, cbs=sp_cb)

# Prune
pruner = Pruner()
pruned_model = pruner.prune_model(student.model)

from fastai.vision import all as fst
from fastai.learner import Learner
from fastai.vision.learner import vision_learner
from fastai.tabular.core import df_shrink
# Code to get the datasets...

# Create dataloaders
data_loaders = data_block.dataloaders(tmp_dir, bs=config["batch_size"])  # Config is a dictionay containing useful info
# Create model
base_model = getattr(fst, config["model"])
model = vision_learner(data_loaders, base_model, pretrained=True, lr=0.001,
                       metrics=fst.error_rate)

# Launch training
model.fine_tune(epochs)

from fasterai.distill.all import *
from fastai.vision.all import *
# FasterAI tutorial replica
student = Learner(data_loaders, getattr(fst, config["model"]),
                  metrics=[accuracy])
kd_cb = KnowledgeDistillationCallback(model, SoftTarget)
student.fit_one_cycle(10, 1e-4, cbs=kd_cb)