dhuynh95

from fastai import *
from fastai.vision import *

from fastai_autoencoder.bottleneck import VAEBottleneck
from fastai_autoencoder.callback import VAEHook
from fastai_autoencoder.util import *
from fastai_autoencoder.vision.learn import VisionAELearner
from mnist_model import create_decoder, create_encoder_denseblock

dhuynh95

train_size = 512
train_size = int(train_size * 1.25)
bs = 128
size = 28

data, valid_data = get_data(train_size,bs=bs,size=size)

# Architectural parameters of our model
conv = nn.Conv2d
act_fn = nn.ReLU

dhuynh95

x,y = data.one_batch()
learn.plot_rec(x)

dhuynh95

n_epoch = 250
lr = 1e-2

learn.fit_one_cycle(n_epoch,lr)

n_epoch = 250
lr = 1e-2 / 2
cbs = [kl_hook]

learn.fit_one_cycle(n_epoch,lr,callbacks=cbs)

dhuynh95

learn.plot_rec_steps(x)

dhuynh95

# Architectural parameters of our model
conv = nn.Conv2d
act_fn = nn.ReLU
bn = nn.BatchNorm2d
rec_loss = "mse"

# Encoder architecture
enc_fn = create_encoder_denseblock
enc_args = {
    "n_dense":3,

dhuynh95

bs_find(learn,lr=lr,bs=bs,num_it=num_it,n_batch=n_batch,beta=beta).plot()

dhuynh95

lr = 1e-3
bs = 64
beta = 0.999
wd = 0.2
num_it = 5000
n_batch = 20

bs_find(learn,lr=lr,num_it=num_it,n_batch=n_batch,bs=bs,beta=beta,wd=wd).plot()

dhuynh95

class CustomDropout(nn.Module):
    """Custom Dropout module to be used as a baseline for MC Dropout"""

    def __init__(self, p:float, activate=True):
        super().__init__()
        self.activate = activate
        self.p = p

    def forward(self, x):
        return nn.functional.dropout(x, self.p, training=self.training or self.activate)

dhuynh95

# Convert nn.Dropout to CustomDropout module
get_args = lambda dp : {"p" : dp.p}
convert_layers(learn.model,nn.Dropout,CustomDropout,get_args)

# Turn on the stochasticity, I use verbose just to make sure it's working fine
switch_custom_dropout(learn.model,True,verbose=True)