DeOldify/ColorizeTraining.ipynb

In [ ]:

import os
os.environ['CUDA_VISIBLE_DEVICES']='0' 

In [ ]:

import fastai
from fastai import *
from fastai.vision import *
from fastai.callbacks import *
from fastai.vision.gan import *
from fasterai.dataset import *
from fasterai.visualize import *
from fasterai.tensorboard import *
from fasterai.loss import *
from fasterai.critics import *
from fasterai.generators import *
from pathlib import Path
from itertools import repeat
plt.style.use('dark_background')
torch.backends.cudnn.benchmark=True
from PIL import ImageFile

In [ ]:

IMAGENET = Path('data/imagenet/ILSVRC/Data/CLS-LOC')
BWIMAGENET = Path('data/imagenet/ILSVRC/Data/CLS-LOC/bandw')

proj_id = 'colorize1'
TENSORBOARD_PATH = Path('data/tensorboard/' + proj_id)

torch.backends.cudnn.benchmark=True

In [ ]:

def decolorize(fn:str, i:int):
    dest = BWIMAGENET/fn.relative_to(IMAGENET)
    dest.parent.mkdir(parents=True, exist_ok=True)
    img = PIL.Image.open(fn).convert('LA').convert('RGB')
    img.save(dest)  

Uncomment the first time you run this notebook.

In [ ]:

#il = ImageItemList.from_folder(IMAGENET/'val')
#parallel(decolorize, il.items, max_workers=16)

In [ ]:

#il = ImageItemList.from_folder(IMAGENET/'train')
#parallel(decolorize, il.items, max_workers=16)

In [ ]:

def get_data(sz:int, bs:int, keep_pct:float):
    return get_colorize_data(sz=sz, bs=bs, crappy_path=BWIMAGENET, good_path=IMAGENET, 
                             random_seed=None, keep_pct=keep_pct,num_workers=16)

In [ ]:

def save():
    learn_gen.save(proj_id + '_gen_' + str(sz))
    learn_crit.save(proj_id + '_crit_' + str(sz))

In [ ]:

def load():
    learn_gen.load(proj_id + '_gen_' + str(sz))
    learn_crit.load(proj_id + '_crit_' + str(sz))

In [ ]:

def colorize_gen_learner_exp(data:ImageDataBunch, gen_loss=FeatureLoss(), arch=models.resnet34):
    return unet_learner3(data, arch, wd=1e-3, blur=True, norm_type=NormType.Spectral,
                        self_attention=True, y_range=(-3.,3.), loss_func=gen_loss)

Training

In [ ]:

#Needed to instantiate critic but not actually used
sz=64
bs=128

data = get_data(sz=sz, bs=bs, keep_pct=1.0)
learn_crit = colorize_crit_learner(data=data, nf=256)
learn_crit.unfreeze()

gen_loss = FeatureLoss()
learn_gen = colorize_gen_learner_exp(data=data)

switcher = partial(AdaptiveGANSwitcher, critic_thresh=0.65)
learn = GANLearner.from_learners(learn_gen, learn_crit, weights_gen=(1.0,2.0), show_img=False, switcher=switcher,
                                 opt_func=partial(optim.Adam, betas=(0.,0.9)), wd=1e-3)

learn.callback_fns.append(partial(GANDiscriminativeLR, mult_lr=5.))
learn.callback_fns.append(partial(GANTensorboardWriter, base_dir=TENSORBOARD_PATH, name='GanLearner', visual_iters=100))

lr=1e-4
unfreeze_fctr=0.1

64px

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

96px

In [ ]:

#lr=lr/2
sz=96
bs=bs//2

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.1)
learn_gen.freeze_to(-1)
learn.fit(1,lr/10)

In [ ]:

save()

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=1.0)

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

128px

In [ ]:

#lr=lr/2
sz=128
bs=bs//2

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.1)
learn_gen.freeze_to(-1)
learn.fit(1,lr/10)

In [ ]:

save()

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=1.0)

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

In [ ]:

load()

160px

In [ ]:

lr=lr/1.5
sz=160
bs=int(bs//1.5)

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.1)
learn_gen.freeze_to(-1)
learn.fit(1,lr/10)

In [ ]:

save()

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.25)

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

192px

In [ ]:

lr=lr/1.5
sz=192
bs=int(bs//1.5)

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.1)
learn_gen.freeze_to(-1)
learn.fit(1,lr/10)

In [ ]:

save()

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.25)

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

224px

In [ ]:

lr=lr/1.5
sz=224
bs=int(bs//1.5)

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.1)
learn_gen.freeze_to(-1)
learn.fit(1,lr/10)

In [ ]:

save()

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.25)

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

In [ ]:

load()

256px

In [ ]:

lr=lr/1.75
sz=256
bs=int(bs//1.5)

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.1)
learn_gen.freeze_to(-1)
learn.fit(1,lr/10)

In [ ]:

save()

In [ ]:

learn.data = get_data(sz=sz, bs=bs, keep_pct=0.25)

In [ ]:

learn_gen.freeze_to(-1)
learn.fit(1,lr)

In [ ]:

save()

In [ ]:

learn_gen.unfreeze()
learn.fit(1,lr*unfreeze_fctr)

In [ ]:

save()

In [ ]: