def main(args):
"""Main train and evaluation function.
Parameters
----------
args: argparse.Namespace
Arguments
"""
formatter = logging.Formatter('%(asctime)s %(levelname)s - %(funcName)s: %(message)s',
"%H:%M:%S")
logger = logging.getLogger(__name__)
logger.setLevel(args.log_level.upper())
stream = logging.StreamHandler()
stream.setLevel(args.log_level.upper())
stream.setFormatter(formatter)
logger.addHandler(stream)
set_seed(args.seed)
device = get_device(is_gpu=not args.no_cuda)
exp_dir = os.path.join(RES_DIR, args.name)
feature_dir = os.path.join(exp_dir, 'training_features')
logger.info("Root directory for saving and loading experiments: {}".format(exp_dir))
if not args.is_eval_only:
create_safe_directory(feature_dir, logger=logger)
# Setting number of epochs to 1, as we need to extract features
args.epochs = 1
args.batch_size = 1
# PREPARES DATA
data_loader = get_dataloaders(args.dataset,
batch_size=args.batch_size,
logger=logger, test=False)
logger.info("Train {} with {} samples".format(args.dataset, len(data_loader.dataset)))
# PREPARES MODEL
args.img_size = get_img_size(args.dataset) # stores for metadata
model = load_model(exp_dir, filename='model.pt')
logger.info('Num parameters in model: {}'.format(get_n_param(model)))
# Extract Features
model = model.to(device) # make sure trainer and viz on same device
fe = FeatureExtractor(model,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar)
fe(data_loader,
epochs=args.epochs,
checkpoint_every=args.checkpoint_every, feature_dir=feature_dir)
# SAVE MODEL AND EXPERIMENT INFORMATION
# save_model(trainer.model, exp_dir, metadata=vars(args))
print('Done.')
def main(args):
set_seed(args.seed)
device = torch.device(
'cuda:{}'.format(args.gpu) if torch.cuda.is_available() else 'cpu')
exp_dir = os.path.join(RES_DIR, args.name)
print("save and load experiments at : {}".format(exp_dir))
if not args.is_eval_only: #train
create_directory(exp_dir)
# PREPARES TRAINING DATA
train_loader = get_dataloaders(args.dataset,
batch_size=args.batch_size)
##############
# PREPARES MODEL
args.img_size = get_img_size(args.dataset) # stores for metadata
cs = [1, 64, 128, 1024]
model = VLAE(args, args.latent_dim, cs)
#TRAINS
optimizer = optim.Adam(model.parameters(), lr=args.lr)
model = model.to(device) # make sure trainer and viz on same device
gif_visualizer = GifTraversalsTraining(model, args.dataset, exp_dir)
reg_coeff = [args.reg_coeff0, args.reg_coeff1, args.reg_coeff2]
trainer = Trainer(model,
optimizer,
reg_coeff,
device=device,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar,
gif_visualizer=gif_visualizer)
trainer(args,
train_loader,
epochs=args.epochs,
checkpoint_every=args.checkpoint_every)
#SAVE MODEL AND EXPERIMENT INFORMATION
save_model(trainer.model, exp_dir, metadata=vars(args))
print("Model has been saved")
def get_samples(dataset, num_samples, idcs=[]):
""" Generate a number of samples from the dataset.
Parameters
----------
dataset : str
The name of the dataset.
num_samples : int, optional
The number of samples to load from the dataset
idcs : list of ints, optional
List of indices to of images to put at the begning of the samples.
"""
data_loader = get_dataloaders(dataset, batch_size=1, shuffle=idcs is None)
idcs += random.sample(range(len(data_loader.dataset)),
num_samples - len(idcs))
samples = torch.stack([data_loader.dataset[i][0] for i in idcs], dim=0)
print("Selected idcs: {}".format(idcs))
return samples
def main(args):
"""Main train and evaluation function.
Parameters
----------
args: argparse.Namespace
Arguments
"""
formatter = logging.Formatter(
'%(asctime)s %(levelname)s - %(funcName)s: %(message)s', "%H:%M:%S")
logger = logging.getLogger(__name__)
logger.setLevel(args.log_level.upper())
stream = logging.StreamHandler()
stream.setLevel(args.log_level.upper())
stream.setFormatter(formatter)
logger.addHandler(stream)
set_seed(args.seed)
device = get_device(is_gpu=not args.no_cuda)
exp_dir = os.path.join(RES_DIR, args.name)
logger.info("Root directory for saving and loading experiments: {}".format(
exp_dir))
if not args.is_eval_only:
create_safe_directory(exp_dir, logger=logger)
if args.loss == "factor":
logger.info(
"FactorVae needs 2 batches per iteration. To replicate this behavior while being consistent, we double the batch size and the the number of epochs."
)
args.batch_size *= 2
args.epochs *= 2
# PREPARES DATA
train_loader = get_dataloaders(args.dataset,
batch_size=args.batch_size,
logger=logger)
logger.info("Train {} with {} samples".format(
args.dataset, len(train_loader.dataset)))
# PREPARES MODEL
args.img_size = get_img_size(args.dataset) # stores for metadata
model = init_specific_model(args.model_type, args.img_size,
args.latent_dim)
logger.info('Num parameters in model: {}'.format(get_n_param(model)))
# TRAINS
optimizer = optim.Adam(model.parameters(), lr=args.lr)
model = model.to(device) # make sure trainer and viz on same device
gif_visualizer = GifTraversalsTraining(model, args.dataset, exp_dir)
loss_f = get_loss_f(args.loss,
n_data=len(train_loader.dataset),
device=device,
**vars(args))
trainer = Trainer(model,
optimizer,
loss_f,
device=device,
logger=logger,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar,
gif_visualizer=gif_visualizer)
trainer(
train_loader,
epochs=args.epochs,
checkpoint_every=args.checkpoint_every,
)
# SAVE MODEL AND EXPERIMENT INFORMATION
save_model(trainer.model, exp_dir, metadata=vars(args))
if args.is_metrics or not args.no_test:
model = load_model(exp_dir, is_gpu=not args.no_cuda)
metadata = load_metadata(exp_dir)
# TO-DO: currently uses train datatset
test_loader = get_dataloaders(metadata["dataset"],
batch_size=args.eval_batchsize,
shuffle=False,
logger=logger)
loss_f = get_loss_f(args.loss,
n_data=len(test_loader.dataset),
device=device,
**vars(args))
use_wandb = False
if use_wandb:
loss = args.loss
wandb.init(project="atmlbetavae", config={"VAE_loss": args.loss})
if loss == "betaH":
beta = loss_f.beta
wandb.config["Beta"] = beta
evaluator = Evaluator(model,
loss_f,
device=device,
logger=logger,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar,
use_wandb=use_wandb)
evaluator(test_loader,
is_metrics=args.is_metrics,
is_losses=not args.no_test)
m1, s1 = _calculate_activation_statistics(dataloader_original, length, model, batch_size, dims)
print("Calculated m1 and s1")
m2, s2 = _calculate_activation_statistics(dataloader_reconstructed, length, model, batch_size, dims)
print("Calculated m2 and s2")
fid_value = _calculate_frechet_distance(m1, s1, m2, s2)
return fid_value
if __name__ == "__main__":
import torch
import random
from disvae.utils.modelIO import load_model
import argparse
import logging
import sys
MODEL_PATH = sys.argv[2] # get the model path (e.g. "results/betaH_mnist")
MODEL_NAME = "model.pt"
GPU_AVAILABLE = True
vae_model = load_model(directory=MODEL_PATH, is_gpu=GPU_AVAILABLE, filename=MODEL_NAME)
mode = sys.argv[1] # get the name of the dataset you want to measure FID for
dataloader1 = get_dataloaders(mode, batch_size=128)[0]
fid_value = get_fid_value(dataloader1, vae_model)
print("FID for ", mode, ": ", fid_value)
def main(args: argparse.Namespace):
"""Main train and evaluation function."""
formatter = logging.Formatter(
'%(asctime)s %(levelname)s - %(funcName)s: %(message)s', "%H:%M:%S")
logger = logging.getLogger(__name__)
logger.setLevel("INFO")
stream = logging.StreamHandler()
stream.setLevel("INFO")
stream.setFormatter(formatter)
logger.addHandler(stream)
set_seed(args.seed)
device = get_device(is_gpu=not args.no_cuda)
exp_dir = os.path.join(RES_DIR, args.name)
logger.info(
f"Root directory for saving and loading experiments: {exp_dir}")
if not args.is_eval_only:
create_safe_directory(exp_dir, logger=logger)
if args.loss == "factor":
logger.info(
"FactorVae needs 2 batches per iteration." +
"To replicate this behavior, double batch size and epochs.")
args.batch_size *= 2
args.epochs *= 2
# PREPARES DATA
train_loader = get_dataloaders(args.dataset,
noise=args.noise,
batch_size=args.batch_size,
logger=logger)
logger.info(
f"Train {args.dataset} with {len(train_loader.dataset)} samples")
# PREPARES MODEL
args.img_size = get_img_size(args.dataset) # stores for metadata
model = VAE(args.img_size, args.latent_dim)
logger.info(f'Num parameters in model: {get_n_param(model)}')
# TRAINS
optimizer = optim.Adam(model.parameters(), lr=args.lr)
model = model.to(device)
gif_visualizer = GifTraversalsTraining(model, args.dataset, exp_dir)
loss_f = get_loss_f(args.loss,
n_data=len(train_loader.dataset),
device=device,
**vars(args))
if args.loss in ['tdGJS', 'tGJS']:
loss_optimizer = optim.Adam(loss_f.parameters(), lr=args.lr)
else:
loss_optimizer = None
print(loss_optimizer)
trainer = Trainer(model,
optimizer,
loss_f,
device=device,
logger=logger,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar,
gif_visualizer=gif_visualizer,
loss_optimizer=loss_optimizer,
denoise=args.noise is not None)
trainer(
train_loader,
epochs=args.epochs,
checkpoint_every=args.checkpoint_every,
)
# SAVE MODEL AND EXPERIMENT INFORMATION
save_model(trainer.model, exp_dir, metadata=vars(args))
# Eval
model = load_model(exp_dir, is_gpu=not args.no_cuda)
metadata = load_metadata(exp_dir)
test_loader = get_dataloaders(metadata["dataset"],
noise=args.noise,
train=False,
batch_size=128,
logger=logger)
loss_f = get_loss_f(args.loss,
n_data=len(test_loader.dataset),
device=device,
**vars(args))
evaluator = Evaluator(model,
loss_f,
device=device,
is_metrics=args.is_metrics,
is_train=False,
logger=logger,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar,
denoise=args.noise is not None)
evaluator(test_loader)
# Train set also
test_loader = get_dataloaders(metadata["dataset"],
train=True,
batch_size=128,
logger=logger)
loss_f = get_loss_f(args.loss,
n_data=len(test_loader.dataset),
device=device,
**vars(args))
evaluator = Evaluator(model,
loss_f,
device=device,
is_metrics=args.is_metrics,
is_train=True,
logger=logger,
save_dir=exp_dir,
is_progress_bar=not args.no_progress_bar)
evaluator(test_loader)
def main(args):
"""Main train and evaluation function.
Parameters
----------
args: argparse.Namespace
Arguments
"""
# Logging info
formatter = logging.Formatter('%(asctime)s %(levelname)s - '
'%(funcName)s: %(message)s',
'%H:%M:%S')
logger = logging.getLogger(__name__)
logger.setLevel('INFO')
stream = logging.StreamHandler()
stream.setLevel('INFO')
stream.setFormatter(formatter)
logger.addHandler(stream)
set_seed(args.seed)
device = torch.device(
'cuda' if torch.cuda.is_available() and args.cuda else 'cpu')
model_name = f'{args.name}_lr{args.lr}_z{args.latent_dim}' \
+ f'_h{args.hidden_dim}_p{args.p_dropout}'
model_dir = os.path.join(args.results, model_name)
logger.info(f'Directory for saving and loading models: {model_dir}')
if not args.eval:
# Model directory
new_model_dir(model_dir, logger=logger)
# Dataloaders
train_loader, valid_loader = get_dataloaders(
args.data, args.t_hours, args.n_bins,
validation=True, dynamic=args.dynamic,
batch_size=args.bs, logger=logger)
logger.info(
f'Train {args.model_type}-{args.t_hours} ' +
f'with {len(train_loader.dataset)} samples')
# Load model
n_tokens = len(np.load(
os.path.join(
args.data, '_dicts', f'{args.t_hours}_{args.n_bins}.npy'),
allow_pickle=True).item())
model = init_model(
args.model_type, n_tokens, args.latent_dim, args.hidden_dim,
p_dropout=args.p_dropout, dt=args.dt,
weighted=args.weighted, dynamic=args.dynamic)
logger.info(f'#params in model: {get_n_param(model)}')
# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
loss_f = BCE()
model = model.to(device)
# Training
trainer = Trainer(
model, loss_f, optimizer,
device=device, logger=logger, save_dir=model_dir, p_bar=args.p_bar)
trainer.train(
train_loader, valid_loader,
epochs=args.epochs, early_stopping=args.early_stopping)
# Save model
metadata = vars(args)
metadata['n_tokens'] = n_tokens
save_model(trainer.model, model_dir, metadata=metadata)
if args.test:
# Load model
model = load_model(model_dir, is_gpu=args.cuda)
metadata = load_metadata(model_dir)
# Dataloader
test_loader, _ = get_dataloaders(
metadata['data'], metadata['t_hours'], metadata['n_bins'],
validation=False, dynamic=metadata['dynamic'], batch_size=128,
shuffle=False, logger=logger)
# Evaluate
loss_f = BCE()
evaluator = Trainer(
model, loss_f,
device=device, logger=logger, save_dir=model_dir, p_bar=args.p_bar)
evaluator._valid_epoch(test_loader)