Myosotis-Researches 0.0.14__tar.gz → 0.0.16__tar.gz

{myosotis_researches-0.0.14 → myosotis_researches-0.0.16/Myosotis_Researches.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Myosotis-Researches
-Version: 0.0.14
+Version: 0.0.16
 Summary: A repository for storing my progress of researches.
 Home-page: https://github.com/Zeyu-Xie/Myosotis-Researches
 Author: Zeyu Xie

{myosotis_researches-0.0.14 → myosotis_researches-0.0.16}/Myosotis_Researches.egg-info/SOURCES.txt

@@ -24,6 +24,18 @@ myosotis_researches/CcGAN/models_256/__init__.py
 myosotis_researches/CcGAN/models_256/autoencoder.py
 myosotis_researches/CcGAN/models_256/cGAN_SAGAN.py
 myosotis_researches/CcGAN/models_256/cGAN_concat_SAGAN.py
+myosotis_researches/CcGAN/train_128/DiffAugment_pytorch.py
+myosotis_researches/CcGAN/train_128/__init__.py
+myosotis_researches/CcGAN/train_128/eval_metrics.py
+myosotis_researches/CcGAN/train_128/opts.py
+myosotis_researches/CcGAN/train_128/pretrain_AE.py
+myosotis_researches/CcGAN/train_128/pretrain_CNN_class.py
+myosotis_researches/CcGAN/train_128/pretrain_CNN_regre.py
+myosotis_researches/CcGAN/train_128/train_ccgan.py
+myosotis_researches/CcGAN/train_128/train_cgan.py
+myosotis_researches/CcGAN/train_128/train_cgan_concat.py
+myosotis_researches/CcGAN/train_128/train_net_for_label_embed.py
+myosotis_researches/CcGAN/train_128/utils.py
 myosotis_researches/CcGAN/utils/__init__.py
 myosotis_researches/CcGAN/utils/concat_image_horizontal.py
 myosotis_researches/CcGAN/utils/concat_image_vertical.py

{myosotis_researches-0.0.14/Myosotis_Researches.egg-info → myosotis_researches-0.0.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Myosotis-Researches
-Version: 0.0.14
+Version: 0.0.16
 Summary: A repository for storing my progress of researches.
 Home-page: https://github.com/Zeyu-Xie/Myosotis-Researches
 Author: Zeyu Xie

myosotis_researches-0.0.16/myosotis_researches/CcGAN/train_128/DiffAugment_pytorch.py

@@ -0,0 +1,76 @@
+# Differentiable Augmentation for Data-Efficient GAN Training
+# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han
+# https://arxiv.org/pdf/2006.10738
+
+import torch
+import torch.nn.functional as F
+
+
+def DiffAugment(x, policy='', channels_first=True):
+    if policy:
+        if not channels_first:
+            x = x.permute(0, 3, 1, 2)
+        for p in policy.split(','):
+            for f in AUGMENT_FNS[p]:
+                x = f(x)
+        if not channels_first:
+            x = x.permute(0, 2, 3, 1)
+        x = x.contiguous()
+    return x
+
+
+def rand_brightness(x):
+    x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
+    return x
+
+
+def rand_saturation(x):
+    x_mean = x.mean(dim=1, keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean
+    return x
+
+
+def rand_contrast(x):
+    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean
+    return x
+
+
+def rand_translation(x, ratio=0.125):
+    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
+    translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(x.size(2), dtype=torch.long, device=x.device),
+        torch.arange(x.size(3), dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
+    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
+    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
+    x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
+    return x
+
+
+def rand_cutout(x, ratio=0.5):
+    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
+    offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
+    grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
+    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
+    mask[grid_batch, grid_x, grid_y] = 0
+    x = x * mask.unsqueeze(1)
+    return x
+
+
+AUGMENT_FNS = {
+    'color': [rand_brightness, rand_saturation, rand_contrast],
+    'translation': [rand_translation],
+    'cutout': [rand_cutout],
+}

myosotis_researches-0.0.16/myosotis_researches/CcGAN/train_128/eval_metrics.py

@@ -0,0 +1,205 @@
+"""
+Compute
+Inception Score (IS),
+Frechet Inception Discrepency (FID), ref "https://github.com/mseitzer/pytorch-fid/blob/master/fid_score.py"
+Maximum Mean Discrepancy (MMD)
+for a set of fake images
+
+use numpy array
+Xr: high-level features for real images; nr by d array
+Yr: labels for real images
+Xg: high-level features for fake images; ng by d array
+Yg: labels for fake images
+IMGSr: real images
+IMGSg: fake images
+
+"""
+
+import os
+import gc
+import numpy as np
+# from numpy import linalg as LA
+from scipy import linalg
+import torch
+import torch.nn as nn
+from scipy.stats import entropy
+from torch.nn import functional as F
+from torchvision.utils import save_image
+
+from .utils import SimpleProgressBar, IMGs_dataset
+
+
+def normalize_images(batch_images):
+    batch_images = batch_images/255.0
+    batch_images = (batch_images - 0.5)/0.5
+    return batch_images
+
+##############################################################################
+# FID scores
+##############################################################################
+# compute FID based on extracted features
+def FID(Xr, Xg, eps=1e-10):
+    '''
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+    '''
+    #sample mean
+    MUr = np.mean(Xr, axis = 0)
+    MUg = np.mean(Xg, axis = 0)
+    mean_diff = MUr - MUg
+    #sample covariance
+    SIGMAr = np.cov(Xr.transpose())
+    SIGMAg = np.cov(Xg.transpose())
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(SIGMAr.dot(SIGMAg), disp=False)#square root of a matrix
+    covmean = covmean.real
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(SIGMAr.shape[0]) * eps
+        covmean = linalg.sqrtm((SIGMAr + offset).dot(SIGMAg + offset))
+
+    #fid score
+    fid_score = mean_diff.dot(mean_diff) + np.trace(SIGMAr + SIGMAg - 2*covmean)
+
+    return fid_score
+
+##test
+#Xr = np.random.rand(10000,1000)
+#Xg = np.random.rand(10000,1000)
+#print(FID(Xr, Xg))
+
+# compute FID from raw images
+def cal_FID(PreNetFID, IMGSr, IMGSg, batch_size = 500, resize = None, norm_img = False):
+    #resize: if None, do not resize; if resize = (H,W), resize images to 3 x H x W
+    
+    PreNetFID.eval()
+
+    nr = IMGSr.shape[0]
+    ng = IMGSg.shape[0]
+
+    nc = IMGSr.shape[1] #IMGSr is nrxNCxIMG_SIExIMG_SIZE
+    img_size = IMGSr.shape[2]
+
+    if batch_size > min(nr, ng):
+        batch_size = min(nr, ng)
+        # print("FID: recude batch size to {}".format(batch_size))
+
+    #compute the length of extracted features
+    with torch.no_grad():
+        test_img = torch.from_numpy(IMGSr[0].reshape((1,nc,img_size,img_size))).type(torch.float).cuda()
+        if resize is not None:
+            test_img = nn.functional.interpolate(test_img, size = resize, scale_factor=None, mode='bilinear', align_corners=False)
+        if norm_img:
+            test_img = normalize_images(test_img)
+        # _, test_features = PreNetFID(test_img)
+        test_features = PreNetFID(test_img)
+        d = test_features.shape[1] #length of extracted features
+
+    Xr = np.zeros((nr, d))
+    Xg = np.zeros((ng, d))
+
+    #batch_size = 500
+    with torch.no_grad():
+        tmp = 0
+        pb1 = SimpleProgressBar()
+        for i in range(nr//batch_size):
+            imgr_tensor = torch.from_numpy(IMGSr[tmp:(tmp+batch_size)]).type(torch.float).cuda()
+            if resize is not None:
+                imgr_tensor = nn.functional.interpolate(imgr_tensor, size = resize, scale_factor=None, mode='bilinear', align_corners=False)
+            if norm_img:
+                imgr_tensor = normalize_images(imgr_tensor)
+            # _, Xr_tmp = PreNetFID(imgr_tensor)
+            Xr_tmp = PreNetFID(imgr_tensor)
+            Xr[tmp:(tmp+batch_size)] = Xr_tmp.detach().cpu().numpy()
+            tmp+=batch_size
+            # pb1.update(min(float(i)*100/(nr//batch_size), 100))
+            pb1.update(min(max(tmp/nr*100,100), 100))
+        del Xr_tmp,imgr_tensor; gc.collect()
+        torch.cuda.empty_cache()
+
+        tmp = 0
+        pb2 = SimpleProgressBar()
+        for j in range(ng//batch_size):
+            imgg_tensor = torch.from_numpy(IMGSg[tmp:(tmp+batch_size)]).type(torch.float).cuda()
+            if resize is not None:
+                imgg_tensor = nn.functional.interpolate(imgg_tensor, size = resize, scale_factor=None, mode='bilinear', align_corners=False)
+            if norm_img:
+                imgg_tensor = normalize_images(imgg_tensor)
+            # _, Xg_tmp = PreNetFID(imgg_tensor)
+            Xg_tmp = PreNetFID(imgg_tensor)
+            Xg[tmp:(tmp+batch_size)] = Xg_tmp.detach().cpu().numpy()
+            tmp+=batch_size
+            # pb2.update(min(float(j)*100/(ng//batch_size), 100))
+            pb2.update(min(max(tmp/ng*100, 100), 100))
+        del Xg_tmp,imgg_tensor; gc.collect()
+        torch.cuda.empty_cache()
+
+
+    fid_score = FID(Xr, Xg, eps=1e-6)
+
+    return fid_score
+
+
+
+
+
+
+##############################################################################
+# label_score
+# difference between assigned label and predicted label
+##############################################################################
+def cal_labelscore(PreNet, images, labels_assi, min_label_before_shift, max_label_after_shift, batch_size = 500, resize = None, norm_img = False, num_workers=0):
+    '''
+    PreNet: pre-trained CNN
+    images: fake images
+    labels_assi: assigned labels
+    resize: if None, do not resize; if resize = (H,W), resize images to 3 x H x W
+    '''
+
+    PreNet.eval()
+
+    # assume images are nxncximg_sizeximg_size
+    n = images.shape[0]
+    nc = images.shape[1] #number of channels
+    img_size = images.shape[2]
+    labels_assi = labels_assi.reshape(-1)
+
+    eval_trainset = IMGs_dataset(images, labels_assi, normalize=False)
+    eval_dataloader = torch.utils.data.DataLoader(eval_trainset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+    labels_pred = np.zeros(n+batch_size)
+
+    nimgs_got = 0
+    pb = SimpleProgressBar()
+    for batch_idx, (batch_images, batch_labels) in enumerate(eval_dataloader):
+        batch_images = batch_images.type(torch.float).cuda()
+        batch_labels = batch_labels.type(torch.float).cuda()
+        batch_size_curr = len(batch_labels)
+
+        if norm_img:
+            batch_images = normalize_images(batch_images)
+
+        batch_labels_pred, _ = PreNet(batch_images)
+        labels_pred[nimgs_got:(nimgs_got+batch_size_curr)] = batch_labels_pred.detach().cpu().numpy().reshape(-1)
+
+        nimgs_got += batch_size_curr
+        pb.update((float(nimgs_got)/n)*100)
+
+        del batch_images; gc.collect()
+        torch.cuda.empty_cache()
+    #end for batch_idx
+
+    labels_pred = labels_pred[0:n]
+
+
+    labels_pred = (labels_pred*max_label_after_shift)-np.abs(min_label_before_shift)
+    labels_assi = (labels_assi*max_label_after_shift)-np.abs(min_label_before_shift)
+
+    ls_mean = np.mean(np.abs(labels_pred-labels_assi))
+    ls_std = np.std(np.abs(labels_pred-labels_assi))
+
+    return ls_mean, ls_std

myosotis_researches-0.0.16/myosotis_researches/CcGAN/train_128/opts.py

@@ -0,0 +1,87 @@
+import argparse
+
+def parse_opts():
+    parser = argparse.ArgumentParser()
+
+    ''' Overall Settings '''
+    parser.add_argument('--root_path', type=str, default='')
+    parser.add_argument('--data_path', type=str, default='')
+    parser.add_argument('--eval_ckpt_path', type=str, default='')
+    parser.add_argument('--seed', type=int, default=2021, metavar='S', help='random seed (default: 2020)')
+    parser.add_argument('--num_workers', type=int, default=0)
+
+
+    ''' Dataset '''
+    ## Data split: RC-49 is split into a train set (the last decimal of the degree is odd) and a test set (the last decimal of the degree is even); the unique labels in two sets do not overlap.
+    parser.add_argument('--data_split', type=str, default='train',
+                        choices=['all', 'train'])
+    parser.add_argument('--min_label', type=float, default=0.0)
+    parser.add_argument('--max_label', type=float, default=90.0)
+    parser.add_argument('--num_channels', type=int, default=3, metavar='N')
+    parser.add_argument('--img_size', type=int, default=128, metavar='N')
+    parser.add_argument('--max_num_img_per_label', type=int, default=50, metavar='N')
+    parser.add_argument('--max_num_img_per_label_after_replica', type=int, default=0, metavar='N')
+    parser.add_argument('--show_real_imgs', action='store_true', default=False)
+    parser.add_argument('--visualize_fake_images', action='store_true', default=False)
+
+
+    ''' GAN settings '''
+    parser.add_argument('--GAN', type=str, default='CcGAN', choices=['cGAN', 'cGAN-concat', 'CcGAN'])
+    parser.add_argument('--GAN_arch', type=str, default='SAGAN', choices=['SAGAN'])
+
+    # label embedding setting
+    parser.add_argument('--net_embed', type=str, default='ResNet34_embed') #ResNetXX_emebed
+    parser.add_argument('--epoch_cnn_embed', type=int, default=200) #epoch of cnn training for label embedding
+    parser.add_argument('--resumeepoch_cnn_embed', type=int, default=0) #epoch of cnn training for label embedding
+    parser.add_argument('--epoch_net_y2h', type=int, default=500)
+    parser.add_argument('--dim_embed', type=int, default=128) #dimension of the embedding space
+    parser.add_argument('--batch_size_embed', type=int, default=256, metavar='N')
+
+    parser.add_argument('--loss_type_gan', type=str, default='hinge')
+    parser.add_argument('--niters_gan', type=int, default=10000, help='number of iterations')
+    parser.add_argument('--resume_niters_gan', type=int, default=0)
+    parser.add_argument('--save_niters_freq', type=int, default=2000, help='frequency of saving checkpoints')
+    parser.add_argument('--lr_g_gan', type=float, default=1e-4, help='learning rate for generator')
+    parser.add_argument('--lr_d_gan', type=float, default=1e-4, help='learning rate for discriminator')
+    parser.add_argument('--dim_gan', type=int, default=128, help='Latent dimension of GAN')
+    parser.add_argument('--batch_size_disc', type=int, default=64)
+    parser.add_argument('--batch_size_gene', type=int, default=64)
+    parser.add_argument('--num_D_steps', type=int, default=4, help='number of Ds updates in one iteration')
+    parser.add_argument('--cGAN_num_classes', type=int, default=20, metavar='N') #bin label into cGAN_num_classes
+    parser.add_argument('--visualize_freq', type=int, default=2000, help='frequency of visualization')
+
+    parser.add_argument('--kernel_sigma', type=float, default=-1.0,
+                        help='If kernel_sigma<0, then use rule-of-thumb formula to compute the sigma.')
+    parser.add_argument('--threshold_type', type=str, default='hard', choices=['soft', 'hard'])
+    parser.add_argument('--kappa', type=float, default=-1)
+    parser.add_argument('--nonzero_soft_weight_threshold', type=float, default=1e-3,
+                        help='threshold for determining nonzero weights for SVDL; we neglect images with too small weights')
+
+    # DiffAugment setting
+    parser.add_argument('--gan_DiffAugment', action='store_true', default=False)
+    parser.add_argument('--gan_DiffAugment_policy', type=str, default='color,translation,cutout')
+
+
+    # evaluation setting
+    '''
+    Four evaluation modes:
+    Mode 1: eval on unique labels used for GAN training;
+    Mode 2. eval on all unique labels in the dataset and when computing FID use all real images in the dataset;
+    Mode 3. eval on all unique labels in the dataset and when computing FID only use real images for GAN training in the dataset (to test SFID's effectiveness on unseen labels);
+    Mode 4. eval on a interval [min_label, max_label] with num_eval_labels labels.
+    '''
+    parser.add_argument('--eval_mode', type=int, default=2)
+    parser.add_argument('--num_eval_labels', type=int, default=-1)
+    parser.add_argument('--samp_batch_size', type=int, default=200)
+    parser.add_argument('--nfake_per_label', type=int, default=200)
+    parser.add_argument('--nreal_per_label', type=int, default=-1)
+    parser.add_argument('--comp_FID', action='store_true', default=False)
+    parser.add_argument('--epoch_FID_CNN', type=int, default=200)
+    parser.add_argument('--FID_radius', type=float, default=0)
+    parser.add_argument('--FID_num_centers', type=int, default=-1)
+    parser.add_argument('--dump_fake_for_NIQE', action='store_true', default=False,
+                        help='Dump fake images for computing NIQE')
+
+    args = parser.parse_args()
+
+    return args

myosotis_researches-0.0.16/myosotis_researches/CcGAN/train_128/pretrain_AE.py

@@ -0,0 +1,268 @@
+
+import os
+import argparse
+import shutil
+import timeit
+import torch
+import torchvision
+import torchvision.transforms as transforms
+import numpy as np
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+import random
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+from torch import autograd
+from torchvision.utils import save_image
+import csv
+from tqdm import tqdm
+import gc
+import h5py
+
+
+#############################
+# Settings
+#############################
+
+parser = argparse.ArgumentParser(description='Pre-train AE for computing FID')
+parser.add_argument('--root_path', type=str, default='')
+parser.add_argument('--data_path', type=str, default='')
+parser.add_argument('--num_workers', type=int, default=0)
+parser.add_argument('--dim_bottleneck', type=int, default=512)
+parser.add_argument('--epochs', type=int, default=200, metavar='N',
+                    help='number of epochs to train CNNs (default: 200)')
+parser.add_argument('--resume_epoch', type=int, default=0)
+parser.add_argument('--batch_size_train', type=int, default=128, metavar='N',
+                    help='input batch size for training')
+parser.add_argument('--batch_size_valid', type=int, default=10, metavar='N',
+                    help='input batch size for testing')
+parser.add_argument('--base_lr', type=float, default=1e-3,
+                    help='learning rate, default=1e-3')
+parser.add_argument('--lr_decay_epochs', type=int, default=50) #decay lr rate every dre_lr_decay_epochs epochs
+parser.add_argument('--lr_decay_factor', type=float, default=0.1)
+parser.add_argument('--lambda_sparsity', type=float, default=1e-4, help='penalty for sparsity')
+parser.add_argument('--weight_dacay', type=float, default=1e-4,
+                    help='Weigth decay, default=1e-4')
+parser.add_argument('--seed', type=int, default=2020, metavar='S',
+                    help='random seed (default: 1)')
+parser.add_argument('--CVMode', action='store_true', default=False,
+                   help='CV mode?')
+parser.add_argument('--img_size', type=int, default=128, metavar='N')
+parser.add_argument('--min_label', type=float, default=0.0)
+parser.add_argument('--max_label', type=float, default=90.0)
+args = parser.parse_args()
+
+wd = args.root_path
+os.chdir(wd)
+from ..models_128 import *
+from .utils import IMGs_dataset, SimpleProgressBar
+
+# some parameters in the opts
+dim_bottleneck = args.dim_bottleneck
+epochs = args.epochs
+base_lr = args.base_lr
+lr_decay_epochs = args.lr_decay_epochs
+lr_decay_factor = args.lr_decay_factor
+resume_epoch = args.resume_epoch
+lambda_sparsity = args.lambda_sparsity
+
+
+# random seed
+random.seed(args.seed)
+torch.manual_seed(args.seed)
+torch.backends.cudnn.deterministic = True
+cudnn.benchmark = False
+np.random.seed(args.seed)
+
+# directories for checkpoint, images and log files
+save_models_folder = wd + '/output/eval_models'
+os.makedirs(save_models_folder, exist_ok=True)
+save_AE_images_in_train_folder = save_models_folder + '/AE_lambda_{}_images_in_train'.format(lambda_sparsity)
+os.makedirs(save_AE_images_in_train_folder, exist_ok=True)
+save_AE_images_in_valid_folder = save_models_folder + '/AE_lambda_{}_images_in_valid'.format(lambda_sparsity)
+os.makedirs(save_AE_images_in_valid_folder, exist_ok=True)
+
+
+###########################################################################################################
+# Data
+###########################################################################################################
+# data loader
+data_filename = args.data_path + '/Ra_' + str(args.img_size) + 'x' + str(args.img_size) + '.h5'
+hf = h5py.File(data_filename, 'r')
+labels = hf['labels'][:]
+labels = labels.astype(float)
+images = hf['images'][:]
+hf.close()
+N_all = len(images)
+assert len(images) == len(labels)
+
+q1 = args.min_label
+q2 = args.max_label
+indx = np.where((labels>q1)*(labels<q2)==True)[0]
+labels = labels[indx]
+images = images[indx]
+assert len(labels)==len(images)
+
+# define training and validation sets
+if args.CVMode:
+    #90% Training; 10% valdation
+    valid_prop = 0.1 #proportion of the validation samples
+    indx_all = np.arange(len(images))
+    np.random.shuffle(indx_all)
+    indx_valid = indx_all[0:int(valid_prop*len(images))]
+    indx_train = indx_all[int(valid_prop*len(images)):]
+
+    trainset = IMGs_dataset(images[indx_train], labels=None, normalize=True)
+    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=args.num_workers)
+    validset = IMGs_dataset(images[indx_valid], labels=None, normalize=True)
+    validloader = torch.utils.data.DataLoader(validset, batch_size=args.batch_size_valid, shuffle=False, num_workers=args.num_workers)
+
+else:
+    trainset = IMGs_dataset(images, labels=None, normalize=True)
+    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=args.num_workers)
+
+
+###########################################################################################################
+# Necessary functions
+###########################################################################################################
+
+def adjust_learning_rate(epoch, epochs, optimizer, base_lr, lr_decay_epochs, lr_decay_factor):
+    lr = base_lr #1e-4
+
+    for i in range(epochs//lr_decay_epochs):
+        if epoch >= (i+1)*lr_decay_epochs:
+            lr *= lr_decay_factor
+
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+def train_AE():
+
+    # define optimizer
+    params = list(net_encoder.parameters()) + list(net_decoder.parameters())
+    optimizer = torch.optim.Adam(params, lr = base_lr, betas=(0.5, 0.999), weight_decay=1e-4)
+
+    # criterion
+    criterion = nn.MSELoss()
+
+    if resume_epoch>0:
+        print("Loading ckpt to resume training AE >>>")
+        ckpt_fullpath = save_models_folder + "/AE_checkpoint_intrain/AE_checkpoint_epoch_{}_lambda_{}.pth".format(resume_epoch, lambda_sparsity)
+        checkpoint = torch.load(ckpt_fullpath)
+        net_encoder.load_state_dict(checkpoint['net_encoder_state_dict'])
+        net_decoder.load_state_dict(checkpoint['net_decoder_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        torch.set_rng_state(checkpoint['rng_state'])
+        gen_iterations = checkpoint['gen_iterations']
+    else:
+        gen_iterations = 0
+
+    start_time = timeit.default_timer()
+    for epoch in range(resume_epoch, epochs):
+
+        adjust_learning_rate(epoch, epochs, optimizer, base_lr, lr_decay_epochs, lr_decay_factor)
+
+        train_loss = 0
+
+        for batch_idx, batch_real_images in enumerate(trainloader):
+
+            net_encoder.train()
+            net_decoder.train()
+
+            batch_size_curr = batch_real_images.shape[0]
+
+            batch_real_images = batch_real_images.type(torch.float).cuda()
+
+
+            batch_features = net_encoder(batch_real_images)
+            batch_recons_images = net_decoder(batch_features)
+
+            '''
+            based on https://debuggercafe.com/sparse-autoencoders-using-l1-regularization-with-pytorch/
+            '''
+            loss = criterion(batch_recons_images, batch_real_images) + lambda_sparsity * batch_features.mean()
+
+            #backward pass
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.cpu().item()
+
+            gen_iterations += 1
+
+            if gen_iterations % 100 == 0:
+                n_row=min(10, int(np.sqrt(batch_size_curr)))
+                with torch.no_grad():
+                    batch_recons_images = net_decoder(net_encoder(batch_real_images[0:n_row**2]))
+                    batch_recons_images = batch_recons_images.detach().cpu()
+                save_image(batch_recons_images.data, save_AE_images_in_train_folder + '/{}.png'.format(gen_iterations), nrow=n_row, normalize=True)
+
+            if gen_iterations % 20 == 0:
+                print("AE+lambda{}: [step {}] [epoch {}/{}] [train loss {}] [Time {}]".format(lambda_sparsity, gen_iterations, epoch+1, epochs, train_loss/(batch_idx+1), timeit.default_timer()-start_time) )
+        # end for batch_idx
+
+        if (epoch+1) % 50 == 0:
+            save_file = save_models_folder + "/AE_checkpoint_intrain/AE_checkpoint_epoch_{}_lambda_{}.pth".format(epoch+1, lambda_sparsity)
+            os.makedirs(os.path.dirname(save_file), exist_ok=True)
+            torch.save({
+                    'gen_iterations': gen_iterations,
+                    'net_encoder_state_dict': net_encoder.state_dict(),
+                    'net_decoder_state_dict': net_decoder.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'rng_state': torch.get_rng_state()
+            }, save_file)
+    #end for epoch
+
+    return net_encoder, net_decoder
+
+
+if args.CVMode:
+    def valid_AE():
+        net_encoder.eval()
+        net_decoder.eval()
+        with torch.no_grad():
+            for batch_idx, images in enumerate(validloader):
+                images = images.type(torch.float).cuda()
+                features = net_encoder(images)
+                recons_images = net_decoder(features)
+                save_image(recons_images.data, save_AE_images_in_valid_folder + '/{}_recons.png'.format(batch_idx), nrow=10, normalize=True)
+                save_image(images.data, save_AE_images_in_valid_folder + '/{}_real.png'.format(batch_idx), nrow=10, normalize=True)
+        return None
+
+
+
+###########################################################################################################
+# Training and validation
+###########################################################################################################
+
+# model initialization
+net_encoder = encoder(dim_bottleneck=args.dim_bottleneck).cuda()
+net_decoder = decoder(dim_bottleneck=args.dim_bottleneck).cuda()
+net_encoder = nn.DataParallel(net_encoder)
+net_decoder = nn.DataParallel(net_decoder)
+
+filename_ckpt = save_models_folder + '/ckpt_AE_epoch_{}_seed_{}_CVMode_{}.pth'.format(args.epochs, args.seed, args.CVMode)
+
+# training
+if not os.path.isfile(filename_ckpt):
+    print("\n Begin training AE: ")
+    start = timeit.default_timer()
+    net_encoder, net_decoder = train_AE()
+    stop = timeit.default_timer()
+    print("Time elapses: {}s".format(stop - start))
+    # save model
+    torch.save({
+    'net_encoder_state_dict': net_encoder.state_dict(),
+    'net_decoder_state_dict': net_decoder.state_dict(),
+    }, filename_ckpt)
+else:
+    print("\n Ckpt already exists")
+    print("\n Loading...")
+    checkpoint = torch.load(filename_ckpt)
+    net_encoder.load_state_dict(checkpoint['net_encoder_state_dict'])
+    net_decoder.load_state_dict(checkpoint['net_decoder_state_dict'])
+
+if args.CVMode:
+    #validation
+    _ = valid_AE()