def test_VAE_saveState(self):
mb_size = 49 # because dummyDataset98.npz is a 98 data size
epoch_nb = 11
datasetName = 'dummyDataset98.npz'
datasetDir = './dummyDataset/'
testDataset = NPZ_Dataset(datasetName, datasetDir, 'Spectrums')
train_loader = torch.utils.data.DataLoader(testDataset,
batch_size=mb_size,
shuffle=True)
# define vae structure
X_dim = 1024
Z_dim = 6
IOh_dims_Enc = [X_dim, 401, Z_dim]
IOh_dims_Dec = [Z_dim, 399, X_dim]
NL_types_Enc = ['relu6']
NL_types_Dec = ['relu6', 'sigmoid']
vae = VAE(X_dim, Z_dim, IOh_dims_Enc, IOh_dims_Dec, NL_types_Enc,
NL_types_Dec, mb_size)
vae.trainVAE(train_loader, epoch_nb)
# save it
if vae.trained:
vae.save(datasetName, datasetDir)
self.assertTrue(vae.created and vae.trained and vae.saved)
def test_gaussianVAE_trainsaveload(self):
mb_size = 49 # because dummyDataset98.npz is a 98 data size
epoch_nb = 5
# if exists remove 'saveloadTest' folder
if os.path.exists('./saveloadTest'):
shutil.rmtree('./saveloadTest')
# create a VAE
X_dim = 1024
Z_dim = 10
IOh_dims_Enc = [X_dim, 600, Z_dim]
IOh_dims_Dec = [Z_dim, 600, X_dim]
NL_types_Enc = ['relu6']
NL_types_Dec = ['relu6']
vae = VAE(X_dim,
Z_dim,
IOh_dims_Enc,
IOh_dims_Dec,
NL_types_Enc,
NL_types_Dec,
mb_size,
bernoulli=False,
gaussian=True)
# prepare dataset
datasetName = 'dummyDataset98.npz'
datasetDir = './dummyDataset/'
saveDir = './dummySaveTest/'
testDataset = NPZ_Dataset(datasetName, datasetDir, 'Spectrums')
train_loader = torch.utils.data.DataLoader(testDataset,
batch_size=mb_size,
shuffle=True)
# train it for 10 epochs
vae.trainVAE(train_loader, epoch_nb)
# save it
savefile = vae.save(datasetName, saveDir)
# reload the savefile of VAE
vae = loadVAE(savefile, saveDir)
# continue training
# vae.trainVAE(train_loader, 10)
# vae.save(datasetName, datasetDir)
self.assertTrue(vae.created and vae.loaded and vae.saved)
IOh_dims_Dec = args.decoderIOdims
NL_types_Enc = args.encoderNL
NL_types_Dec = args.decoderNL
if args.type == 'bernoulli':
vae = VAE(X_dim, Z_dim, IOh_dims_Enc,
IOh_dims_Dec, NL_types_Enc, NL_types_Dec, mb_size,
bernoulli=True, gaussian=False, beta=beta, Nwu=Nwu, noiseIn=noise, noiseGain=noiseGain)
elif args.type == 'gaussian':
vae = VAE(X_dim, Z_dim, IOh_dims_Enc,
IOh_dims_Dec, NL_types_Enc, NL_types_Dec, mb_size,
bernoulli=False, gaussian=True, beta=beta, Nwu=Nwu, noiseIn=noise, noiseGain=noiseGain)
else:
print("ERROR script: Chose VAE type -> either bernoulli or gaussian")
# train it for N epochs
vae.trainVAE(train_loader, epoch_nb)
# save it
vae.save(datasetName, saveDir)
elif mode == "load":
# get savefile path
directory = args.vae_path
savefile = directory.split("/")[-1]
directory = directory.replace(savefile, "")
# load vae
vaeLoaded = loadVAE(savefile, directory)
# Now do some stuff (e.g. PCA, t-sne, generate ...)
# Here for example it generates 200 images. Each image is 200x1024.
