def load_model(directory, is_gpu=True, filename='model.pt'):
"""Load a trained model.
Parameters
----------
directory : string
Path to folder where model is saved.
is_gpu : bool
Whether to load on GPU is available.
"""
device = get_device(is_gpu=is_gpu)
metadata = load_metadata(directory)
model = init_model(metadata['model_type'],
metadata['n_tokens'],
metadata['latent_dim'],
metadata['hidden_dim'],
dt=metadata['dt'],
weighted=metadata['weighted'],
dynamic=metadata['dynamic']).to(device)
model.load_state_dict(torch.load(os.path.join(directory, filename)),
strict=False)
model.eval()
return model
def __init__(self, input_dim, hidden_dim, layer_norm=True):
"""
LSTM class with layer normalization option.
Parameters
----------
input_dim : int
Dimensionality of LSTM input.
hidden_dim : int
Dimensionality of LSTM hidden state.
layer_norm : bool, optional
Whether to use layer normalized version of LSTM.
"""
super(LSTM, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.layer_norm = layer_norm
self.lstm_cell = LSTMCell(input_dim,
hidden_dim,
layer_norm=self.layer_norm)
self.device = get_device()
def __init__(self, n_tokens, latent_dim, dt=1.0, weighted=True):
"""Token embedder.
Parameters
----------
n_tokens: int
Number of tokens in vocabulary
latent_dim: int
Dimensionality of latent embedding.
dt: float
Time increment between sequence steps.
weighted: bool
Whether or not to add embedding weights.
"""
super(Embedder, self).__init__()
self.n_tokens = n_tokens
self.latent_dim = latent_dim
self.dt = dt
self.weighted = weighted
self.device = get_device()
self.embedX = Embedding(n_tokens + 1, latent_dim, padding_idx=0)
if self.weighted:
self.embedW = Embedding(n_tokens + 1, 1)
self.softmax = nn.Softmax(dim=1)
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):
"""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))
# Initialize a new sweep
# Arguments:
# – sweep_config: the sweep config dictionary defined above
# – entity: Set the username for the sweep
# – project: Set the project name for the sweep
config = wandb.config
with wandb.init(name="sweep-reg_anneal-seed",
notes='This is a test run',
tags=['btcvae', 'dsprites'],
entity='neonkitchen',
project]'sweep'
config = config):
def main():
# Initialise
config = parse_args(sys.argv[1:])
device = get_device()
set_seed(config.seed)
# Data
arrs = np.load(config.data, allow_pickle=True).item()
train_data = EHR(arrs['X_train'], arrs['Y_train'])
valid_data = EHR(arrs['X_valid'], arrs['Y_valid'])
train_loader = DataLoader(train_data,
batch_size=config.batch_size,
shuffle=True,
pin_memory=True)
valid_loader = DataLoader(valid_data,
batch_size=128,
shuffle=False,
pin_memory=True)
# Model
n_tokens = int(arrs['X_train'].max()) + 1 # +1 to embedding size for 0 emb
model = Net(n_tokens,
config.emb_dim,
config.rnn_dim,
variational=config.variational,
layer_norm=config.layer_norm).to(device)
optimizer = optim.Adam(model.parameters(),
lr=config.learning_rate,
weight_decay=config.weight_decay)
# Train
for epoch in range(config.epochs):
t_loss = train_epoch(model, device, train_loader, optimizer)
v_loss, v_auroc, v_auprc = test(model, device, valid_loader)
# Save
new_model_dir(os.path.join('results', config.name))
torch.save(model.state_dict(), os.path.join(RES_DIR, "model.h5"))
def __init__(self, n_tokens, emb_dim, T=48.0):
"""
Embedding aggregation class based in equal time intervals
Parameters
----------
n_tokens : int
Number of tokens to embed
emb_dim : int
Dimensionality embedding space
T : float, optional
Total period of time
"""
super(Aggregator, self).__init__()
self.n_tokens = n_tokens
self.emb_dim = emb_dim
self.T = T
self.device = get_device()
self.embed = Embedding(n_tokens, emb_dim, variational=False)
def __init__(self, n_tokens, emb_dim, T=48):
"""
Embedding class with built-in variational option.
Parameters
----------
n_tokens : int
Number of tokens to embed
emb_dim : int
Dimensionality embedding space
T : float, optional
Total period of time
"""
super(VariationalAggregator, self).__init__()
self.n_tokens = n_tokens
self.emb_dim = emb_dim
self.T = T
self.device = get_device()
self.embed = Embedding(n_tokens, emb_dim, variational=True)
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)
def _reconstruction_loss(data,
recon_data,
distribution="bernoulli",
storer=None):
"""
Calculates the per image reconstruction loss for a batch of data. I.e. negative
log likelihood.
Parameters
----------
data : torch.Tensor
Input data (e.g. batch of images). Shape : (batch_size, n_chan,
height, width).
recon_data : torch.Tensor
Reconstructed data. Shape : (batch_size, n_chan, height, width).
distribution : {"bernoulli", "gaussian", "laplace"}
Distribution of the likelihood on the each pixel. Implicitely defines the
loss Bernoulli corresponds to a binary cross entropy (bse) loss and is the
most commonly used. It has the issue that it doesn't penalize the same
way (0.1,0.2) and (0.4,0.5), which might not be optimal. Gaussian
distribution corresponds to MSE, and is sometimes used, but hard to train
ecause it ends up focusing only a few pixels that are very wrong. Laplace
distribution corresponds to L1 solves partially the issue of MSE.
storer : dict
Dictionary in which to store important variables for vizualisation.
Returns
-------
loss : torch.Tensor
Per image cross entropy (i.e. normalized per batch but not pixel and
channel)
"""
batch_size, n_chan, height, width = recon_data.size()
is_colored = n_chan == 3
if distribution == "bernoulli":
loss = F.binary_cross_entropy(recon_data, data, reduction="sum")
elif distribution == "gaussian":
# loss in [0,255] space but normalized by 255 to not be too big
loss = F.mse_loss(recon_data * 255, data * 255, reduction="sum") / 255
elif distribution == "laplace":
# loss in [0,255] space but normalized by 255 to not be too big but
# multiply by 255 and divide 255, is the same as not doing anything for L1
loss = F.l1_loss(recon_data, data, reduction="sum")
loss = loss * 3 # emperical value to give similar values than bernoulli => use same hyperparam
loss = loss * (loss != 0) # masking to avoid nan
elif distribution == "dfc_vae":
loss = 0
device = get_device(is_gpu=data.is_cuda)
descriptor = VGG19().to(device)
data_features = [_data.to(device) for _data in descriptor(data)]
recon_features = [_data.to(device) for _data in descriptor(recon_data)]
for f, target in zip(recon_features, data_features):
loss += F.mse_loss(f * 255, target * 255, reduction="sum") / 255
else:
assert distribution not in RECON_DIST
raise ValueError("Unkown distribution: {}".format(distribution))
loss = loss / batch_size
if storer is not None:
storer['recon_loss'].append(loss.item())
return loss
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)
import gym
import numpy as np
import torch
from torch.nn import functional as F
from models.decoder import StateTransitionDecoder, RewardDecoder, TaskDecoder
from models.encoder import RNNEncoder
from utils.storage_vae import RolloutStorageVAE
from utils import helpers as utl
device = utl.get_device()
class VaribadVAE:
"""
VAE of variBAD:
- has an encoder and decoder,
- can compute the ELBO loss
- can update the VAE part of the model
"""
def __init__(self, args, logger, get_iter_idx):
self.args = args
self.logger = logger
self.get_iter_idx = get_iter_idx
self.task_dim = self.get_task_dim()
# initialise the encoder
self.encoder = self.initialise_encoder()
# initialise the decoders (returns None for unused decoders)
import logging
import torch
import os
from dataloader.utils import load_dataset
from decoding.generation import generate_beam
from decoding.greedy import greedy_search
from architecture.model import CDS
from utils.helpers import get_device, calculate_rouge, set_seed
from utils.cmdopt import parse_sum_args
logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)
opt = parse_sum_args()
device = get_device()
set_seed(55)
def test(dataset, fields, model):
already, hypothesis, references = 0, [], []
for batch in dataset:
if opt.tf:
predictions = greedy_search(model, opt, batch.src, fields,
opt.max_length)
else:
predictions, _ = generate_beam(5, model, opt, batch.src, fields)
predictions = [p for p, _ in predictions]
