<jupyter_start><jupyter_text>Autoencoder using fastai V2  setup<jupyter_code>#!pip install -Uqq fastai fastcore
import fastai
fastai.__version__
from fastai.test_utils import show_install
show_install()
from fastai.vision.all import *<jupyter_output><empty_output><jupyter_text>Using V2- Datablock API- check that tfms are applied to input/target- labeling func: `lambda o: o`<jupyter_code>#path = untar_data(URLs.MNIST_SAMPLE)

# load our core data
filename = "zappos-50k-simplified_sort" # no features
filename = "full_db" # no features


df = pd.read_pickle(f"data/{filename}.pkl")

#DATA_path = Path.home()/'Projects/DATABASE'
#path_images = Path(config.d['data_path']+'/ut-zap50k-images')


# Lets test a few different sizes

IMG_SIZE_LG = 224
IMG_SIZE_MD = 160
IMG_SIZE_SM = 128

IMG_SIZE = IMG_SIZE_MD

IMG_SIZES = {"small":  IMG_SIZE_SM,
             "medium": IMG_SIZE_MD,
             "large":  IMG_SIZE_LG}
np.random.seed(3333)
torch.manual_seed(3333)

image_path = Path("/home/ergonyc/Projects/Project2.0/snkr-finder/data/")
#image_path = Path("data/")
batchsize = 32


df = df[["path","train","test","validate","t_t_v"]]
#DATA_path = Path.home()/'Projects/DATABASE'
#path_images = Path(config.d['data_path']+'/ut-zap50k-images')

# alternate strategy:
# Drop all the "test" rows for now, and create an "is_valid" column...
# should probably drop a ton of columns to jus tkeep the file paths...

df.loc[:,'is_valid']= df.test | df.validate

L([image_path/d for d in df.path])
# some helper functions borrowed from validating the feature embedding
def get_x(r): return image_path/r['path']
#def get_y(r): return r['Category']
get_y = get_x  

def splitter(df):
    train = df.index[df['train']]
    valid = df.index[df['validate']]
    return train,valid


batch_tfms = Normalize.from_stats(*imagenet_stats)
rand_tfms = aug_transforms(mult=1.0, 
               do_flip=True, 
               flip_vert=False, 
               max_rotate=3.0, 
               min_zoom=.95, 
               max_zoom=1.01, 
               max_lighting=0.1, 
               max_warp=0.023, 
               p_affine=0.66, 
               p_lighting=0.2, 
               xtra_tfms=None, 
               size=None, 
               mode='bilinear', 
               pad_mode='border', 
               align_corners=True, 
               batch=False, 
               min_scale=1.0)

btfms = rand_tfms+[batch_tfms]
tr,vl = splitter(df)
len(vl),len(tr),df.shape
doc(ColReader)
ColSplitter('is_valid')(df),splitter(df)
get_x(df)[-10:],tr[-10:]
block = DataBlock(blocks=(ImageBlock(cls=PILImage), ImageBlock(cls=PILImage)),
                  get_x=get_x, 
                  get_y=get_y,
                  splitter=ColSplitter('is_valid'), 
                  item_tfms= Resize(IMG_SIZE,method='pad', pad_mode='border'),
                  batch_tfms = btfms,
)
# # Alternate way to call... just gobble all the images...
# block1 = DataBlock(blocks=(ImageBlock(cls=PILImage), ImageBlock(cls=PILImage)),
#                   get_items = lambda o: get_image_files(o,folders=['scraped','zappos']),
#                   get_y = lambda o: o,
#                   splitter=ColSplitter('is_valid'),
#                   item_tfms=Resize(IMG_SIZE,method='pad', pad_mode='border'),
#                   batch_tfms = btfms,
# )

# im_path = Path("/home/ergonyc/Projects/Project2.0/snkr-finder/imgdb/")

# fls = get_image_files(im_path,folders=['scraped','zappos']),

# len(fls[0])

# type(get_image_files(im_path,folders=['scraped','zappos'])[0])

doc(get_image_files)
# block.summary(im_path)
# dls = block.dataloaders(image_path, batch_size=32)

#block.summary(df)
dls = block.dataloaders(df, batch_size=32)

#block1.summary(im_path)
#dls = block1.dataloaders(image_path, batch_size=32)

dls.show_batch()
rnet_body = create_body(xresnet18, n_in=3)
# need to pool 512*5*5 to latent_size+latent_sizek
x,y = dls.one_batch()
x.shape, y.shape
#rnet_body.cuda()(x).shape
doc(xresnet18)<jupyter_output><empty_output><jupyter_text>We will define a basic architecture to decode the image.<jupyter_code>class UpsampleBlock(Module):
    def __init__(self, up_in_c:int, final_div:bool=True, blur:bool=False, leaky:float=None, **kwargs):
        """
        up_in_c: upsample input channels
            -> maps to ni, nf=ni//2 
        """
        self.shuf = PixelShuffle_ICNR(up_in_c, up_in_c//2, blur=blur, **kwargs)
        ni = up_in_c//2
        nf = ni if final_div else ni//2
        self.conv1 = ConvLayer(ni, nf, **kwargs)
        self.conv2 = ConvLayer(nf, nf, **kwargs)
        self.relu = nn.ReLU()

    def forward(self, up_in:Tensor) -> Tensor:
        up_out = self.shuf(up_in)
        cat_x = self.relu(up_out)
        
        return self.conv2(self.conv1(cat_x))
    
    def shuff(self,up_in:Tensor) -> Tensor:
        up_out = self.shuf(up_in)
        cat_x = self.relu(up_out)
        return cat_x
    
doc(PixelShuffle_ICNR)
def decoder_resnet(y_range, n_out=3):
    return nn.Sequential(UpsampleBlock(512), 
                         UpsampleBlock(256),
                         UpsampleBlock(128),
                         UpsampleBlock(64),
                         UpsampleBlock(32),
                         nn.Conv2d(16, n_out, 1),
                         SigmoidRange(*y_range)
                        )
                         
y_range = (-3.,3.)
rnet_dec = decoder_resnet(y_range).cuda()

# deterministic autoencoder...
rnet_ae = nn.Sequential(rnet_body, rnet_dec)
rnet_dec[:6].cuda()(rnet_body.cuda()(x)).shape
rnet_dec[0].cuda().shuff(rnet_body.cuda()(x)).shape
def my_xresnet(pretrained, expansion, layers, **kwargs):
    # TODO pretrain all sizes. Currently will fail with non-xrn50
    url = 'https://s3.amazonaws.com/fast-ai-modelzoo/xrn50_940.pth'
    res = XResNet(ResBlock, expansion, layers, **kwargs)
    if pretrained: res.load_state_dict(load_state_dict_from_url(url, map_location='cpu')['model'], strict=False)
    return res
doc(XResNet)
# def _update_first_layer(model, n_in, pretrained):
#     "Change first layer based on number of input channels"
#     if n_in == 3: return 
    
#     first_layer, parent, name = _get_first_layer(model)
#     assert isinstance(first_layer, nn.Conv2d), f'Change of input channels only supported with Conv2d, found {first_layer.__class__.__name__}'
#     assert getattr(first_layer, 'in_channels') == 3, f'Unexpected number of input channels, found {getattr(first_layer, "in_channels")} while expecting 3'
#     params = {attr:getattr(first_layer, attr) for attr in 'out_channels kernel_size stride padding dilation groups padding_mode'.split()}
#     params['bias'] = getattr(first_layer, 'bias') is not None
#     params['in_channels'] = n_in
#     new_layer = nn.Conv2d(**params)
#     if pretrained:
#         _load_pretrained_weights(new_layer, first_layer)
#     setattr(parent, name, new_layer)

# def my_create_body(arch, n_in=3, pretrained=True, cut=None):
#     "Cut off the body of a typically pretrained `arch` as determined by `cut`"
#     model = arch(pretrained=pretrained)
#     _update_first_layer(model, n_in, pretrained)
#     #cut = ifnone(cut, cnn_config(arch)['cut'])
#     if cut is None:
#         ll = list(enumerate(model.children()))
#         cut = next(i for i,o in reversed(ll) if has_pool_type(o))
#     if   isinstance(cut, int):      return nn.Sequential(*list(model.children())[:cut])
#     elif callable(cut): return cut(model)
#     else:                           raise NamedError("cut must be either integer or a function")


def simple_cut_arch(arch):
    model = arch(pretrained=True)
    ll = list(enumerate(model.children()))
    cut = next(i for i,o in reversed(ll) if has_pool_type(o))
    return nn.Sequential(*list(model.children())[:cut])



def simple_cut_model(model):
    ll = list(enumerate(model.children()))
    cut = next(i for i,o in reversed(ll) if has_pool_type(o))
    return nn.Sequential(*list(model.children())[:cut])

#rnet_body = create_body(xresnet18, n_in=3)

net1 = my_xresnet(False, 1, [2, 2,  2]) #torch.Size([32, 1000])
rnet_body = simple_cut_model(net1)
#net1 = my_xresnet(False, 1, [3, 4,  6, 3]) #torch.Size([32, 1000])
#net1 = my_xresnet(False, 4, [3, 4,  6, 3]) #torch.Size([32, 1000])
arch = simple_cut(xresnet18)

    ll = list(enumerate(model.children()))
    cut = next(i for i,o in reversed(ll) if has_pool_type(o))
    return nn.Sequential(*list(model.children())[:cut])


[x for x in range(10)][:9]

#rnet_body.cuda()(x).shape
for i in range(11):
    print(rnet_body[:i].cuda()(x).shape)
    #print(arch[:i])

    #arch

    
    
    
#net1[:5].cuda()(x).shape
# need to pool 512*5*5 to latent_size+latent_sizek
#
# rnet_dec(arch_rnet(x)).shape,arch_rnet(x).shape,rnet_dec[0](arch_rnet(x)).shape
# doc(MSELossFlat)

def xresnet18 (pretrained=False, **kwargs): return _xresnet(pretrained, 1, [2, 2,  2, 2], **kwargs)
def xresnet34 (pretrained=False, **kwargs): return _xresnet(pretrained, 1, [3, 4,  6, 3], **kwargs)
def xresnet50 (pretrained=False, **kwargs): return _xresnet(pretrained, 4, [3, 4,  6, 3], **kwargs)
def xresnet101(pretrained=False, **kwargs): return _xresnet(pretrained, 4, [3, 4, 23, 3], **kwargs)
def xresnet152(pretrained=False, **kwargs): return _xresnet(pretrained, 4, [3, 8, 36, 3], **kwargs)
def xresnet18_deep  (pretrained=False, **kwargs): return _xresnet(pretrained, 1, [2,2,2,2,1,1], **kwargs)
def xresnet34_deep  (pretrained=False, **kwargs): return _xresnet(pretrained, 1, [3,4,6,3,1,1], **kwargs)
def xresnet50_deep  (pretrained=False, **kwargs): return _xresnet(pretrained, 4, [3,4,6,3,1,1], **kwargs)
def xresnet18_deeper(pretrained=False, **kwargs): return _xresnet(pretrained, 1, [2,2,1,1,1,1,1,1], **kwargs)
def xresnet34_deeper(pretrained=False, **kwargs): return _xresnet(pretrained, 1, [3,4,6,3,1,1,1,1], **kwargs)
def xresnet50_deeper(pretrained=False, **kwargs): return _xresnet(pretrained, 4, [3,4,6,3,1,1,1,1], **kwargs)


learn = Learner(dls, rnet_ae, loss_func=MSELossFlat()).to_detach
learn.fit_one_cycle(3)

learn.show_results()
for i in range(0,10):
    b = learn.dls.one_batch()
    pred = learn.model(b[0])
    e = learn.loss_func(pred,b[1])
    print(f"mse: {e.mean()}" )
learn.model(b[0])
learn.recorder.plot_losses()

learn.show_results()
# learn.recorder.plot_sched()<jupyter_output><empty_output><jupyter_text>VAEVariational Auto-Encoder for FastAI<jupyter_code># CLASSES
    
class ConvBnRelu(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size = 3, stride = 1, bias = False, conv = nn.Conv2d,
                 bn = nn.BatchNorm2d, act_fn = nn.ReLU,**kwargs):
        super(ConvBnRelu,self).__init__()
        self.conv = conv(in_channels=in_channels,out_channels=out_channels,
                         kernel_size=kernel_size,stride = stride, bias=bias,**kwargs)
        if bn:
            self.bn = bn(out_channels)
        else:
            self.bn = None
        if act_fn:
            self.act_fn = act_fn(inplace = True)
        else:
            self.act_fn = None
            
    def forward(self,x):
        output = self.conv(x)
        if self.bn:
            output = self.bn(output)
        if self.act_fn:
            output = self.act_fn(output)
        return output
                          
class Downsample(nn.Module):
    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
        super(Downsample, self).__init__()
        self.filt_size = filt_size
        self.pad_off = pad_off
        self.pad_sizes = [int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)), int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2))]
        self.pad_sizes = [pad_size+pad_off for pad_size in self.pad_sizes]
        self.stride = stride
        self.off = int((self.stride-1)/2.)
        self.channels = channels

        #print('Filter size [%i]'%filt_size)
        if(self.filt_size==1):
            a = np.array([1.,])
        elif(self.filt_size==2):
            a = np.array([1., 1.])
        elif(self.filt_size==3):
            a = np.array([1., 2., 1.])
        elif(self.filt_size==4):    
            a = np.array([1., 3., 3., 1.])
        elif(self.filt_size==5):    
            a = np.array([1., 4., 6., 4., 1.])
        elif(self.filt_size==6):    
            a = np.array([1., 5., 10., 10., 5., 1.])
        elif(self.filt_size==7):    
            a = np.array([1., 6., 15., 20., 15., 6., 1.])

        filt = torch.Tensor(a[:,None]*a[None,:])
        filt = filt/torch.sum(filt)
        self.register_buffer('filt', filt[None,None,:,:].repeat((self.channels,1,1,1)))

        self.pad = nn.ReflectionPad2d(self.pad_sizes)

    def forward(self, inp):
        if(self.filt_size==1):
            if(self.pad_off==0):
                return inp[:,:,::self.stride,::self.stride]    
            else:
                return self.pad(inp)[:,:,::self.stride,::self.stride]
        else:
            return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])

# resnet building blocks...
class ResBlock(nn.Module):
    def __init__(self,in_channels, kernel_size=3, stride=1, padding=1, conv=ConvBnRelu,**kwargs):
        super(ResBlock,self).__init__()
        self.conv1 = conv(in_channels=in_channels,out_channels=in_channels//2,kernel_size=kernel_size,
                        stride=stride,padding=padding,**kwargs)
        self.conv2 = conv(in_channels=in_channels//2,out_channels=in_channels,kernel_size=kernel_size,
                        stride=stride,padding=padding,**kwargs)
        
    def forward(self,x):
        h = self.conv2(self.conv1(x))
        return x + h
    
class DenseBlock(nn.Module):
    def __init__(self,in_channels, kernel_size=3, stride=1, padding=1, conv=ConvBnRelu,**kwargs):
        super(DenseBlock,self).__init__()
        self.conv = ResBlock(in_channels,kernel_size=kernel_size,stride=stride,padding=padding,conv=conv,**kwargs)
        
    def forward(self,x):
        h = self.conv(x)
        output = torch.cat([x,h],dim=1)
        return output

# this is the "Variational" magic.  aka the "reparamaterization trick"
class NormalSampler(nn.Module):
    def __init__(self):
        super(NormalSampler,self).__init__()
    
    def forward(self,mu,logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mu + eps * std
        return z

class VAELinear(nn.Module):
    def __init__(self,in_features,out_features):
        super(VAELinear,self).__init__()
        self.linear = nn.Linear(in_features,out_features)
        self.mu_linear = nn.Linear(out_features,out_features)
        self.logvar_linear = nn.Linear(out_features,out_features)

    def forward(self,x):
        x = self.linear(x)
        mu = self.mu_linear(x)
        logvar = self.logvar_linear(x)

        return mu,logvar

class VAELayer(nn.Module):
    def __init__(self,in_features,out_features):
        super(VAELayer,self).__init__()
        self.blinear = VAELinear(in_features,out_features)
        self.sampler = NormalSampler()

    def forward(self,x):
        mu,logvar = self.blinear(x)
        z = self.sampler(mu,logvar)
        return z
    
class VAEBottleneck(nn.Module):
    """    
    Wrapper to make the VAELayer play nice with the input (from encoder)
    and output (to decoder/generator)

    __init__ takes a list of "features" nominally an inputshape and output shape.  
    nfs : [128*3,latent_dimension]
    
    """
    def __init__(self,nfs:list):
        super(VAEBottleneck,self).__init__()

        n = len(nfs)
        layers = [nn.Linear(in_features=nfs[i], out_features=nfs[i+1]) if i < n -2 
             else VAELayer(in_features=nfs[i], out_features=nfs[i+1]) for i in range(n-1)]
        # just makes a single VAELayer for nfs=[128,32]
        self.fc = nn.Sequential(*layers)
    
    def forward(self,x):
        bs = x.size(0)
        x = x.view(bs, -1)
        z = self.fc(x)
        
        return z
    
    
class PointwiseConv(nn.Module):
    def __init__(self,in_channels,out_channels,bias = True,conv = nn.Conv2d,**kwargs):
        super(PointwiseConv,self).__init__()
        self.conv = conv(in_channels=in_channels,out_channels=out_channels,kernel_size=1,
                         stride=1,padding=0,bias = bias)

    def forward(self,x):
        output = self.conv(x)
        return output

class SpatialDecoder2D(nn.Module):
    def __init__(self,size):
        super(SpatialDecoder2D,self).__init__()
        
        # Coordinates for the broadcast decoder
        # jah changed from 1 to 3 below
        self.size = size
        x = torch.linspace(-1, 3, size)
        y = torch.linspace(-1, 3, size)
        
        x_grid, y_grid = torch.meshgrid(x, y)
        # Add as constant, with extra dims for N and C
        self.register_buffer('x_grid', x_grid.view((1, 1) + x_grid.shape))
        self.register_buffer('y_grid', y_grid.view((1, 1) + y_grid.shape))
        
    def forward(self,z):
        batch_size = z.size(0)
        # View z as 4D tensor to be tiled across new H and W dimensions
        # Shape: NxDx1x1
        z = z.view(z.shape + (1, 1))

        # Tile across to match image size
        # Shape: NxDx64x64
        z = z.expand(-1, -1, self.size, self.size)

        # Expand grids to batches and concatenate on the channel dimension
        # Shape: Nx(D+2)x64x64
        x = torch.cat((self.x_grid.expand(batch_size, -1, -1, -1),
                       self.y_grid.expand(batch_size, -1, -1, -1), z), dim=1)

        return x
    

## ARCHITECUTE
def create_encoder(nfs,ks,conv=nn.Conv2d,bn=nn.BatchNorm2d,act_fn = nn.ReLU):
    """
    constructor for generic convolutional encoder 
    """
    n = len(nfs)
    conv_layers = [nn.Sequential(ConvBnRelu(nfs[i],nfs[i+1],kernel_size=ks[i],
                                            conv = conv,bn=bn,act_fn=act_fn, padding = ks[i] //2 ),
                                 Downsample(channels=nfs[i+1],filt_size=3,stride=2))
                                   for i in range(n-1)]        
    convs = nn.Sequential(*conv_layers)
    return convs

def create_encoder_denseblock(n_dense,c_start):
    """
    constructor for resnet with dense blocks  (?) 

    n_dense": 3,
    "c_start": 4
    """
    first_layer = nn.Sequential(ConvBnRelu(3,c_start,kernel_size=3,padding = 1),
                                ResBlock(c_start),
                                Downsample(channels=4,filt_size=3,stride=2))
    
    layers = [first_layer] + [
        nn.Sequential(
            DenseBlock(c_start * (2**c)),
            Downsample(channels=c_start * (2**(c+1)),filt_size=3,stride=2)) for c in range(n_dense)
    ]
    
    model = nn.Sequential(*layers)
    
    return model

def create_decoder(nfs, ks, size, conv=nn.Conv2d, bn=nn.BatchNorm2d, act_fn=nn.ReLU):
    """
    CURR VALUES:
    "nfs":[66,3*32,3*16,3*8,3*4,3*2,3,1,3],
    "ks": [ 3, 1, 3,1,3,1,3,1],   
    "size": IMG_SIZE 
    """
    n = len(nfs)
    
    # We add two channels to the first layer to include x and y channels
    first_layer = ConvBnRelu(nfs[0], #input size 
                             nfs[1], # output size
                             conv = PointwiseConv,
                             bn=bn,
                             act_fn=act_fn)

    conv_layers = [first_layer] + [ConvBnRelu(nfs[i],nfs[i+1],kernel_size=ks[i-1],
                                              padding = ks[i-1] // 2,conv = conv,bn=bn,act_fn=act_fn)
                                   for i in range(1,n - 1)]        
    dec_convs = nn.Sequential(*conv_layers)
    
    dec = nn.Sequential(SpatialDecoder2D(size),dec_convs)
    #SigmoidRange(*y_range)
    return dec

def decoder_simple(y_range=[-3,3], n_out=3):
    return nn.Sequential(#UpsampleBlock(64),
                         UpsampleBlock(32),
                         nn.Conv2d(16, n_out, 1),
                         SigmoidRange(*y_range)
                        )
# Encoder architecture
# We create each part of the autoencoder
enc_args = {
    "n_dense": 3,
    "c_start": 4
}
enc = create_encoder_denseblock(**enc_args)
# enc_args_simple = {"nfs":[66,3*32,3*16,3*8,3*4,3*2,3],
#     "ks": [ 3, 1, 3,1,3,3]
#     }
# enc = create_encoder(**enc_args_simple)


# VAE "Bottleneck" architecture
latent_dimension = 64 # actually 128=64mu+64logvar
bn_args = {
    "nfs":[32*10*10,latent_dimension]
}
# 2**(c_start+1)*(IMG_SIZE*2**-c_start)**2
bn = VAEBottleneck(**bn_args)


# Decoder (generator) architecture
# standard 3,1 pair of kernels
dec_args = {
#     "nfs":[66,3*32,3*16,3*8,3*4,3*2,3,1,3],
#     "ks": [ 3, 1, 3,1,3,1,3,1], 
    "nfs":[66,3*32,3*16,3*8,3*4,3*2,3],
    "ks": [ 3, 1, 3,1,3,3], 
    "size": IMG_SIZE #10 #IMG_SIZE
}
dec = create_decoder(**dec_args)

# # simple "UpsampleBlock" decoder 
# dec_simple = decoder_simple()
arch = nn.Sequential(enc,bn,dec)
#arch2 = nn.Sequential(enc,bn,dec_simple)
v_encoder = nn.Sequential(enc,bn) #variational 

z = v_encoder.cuda()(x)
z.shape,dec_simple[0]

#dec[0](y).shape,y.shape
#dec_simple(z).shape
rnet_bn = VAEBottleneck(nfs=[512*5*5,latent_dimension])
rnet_enc = nn.Sequential(rnet_body,rnet_bn)

rnet_vae = nn.Sequential(rnet_body,rnet_bn,dec)
#rnet_vae2 = nn.Sequential(rnet_body,rnet_bn,dec_simple)  # need to fix the dec_simple...
#y = rnet_body.cuda()(x)

#rnet_body.cuda()(x).shape

#y.shape,rnet_enc(x).shape

arch = arch.cuda()
x_hat1 = arch(x)

x_hat1.max(),x_hat1.min(),x_hat1.mean(),x_hat1.std()
x.max(),x.min(),x.mean(),x.std()
rnet_vae = rnet_vae.cuda()
x_hat2 = rnet_vae(x)
x_hat2.max(),x_hat2.min(),x_hat2.mean(),x_hat2.std()

(IMG_SIZE*2**-4)**2
learn = Learner(dls, rnet_vae, loss_func=MSELossFlat())
lr_min, lr_steep = learn.lr_find()

lr_min,lr_steep
learn.fit_one_cycle(25,lr_max =.0007)

learn.show_results()
#input (32, 3, 160, 160)

#enc[0] (32, 4, 80, 80)
#enc[1] (32, 4, 80, 80
#enc[2] (32, 8, 40, 40
#enc[3] (32, 16, 20, 20
#enc[4] (32, 32, 10, 10


#bn input dim: 32*10*10
# output: (32, 64)

# input     (32,64)
# dec[0]    (32, 66, 160, 160)
# dec[1][0] (32, 96, 160, 160)
# dec[1][1] (32, 48, 160, 160)
# dec[1][2] (32, 24, 160, 160)
# dec[1][3] (32, 12, 160, 160)
# dec[1][4] (32, 6, 160, 160)
# dec[1][5] (32, 3, 160, 160)


#y_range = (-3.,3.)
my_ac = encoder.cuda()

my_dec = dec.cuda()


y = encoder(x)
y = my_ac(x)

dec[0](y).shape,y.shape
#my_deencoder(x)).shape
#my_ac(x).shape
class AutoEncoderLearner(Learner):
    def __init__(self,data:DataBlock,rec_loss:str,enc:nn.Module,bn:nn.Module,dec:nn.Module,**kwargs):
        self.enc = enc
        self.bn = bn
        self.dec = dec
        
        assert rec_loss in ["mse","ce"],"Loss function must be mse or ce"
        if rec_loss == "mse":
            self.rec_loss = nn.MSELoss(reduction="none")
        else:
            self.rec_loss = nn.CrossEntropyLoss(reduction="none")
        self.encode = nn.Sequential(enc,bn)
        
        self.inferring = False
        
        ae = nn.Sequential(enc,bn,dec)
        
        super().__init__(data, ae, loss_func=self.loss_func, **kwargs)
        
        # Callback to replace y with x during the training loop
        replace_cb = ReplaceTargetCallback(self)
        self.callbacks.append(replace_cb)
        
    def loss_func(self,x_rec,x,**kwargs):
        bs = x.shape[0]
        if isinstance(self.rec_loss,nn.MSELoss):
            l = self.rec_loss(x, x_rec).view(bs, -1).sum(dim=-1).mean()
        else:
            # First we discretize x and turn it from (B,1,H,W) to (B,H,W)
            x = (x * 256).long().squeeze(1)
            l = self.rec_loss(x_rec,x).view(bs,-1).sum(dim=-1).mean()
        return l
    
    def decode(self,z):
        x_rec = self.dec(z)
        return x_rec
    
#     def plot_2d_latents(self,ds_type:DatasetType=DatasetType.Valid, n_batch = 10,return_fig=False):
#         """Plot a 2D map of latents colored by class"""
#         z,y = self.get_latents(ds_type = ds_type, n_batch = n_batch)
#         z,y = z.numpy(), y.numpy()
#         print("Computing the TSNE projection")
#         zs = TSNE(n_components=2).fit_transform(z)
#         fig,ax = plt.subplots(1,figsize = (16,12))
#         ax.scatter(zs[:,0],zs[:,1],c = y)
#         ax.set_title("TSNE projection of latents on two dimensions")
#         if return_fig:
#             return fig
        
    def plot_rec(self,x=None,i=0,sz = 64,gray = True):
        """Plot a """
        self.model.cpu()
        if not isinstance(x,torch.Tensor):
            x,y = self.data.one_batch()
            
        x_rec = self.model(x)
        
        if isinstance(self.rec_loss,nn.CrossEntropyLoss):
            x_rec = x_rec.argmax(dim = 1,keepdim = True)
            
        x_rec = x_rec[i].squeeze(1)
        x = x[i].squeeze(1)

        img = x.permute(1,2,0).numpy()
        img_r = x_rec.permute(1,2,0).detach().numpy()
        
        if gray:
            img = np.concatenate((img,)*3,axis = -1)
            img_r = np.concatenate((img_r,)*3,axis = -1)

        fig,ax = plt.subplots(2,figsize = (16,16))
        ax[0].imshow(img,cmap = "gray")
        ax[0].set_title("Original")

        ax[1].imshow(img_r,cmap = "gray")
        ax[1].set_title("Reconstruction")

        self.model.cuda()

    # We perturb the latent variables on each variable with different magnitudes 
    def plot_shades(self,x=None,i=0,n_perturb = 13, mag_perturb = 3, ax = None):
        """Plot the reconstruction of the """
        self.model.cpu()
        if not isinstance(x,torch.Tensor):
            x,y = self.data.one_batch()

        x = x[i].unsqueeze(0)

        # We get the latent code
        z = self.encode(x)
        n_z = z.shape[1]

        # We create a scale of perturbations
        mag_perturb = np.abs(mag_perturb)
        scale_perturb = np.linspace(-mag_perturb,mag_perturb,n_perturb)
        
        if not ax:
            fig, ax = plt.subplots(n_z,n_perturb, figsize=(16,12))
            fig.tight_layout()

        for i in range(n_z):
            for (j,perturb) in enumerate(scale_perturb):
                # For each z_i we add a minor perturbation
                z_perturb = z.clone()
                z_perturb[0,i] += perturb

                # We reconstruct our image
                x_rec = self.decode(z_perturb)

                # We plot it in the grid
                img = x_rec.squeeze(0).permute(1,2,0).detach().numpy()
                ax[i][j].set_axis_off()
                ax[i][j].imshow(img)
                ax[i][j].set_title(f"z_{i} with {round(perturb * 1e2) / 1e2}")
                
        self.model.cuda()
        
    def get_latents(self, ds_type:DatasetType=DatasetType.Valid, activ:nn.Module=None,
                  with_loss:bool=False, n_batch=None, pbar=None):
        "Return predictions and targets on `ds_type` dataset."
        lf = self.loss_func if with_loss else None
        self.inferring = True
        output = get_preds(self.encode, self.dl(ds_type), cb_handler=CallbackHandler(self.callbacks),
                         activ=activ, loss_func=lf, n_batch=n_batch, pbar=pbar)
        self.inferring = False
        return output
    
    def get_error(self, ds_type:DatasetType=DatasetType.Valid,activ:nn.Module=None, n_batch=None, pbar=None):
        "Return predictions and targets on `ds_type` dataset."
        
        x_rec,x = get_preds(self.model, self.dl(ds_type), cb_handler=CallbackHandler(self.callbacks),
                         activ=activ, loss_func=None, n_batch=n_batch, pbar=pbar)
        loss_func = lambda x_rec,x : self.rec_loss(x, x_rec, reduction='none').view(x.shape[0], -1).sum(dim=-1)
        l = loss_func(x_rec,x)
        return l
    
    
def gray_to_rgb(img):
    return np.concatenate((img,)*3,axis=-1)

def register_output(m,i,o):
    return o

class VisionAELearner(AutoEncoderLearner):
    def plot_rec(self,x=None,random=False,n=5,gray = True,return_fig=False):
        """Plot a """
        self.model.cpu()
        if not isinstance(x,torch.Tensor):
            x,y = self.data.one_batch()
        
        bs = self.data.batch_size
        
        assert n < bs, f"Number of images to plot larger than batch size {n}>{bs}"
        
        if random:
            idx = np.random.choice(bs,n,replace = False)
        else:
            idx = np.arange(n)
        x = x[idx]
        
        x_rec = self.model(x)
    
        if isinstance(self.rec_loss,nn.CrossEntropyLoss):
            x_rec = x_rec.argmax(dim = 1,keepdim = True)
        
        fig,ax = plt.subplots(n,2,figsize = (20,20))
        
        for i in range(n):
            img = x[i].permute(1,2,0).numpy()
            img_r = x_rec[i].permute(1,2,0).detach().numpy()
    
            if gray:
                img = np.concatenate((img,)*3,axis = -1)
                img_r = np.concatenate((img_r,)*3,axis = -1)
                
            ax[i][0].imshow(img,cmap = "gray")
            ax[i][0].set_title("Original")
    
            ax[i][1].imshow(img_r,cmap = "gray")
            ax[i][1].set_title("Reconstruction")
        
        self.model.cuda()
        if return_fig:
            return fig
        
    def set_dec_modules(self,dec_modules:list):
        self.dec_modules = dec_modules

    def plot_rec_steps(self,x=None,random=False,n=5,gray = True,return_fig=False):
        """Plot a """
        self.model.cpu()
        if not isinstance(x,torch.Tensor):
            x,y = self.data.one_batch()
        
        bs = self.data.batch_size
        
        assert n < bs, f"Number of images to plot larger than batch size {n}>{bs}"
        
        if random:
            idx = np.random.choice(bs,n,replace = False)
        else:
            idx = np.arange(n)
        x = x[idx]
        
        modules = self.dec_modules
        
        with Hooks(modules,register_output) as hooks:
            x_rec = self.model(x)
            outputs = hooks.stored
            
            n_layers = len(outputs)
        
            fig,ax = plt.subplots(n,n_layers+1,figsize = (20,20))

            for i in range(n):
                img = x[i].permute(1,2,0).numpy()
                if gray : img = gray_to_rgb(img)
                
                ax[i][0].imshow(img,cmap = "gray")
                ax[i][0].set_title("Original")
                
                for j in range(n_layers):
                    img_r = outputs[j][i].mean(dim=0,keepdim=True).permute(1,2,0).detach().numpy()
                    if gray : img_r = gray_to_rgb(img_r)
                
                    ax[i][j+1].imshow(img_r,cmap = "gray")
                    ax[i][j+1].set_title(f"Step {j}")
            self.model.cuda()
            if return_fig:
                return fig



learn = Learner(dls, arch, loss_func=MSELossFlat(),metrics=[MSELossFlat])
# We wrap the whole thing in a learner, and add a hook for the KL loss
learn = VisionAELearner(data,rec_loss,enc,bn,dec)
kl_hook = VAEHook(learn,beta=1)

# We add this code to plot the reconstructions
dec_modules = list(learn.dec[1].children())
#############
# mnist_vae.py

#import torch.nn as nn
#import torch
#from fastai_autoencoder.util import *
#from fastai_autoencoder.decoder import SpatialDecoder2D


#bottleneck.py
import torch.nn as nn
import torch
import torch.nn.functional as F

class NormalSampler(nn.Module):
    def __init__(self):
        super(NormalSampler,self).__init__()
    
    def forward(self,mu,logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mu + eps * std
        return z

class VAELinear(nn.Module):
    def __init__(self,in_features,out_features):
        super(VAELinear,self).__init__()
        self.linear = nn.Linear(in_features,out_features)
        self.mu_linear = nn.Linear(out_features,out_features)
        self.logvar_linear = nn.Linear(out_features,out_features)

    def forward(self,x):
        x = self.linear(x)
        mu = self.mu_linear(x)
        logvar = self.logvar_linear(x)

        return mu,logvar

class VAELayer(nn.Module):
    def __init__(self,in_features,out_features):
        super(VAELayer,self).__init__()
        self.blinear = VAELinear(in_features,out_features)
        self.sampler = NormalSampler()

    def forward(self,x):
        mu,logvar = self.blinear(x)
        z = self.sampler(mu,logvar)
        return z

# class VQVAEBottleneck(nn.Module):
#     def __init__(self, num_embeddings, embedding_dim, commitment_cost=0.25):
#         super(VQVAEBottleneck, self).__init__()
        
#         self._embedding_dim = embedding_dim
#         self._num_embeddings = num_embeddings
        
#         self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
#         self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings)
#         self._commitment_cost = commitment_cost

#         self.inferring = False
        
#     def forward(self, inputs):
#         # convert inputs from BxCxHxW to BxHxWxC

#         inputs = inputs.permute(0, 2, 3, 1).contiguous()
        
#         # Store the BxHxWxC shape
#         input_shape = inputs.shape
        
#         # Flatten input from BxHxWxC to BHWxC
#         flat_input = inputs.view(-1, self._embedding_dim)
        
#         # Calculate distances between our input BHWxC and the embeddings KxC which outputs a BHWxK matrix
#         distances = (torch.sum(flat_input**2, dim=1, keepdim=True) 
#                     + torch.sum(self._embedding.weight**2, dim=1)
#                     - 2 * torch.matmul(flat_input, self._embedding.weight.t()))
            
#         # Computing the closest BHWxK to BHWx1
#         encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
        
#         # If we are on inferrence mode we just output the discrete encoding in shape BxHxW
#         if self.inferring:
#             return encoding_indices.view(input_shape[:-1])
        
#         # We reshape from BHWxC to BxHxWxC
#         quantized = self._embedding(encoding_indices).view(input_shape)
        
#         # We compute the loss
#         e_latent_loss = torch.mean((quantized.detach() - inputs)**2)
#         q_latent_loss = torch.mean((quantized - inputs.detach())**2)
#         self.loss = q_latent_loss + self._commitment_cost * e_latent_loss
        
#         # We switch the gradients
#         quantized = inputs + (quantized - inputs).detach()
        
#         # We permute from BxHxWxC to BxCxHxW
#         quantized = quantized.permute(0,3,1,2).contiguous()
        
#         return quantized

class VAEBottleneck(nn.Module):
    def __init__(self,nfs:list):
        super(VAEBottleneck,self).__init__()

        n = len(nfs)
        layers = [nn.Linear(in_features=nfs[i], out_features=nfs[i+1]) if i < n -2 
             else VAELayer(in_features=nfs[i], out_features=nfs[i+1]) for i in range(n-1)]
        # just makes a single VAELayer for nfs=[128,32]
        self.fc = nn.Sequential(*layers)
    
    def forward(self,x):
        bs = x.size(0)
        x = x.view(bs, -1)
        z = self.fc(x)
        
        return z

class Bottleneck(nn.Module):
    def __init__(self,nfs,activation):
        super(Bottleneck,self).__init__()

        n = len(nfs)
        layers = [nn.Linear(in_features=nfs[i], out_features=nfs[i+1]) for i in range(n-1)]

        self.fc = nn.Sequential(*layers)
    
    def forward(self,x):
        bs = x.size(0)
        x = x.view(bs, -1)
        z = self.fc(x)
        
        return z


    
######## callback.py
from fastai.callbacks import LearnerCallback
from fastai.basic_train import Learner
from fastai.callbacks.hooks import HookCallback
from fastai_autoencoder.bottleneck import VAELinear
import torch.nn as nn
import torch
import numpy as np

def get_layer(m,buffer,layer):
    """Function which takes a list and a model append the elements"""
    for c in m.children():
        if isinstance(c,layer):
            buffer.append(c)
        get_layer(c,buffer,layer)

class ReplaceTargetCallback(LearnerCallback):
    """Callback to modify the loss of the learner to compute the loss against x"""
    _order = 9999
    
    def __init__(self,learn:Learner):
        super().__init__(learn)
        
    def on_batch_begin(self,last_input,last_target,train,**kwargs):
        # We keep the original x to compute the reconstruction loss
        if not self.learn.inferring:
            return {"last_input" : last_input,"last_target" : last_input}
        else:
            return {"last_input" : last_input,"last_target" : last_target} 
        
# class VQVAEHook(HookCallback):
#     """Callback to modify the loss of the learner to compute the loss against x"""
#     _order = 10000
    
#     def __init__(self, learn,beta = 1,do_remove:bool=True):
#         super().__init__(learn)
        
#         # We look for the VAE bottleneck layer
#         self.learn = learn
        
#         self.loss = []
        
#         buffer = []
#         get_layer(learn.model,buffer,VQVAEBottleneck)
#         if not buffer:
#             raise NotImplementedError("No VQ VAE Bottleneck found")
            
#         self.modules = buffer
#         self.do_remove = do_remove
        
#     def on_backward_begin(self,last_loss,**kwargs):
#         total_loss = last_loss + self.current_loss
        
#         return {"last_loss" : total_loss}
    
#     def hook(self, m:nn.Module, i, o):
#         "Save the latents of the bottleneck"
#         self.current_loss = m.loss
#         self.loss.append(m.loss)
        
class VAEHook(HookCallback):
    """Hook to register the parameters of the latents during the forward pass to compute the KL term of the VAE"""
    
    def __init__(self, learn,beta = 1,do_remove:bool=True):
        super().__init__(learn)
        
        # We look for the VAE bottleneck layer
        self.learn = learn
        self.beta = beta
        
        self.loss = []
        
        buffer = []
        get_layer(learn.model,buffer,VAELinear)
        if not buffer:
            raise NotImplementedError("No Bayesian Linear found")
            
        self.modules = buffer
        self.do_remove = do_remove
    
    def on_backward_begin(self,last_loss,**kwargs):
        n = len(self.learn.mu)
        # We add the KL term of each Bayesian Linear Layer
        
        
        mu = self.learn.mu
        logvar = self.learn.logvar
        kl = self.beta * (-0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(dim=-1).mean())
        
        total_loss = last_loss + kl
        self.loss.append({"rec_loss":last_loss.cpu().detach().numpy(),"kl_loss":kl.cpu().detach().numpy(),
                          "total_loss":total_loss.cpu().detach().numpy()})
        
        return {"last_loss" : total_loss}
        
    def hook(self, m:nn.Module, i, o):
        "Save the latents of the bottleneck"
        mu,logvar = o
        self.learn.mu = mu
        self.learn.logvar = logvar

class HighFrequencyLoss(LearnerCallback):
    def __init__(self, learn,low_ratio = 0.15,threshold = 1e-2,mul=1,scaling=True,debug=False):
        super().__init__(learn)
        
        # We look for the VAE bottleneck layer
        assert low_ratio < 0.5, "Low ratio too high"
        self.low_ratio = low_ratio
        self.window_size = int(28 * low_ratio)
        self.threshold = threshold
        self.scaling = scaling
        self.mul = mul
        self.debug = debug
        
    def get_exponent(self,x): return np.floor(np.log10(np.abs(x))).astype(int)
        
    def on_backward_begin(self,last_loss,**kwargs):
        
        x = kwargs["last_input"]
        x_rec = kwargs["last_output"]
        
        # First we get the fft of the batch
        f = np.fft.fft2(x.squeeze(1))
        fshift = np.fft.fftshift(f,axes=(1,2))

        # We zero the low frequencies
        rows, cols = x.shape[-2], x.shape[-1]
        crow,ccol = rows//2 , cols//2
        fshift[:,crow-self.window_size:crow+self.window_size, ccol-self.window_size:ccol+self.window_size] = 0
        
        # We reconstruct the image
        f_ishift = np.fft.ifftshift(fshift,axes=(1,2))
        img_back = np.fft.ifft2(f_ishift)
        
        # We keep the indexes of pixels with high values
        img_back = np.abs(img_back)
        img_back = img_back / img_back.sum(axis=(1,2),keepdims=True)
        idx = (img_back > self.threshold)
        
        img_back = torch.tensor(img_back[idx]).cuda()
        mask = torch.ByteTensor(idx).cuda()
        
        # We select only the pixels with high values
        x_hf = torch.masked_select(x.view_as(mask),mask)
        x_rec_hf = torch.masked_select(x_rec.view_as(mask),mask)
        
        bs = x.shape[0]
        diff = img_back * self.learn.rec_loss(x_hf,x_rec_hf)
        
        hf_loss = diff.sum() / bs
        
        # If we scale it we put both losses on the same scale
        if self.scaling:
            exponent = max(0,self.get_exponent(last_loss.item()) - self.get_exponent(hf_loss.item()))
            rescale_factor = 10**(exponent)
            hf_loss *= rescale_factor
            
        hf_loss *= self.mul
        total_loss = last_loss + hf_loss
        
        output = {"last_loss" : total_loss}
        if self.debug:
            print(f"Using High Frequency Loss")
            print(f"Loss before : {last_loss}")
            print(f"High frequency loss : {hf_loss}")
            print(output)
        return output
    
#util.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from fastai.torch_core import requires_grad
from fastai.callbacks import Hooks
from fastai.torch_core import flatten_model
import numpy as np


class ConvBnRelu(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size = 3, stride = 1, bias = False, conv = nn.Conv2d,
                 bn = nn.BatchNorm2d, act_fn = nn.ReLU,**kwargs):
        super(ConvBnRelu,self).__init__()
        self.conv = conv(in_channels=in_channels,out_channels=out_channels,
                         kernel_size=kernel_size,stride = stride, bias=bias,**kwargs)

        if bn:
            self.bn = bn(out_channels)
        else:
            self.bn = None
        if act_fn:
            self.act_fn = act_fn(inplace = True)
        else:
            self.act_fn = None
    def forward(self,x):
        output = self.conv(x)
        if self.bn:
            output = self.bn(output)
        if self.act_fn:
            output = self.act_fn(output)
        return output

  


class DepthwiseConv(nn.Module):
    def __init__(self,in_channels,kernel_size=3,stride=1,bias=True,conv = nn.Conv2d,**kwargs):
        super(DepthwiseConv,self).__init__()
        self.conv = conv(in_channels=in_channels,out_channels=in_channels,kernel_size=kernel_size,
                         padding=kernel_size //2,stride = stride,groups = in_channels, bias = bias)

    def forward(self,x):
        output = self.conv(x)
        return output


class MobileConv(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size=3,stride = 1,bias = True,conv = nn.Conv2d,
                dw_conv = DepthwiseConv, pw_conv = PointwiseConv,**kwargs):
        super(MobileConv,self).__init__()
        self.depth_conv = dw_conv(in_channels,kernel_size,stride,bias,conv)
        self.point_conv = pw_conv(in_channels,out_channels,bias,conv,)

    def forward(self,x):
        output = self.point_conv(self.depth_conv(x))
        return output

def zero_bias(conv_layer,conv_type = nn.Conv2d):
    for c in conv_layer.children():
        if isinstance(c,conv_type):
            c.bias.requires_grad = False
            c.bias.data = torch.zeros_like(c.bias)
        else:
            zero_bias(c)
  
    
def register_stats(m,i,o):
    d = o.data
    mean, std = d.mean().item(),d.std().item()
    output = {"mean":mean,"std":std,"m":m}
    return output

def lsuv_init(model,xb,layers=None,types = (nn.Linear,nn.Conv2d),unfreeze = True):
    """LSUV init of the model.
    """
    output = []
    
    # If no layers are specified we will grab the CNN and Linear layers
    if not layers:
        layers = flatten_model(model)
        layers = [layer for layer in layers if isinstance(layer,types) ]
    
    requires_grad(model,False)
    print("Freezing all layers")
    with Hooks(layers,register_stats) as hooks:
        for i,hook in enumerate(hooks):
            # We first get the module
            model(xb)
            m = hook.stored["m"]
            
            while model(xb) is not None and abs(hook.stored["mean"])  > 1e-3: m.bias.data -= hook.stored["mean"]
            while model(xb) is not None and abs(hook.stored["std"]-1) > 1e-3: m.weight.data /= hook.stored["std"]
            
            output.append(f"Layer {i} named {str(m)} with m={hook.stored['mean']},std={hook.stored['std']}")
    
    if unfreeze:
        print("Unfreezing all layers")
        requires_grad(model,True)
    return output

# # decoder.py
# import torch.nn as nn
# import torch



# from fastai_autoencoder.util import PointwiseConv
# class SpatialDecoder1D(nn.Module):
#     def __init__(self,cnn,nfs,activation,im_size):
#         super(Decoder1D,self).__init__()
#         self.cnn = cnn
        
#         # Coordinates for the broadcast decoder
#         self.im_size = im_size
#         x = torch.linspace(-1, 1, im_size)
#         x_grid, y_grid = torch.meshgrid(x, y)
#         # Add as constant, with extra dims for N and C
#         self.register_buffer('x_grid', x_grid.view((1, 1) + x_grid.shape))

#         n = len(nfs)
#         # We add two channels to the first layer to include x and y channels
#         first_layer = PointwiseConv(nfs[0]+2,nfs[1],activation,conv = cnn)

#         conv_layers = [first_layer] + [cnn(nfs[i],nfs[i+1],activation, 3) for i in range(1,n - 1)]        
#         self.dec_convs = nn.Sequential(*conv_layers)
        
#     def forward(self,z):
#         batch_size = z.size(0)
#         # View z as 4D tensor to be tiled across new H and W dimensions
#         # Shape: NxDx1x1
#         z = z.view(z.shape + (1, 1))

#         # Tile across to match image size
#         # Shape: NxDx64x64
#         z = z.expand(-1, -1, self.im_size)

#         # Expand grids to batches and concatenate on the channel dimension
#         # Shape: Nx(D+2)x64x64
#         x = torch.cat((self.x_grid.expand(batch_size, -1, -1, -1), z), dim=1)

#         x = self.dec_convs(x)

#         return x
    
#learner.py    
from fastai.basic_train import Learner
from fastai.basic_data import DataBunch, DatasetType
from fastai.basic_train import get_preds
from fastai.callback import CallbackHandler

from fastai_autoencoder.callback import ReplaceTargetCallback, VAEHook

import torch.nn as nn
import torch
import gc
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np


# 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,
    "c_start" :4
}

# Bottleneck architecture
bn_fn = VAEBottleneck
bn_args = {
    "nfs":[128,32]
}

# Decoder architecture
dec_fn = create_decoder
dec_args = {
    "nfs":[32,64,32,16,8,4,2,1],
    "ks":[3,1,3,1,3,1],   
    "size": 28
}

# We create each part of the autoencoder
enc = enc_fn(**enc_args)
bn = bn_fn(**bn_args)
dec = dec_fn(**dec_args)

# We wrap the whole thing in a learner, and add a hook for the KL loss
learn = VisionAELearner(data,rec_loss,enc,bn,dec)
kl_hook = VAEHook(learn,beta=1)

# We add this code to plot the reconstructions
dec_modules = list(learn.dec[1].children())
## ---------------------- ResNet VAE ---------------------- ##

class ResNet_VAE(nn.Module):
    def __init__(self, fc_hidden1=1024, fc_hidden2=768, drop_p=0.3, CNN_embed_dim=256):
        super(ResNet_VAE, self).__init__()

        self.fc_hidden1, self.fc_hidden2, self.CNN_embed_dim = fc_hidden1, fc_hidden2, CNN_embed_dim

        # CNN architechtures
        self.ch1, self.ch2, self.ch3, self.ch4 = 16, 32, 64, 128
        self.k1, self.k2, self.k3, self.k4 = (5, 5), (3, 3), (3, 3), (3, 3)      # 2d kernal size
        self.s1, self.s2, self.s3, self.s4 = (2, 2), (2, 2), (2, 2), (2, 2)      # 2d strides
        self.pd1, self.pd2, self.pd3, self.pd4 = (0, 0), (0, 0), (0, 0), (0, 0)  # 2d padding

        # encoding components
        resnet = models.resnet152(pretrained=True)
        modules = list(resnet.children())[:-1]      # delete the last fc layer.
        self.resnet = nn.Sequential(*modules)
        self.fc1 = nn.Linear(resnet.fc.in_features, self.fc_hidden1)
        self.bn1 = nn.BatchNorm1d(self.fc_hidden1, momentum=0.01)
        self.fc2 = nn.Linear(self.fc_hidden1, self.fc_hidden2)
        self.bn2 = nn.BatchNorm1d(self.fc_hidden2, momentum=0.01)
        # Latent vectors mu and sigma
        self.fc3_mu = nn.Linear(self.fc_hidden2, self.CNN_embed_dim)      # output = CNN embedding latent variables
        self.fc3_logvar = nn.Linear(self.fc_hidden2, self.CNN_embed_dim)  # output = CNN embedding latent variables

        # Sampling vector
        self.fc4 = nn.Linear(self.CNN_embed_dim, self.fc_hidden2)
        self.fc_bn4 = nn.BatchNorm1d(self.fc_hidden2)
        self.fc5 = nn.Linear(self.fc_hidden2, 64 * 4 * 4)
        self.fc_bn5 = nn.BatchNorm1d(64 * 4 * 4)
        self.relu = nn.ReLU(inplace=True)

        # Decoder
        self.convTrans6 = nn.Sequential(
            nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=self.k4, stride=self.s4,
                               padding=self.pd4),
            nn.BatchNorm2d(32, momentum=0.01),
            nn.ReLU(inplace=True),
        )
        self.convTrans7 = nn.Sequential(
            nn.ConvTranspose2d(in_channels=32, out_channels=8, kernel_size=self.k3, stride=self.s3,
                               padding=self.pd3),
            nn.BatchNorm2d(8, momentum=0.01),
            nn.ReLU(inplace=True),
        )

        self.convTrans8 = nn.Sequential(
            nn.ConvTranspose2d(in_channels=8, out_channels=3, kernel_size=self.k2, stride=self.s2,
                               padding=self.pd2),
            nn.BatchNorm2d(3, momentum=0.01),
            nn.Sigmoid()    # y = (y1, y2, y3) \in [0 ,1]^3
        )


    def encode(self, x):
        x = self.resnet(x)  # ResNet
        x = x.view(x.size(0), -1)  # flatten output of conv

        # FC layers
        x = self.bn1(self.fc1(x))
        x = self.relu(x)
        x = self.bn2(self.fc2(x))
        x = self.relu(x)
        # x = F.dropout(x, p=self.drop_p, training=self.training)
        mu, logvar = self.fc3_mu(x), self.fc3_logvar(x)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = torch.autograd.Variable(std.data.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        x = self.relu(self.fc_bn4(self.fc4(z)))
        x = self.relu(self.fc_bn5(self.fc5(x))).view(-1, 64, 4, 4)
        x = self.convTrans6(x)
        x = self.convTrans7(x)
        x = self.convTrans8(x)
        x = F.interpolate(x, size=(224, 224), mode='bilinear')
        return x

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_reconst = self.decode(z)

        return x_reconst, z, mu, logvar

def conv2D_output_size(img_size, padding, kernel_size, stride):
    # compute output shape of conv2D
    outshape = (np.floor((img_size[0] + 2 * padding[0] - (kernel_size[0] - 1) - 1) / stride[0] + 1).astype(int),
                np.floor((img_size[1] + 2 * padding[1] - (kernel_size[1] - 1) - 1) / stride[1] + 1).astype(int))
    return outshape

def convtrans2D_output_size(img_size, padding, kernel_size, stride):
    # compute output shape of conv2D
    outshape = ((img_size[0] - 1) * stride[0] - 2 * padding[0] + kernel_size[0],
                (img_size[1] - 1) * stride[1] - 2 * padding[1] + kernel_size[1])
    return outshape


# EncoderCNN architecture
CNN_fc_hidden1, CNN_fc_hidden2 = 512, 512
CNN_embed_dim = 128     # latent dim extracted by 2D CNN
res_size = 160        # ResNet image size
dropout_p = 0.2       # dropout probability
# Create model
resnet_vae = ResNet_VAE(fc_hidden1=CNN_fc_hidden1, fc_hidden2=CNN_fc_hidden2, drop_p=dropout_p, CNN_embed_dim=CNN_embed_dim).cuda()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13396139705882354, 'fraction_numerical': 0.024528952205882353, 'mean_word_length': 2.4805064978340554, 'pattern_counts': {'":': 30, '<': 22, '<?xml version=': 0, '>': 25, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 15, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '438679', 'license': 'Apache-2.0', 'n_tokens_mistral': 19294, 'n_tokens_neox': 18367, 'n_words': 3644, 'path': 'archive/_AutoEncoder.ipynb', 'repo_name': 'ergonyc/snkr-finder', 'size': '52224'}