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]