from pathlib import Path

import json

import argparse

import fastai

from fastai.vision import *

from fastai.vision.gan import *

from fastai.metrics import *

from fastai.callbacks import *

from fastai.callbacks.tensorboard import LearnerTensorboardWriter

from fastai.vision.models import *

from fastai.distributed import *

from pytorchcv.model_provider import get_model as ptcv_get_model

from mmcv import Config as Cfg

from mmcv.runner import load_checkpoint

import data

import loss

from utils import *

def parse_args():

return args

def set_env(nbatch):

return nbatch, device, ncuda

if __name__ == "__main__":

gc.collect()

args = parse_args()

cfg = Cfg.fromfile(args.config)

_, device, _ = set_env(cfg.total_bs)

try:

model = ptcv_get_model(cfg.model.type, pretrained=cfg.model.pretrained)

except:

print("Model Not Implemented. ")

exit()

try:

dataset = getattr(data, cfg.dataset_type)

dataset = dataset(batch_size=cfg.total_bs, imgsize=cfg.img_size)

dataset.show_batch(rows=3, figsize=(10,10))

plt.show()

except:

print("Dataset Type {} are Not Implemented. ".format(cfg.dataset_type))

exit()

metrics = [getattr(fastai.metrics, met) for met in cfg.metric]

try:

loss = getattr(nn, cfg.loss)

except AttributeError as error:

loss = getattr(loss, cfg.loss)

optimizer = get_optimizer(cfg.optimizer)

learner = Learner(dataset, model, opt_func=optimizer, loss_func=loss().to(device),

metrics=metrics, bn_wd=False, true_wd=True,

wd=cfg.optimizer.weight_decay, path=cfg.work_dir) # Custom Learner

# models : Darknet, resnet18,34,50,101,152,xresnet18,34,50,101,152,squeezenet1_0,squeezenet1_1, densenet121

# learner = cnn_learner(dataset, models.resnet18, wd=cfg.optimizer.weight_decay,metrics=metrics)

# learner = unet_learner(dataset, models.resnet34, metrics=partial(foreground_acc,void_code=30), wd=cfg.optimizer.weight_decay)

# GAN

# generator = basic_generator(in_size=64, n_channels=3, n_extra_layers=1)

# critic = basic_critic(in_size=64, n_channels=3, n_extra_layers=1)

# learner = GANLearner.wgan(dataset, generator, critic, switch_eval=False,opt_func=partial(optim.Adam, betas=(0., 0.99)), wd=0.)

# NLP

# languagemodellearner()

# textclassifierlearner()

# tabular_learner()

if args.local_rank > -1:

learner.to_distributed(args.local_rank)

else:

learner.to_parallel()

if args.validate:

learner.model.eval()

print(learner.validate())

exit()

learner.to_fp16()

callbacks = [

SaveModelCallback(learner, name='model_best', monitor='accuracy'),

# LearnerTensorboardWriter(learner,name=Path('runs/'),base_dir=Path(cfg.work_dir))

]

for callback in callbacks:

learner.callbacks.append(callback)

# learner.mixup()

# learner.cutmix()

# learner.ricap()

# learner.blend()

# learner.show_tfms()

print(cfg.text)

# learner.lr_find(num_it=200)

# learner.recorder.plot(suggestion=True,show_moms=False)

# learner.sched.plot_lr(show_moms=False)

# exit()

if cfg.lr_config.policy.lower() == 'fit':

learner.fit(int(cfg.total_epochs),cfg.optimizer.lr)

if cfg.lr_config.policy.lower() == 'cyclic':

learner.fit_one_cycle(int(cfg.total_epochs), cfg.optimizer.lr, tot_epochs=int(cfg.total_epochs),

div_factor=cfg.lr_config.div_factor,pct_start=cfg.lr_config.pct_start)

if cfg.lr_config.policy.lower() == 'cosine':

fit_warmup_cosannealing(learner, cfg.optimizer.lr, cfg.total_epochs, warmup_ratio=cfg.lr_config.warmup_ratio)

if cfg.lr_config.policy.lower() == 'step':

fit_warmup_multistep(learner, cfg.optimizer.lr, cfg.lr_config.gamma, cfg.lr_config.step, cfg.total_epochs,

warmup_ratio=cfg.lr_config.warmup_ratio)

# if cfg.lr_config.policy.lower() == 'warm_restart':

# fit_warmup_restart(learner,n_cycles, cfg.optimizer.lr, cfg.optimizer.mom, cycle_len, cycle_mult)

# show results

learner.gan_trainer.switch(gen_mode=True)

learner.show_results(ds_type=DatasetType.Train, rows=16, figsize=(8, 8))

# Interpretation

# learner.to_fp32()

# learner.model.eval()

# interp = ClassificationInterpretation.from_learner(learner, DatasetType.Valid, tta=True)

# interp.plot_top_losses(16, figsize=(10,10))

# interp.plot_confusion_matrix(slice_size=10)

# interp.most_confused(min_val=2, slice_size=10)

# tta_params = {'beta':0.12, 'scale':1.0}

# preds, y = learner.TTA(ds_type=DatasetType.Valid,**tta_params)

# preds.argmax(1)