def main_flags():
# Data & model config
flags.DEFINE_string('data_config', 'datasets/multid_config.py',
'Path to a data config file.')
flags.DEFINE_string('model_config', 'models/genesis_config.py',
'Path to a model config file.')
# Logging config
flags.DEFINE_string('results_dir', 'checkpoints',
'Top directory for all experimental results.')
flags.DEFINE_string('run_name', 'test',
'Name of this job and name of results folder.')
flags.DEFINE_integer(
'report_loss_every', 1000,
'Number of iterations between reporting minibatch loss.')
flags.DEFINE_integer('run_validation_every', 10000,
'How many equally spaced validation runs to do.')
flags.DEFINE_integer('num_checkpoints', 40,
'How many equally spaced model checkpoints to save.')
flags.DEFINE_boolean('resume', False, 'Tries to resume a job if True.')
flags.DEFINE_boolean(
'log_grads_and_weights', False,
'Log gradient and weight histograms - storage intensive!')
flags.DEFINE_boolean(
'log_distributions', False,
'Log mu and sigma of posterior and prior distributions.')
# Optimisation config
flags.DEFINE_integer('train_iter', 2000000,
'Number of training iterations.')
flags.DEFINE_integer('batch_size', 32, 'Mini-batch size.')
flags.DEFINE_string('optimiser', 'adam', 'Optimiser for updating weights.')
flags.DEFINE_float('learning_rate', 0.0001, 'Learning rate.')
flags.DEFINE_integer('N_eval', 10000,
'Number of samples to run evaluation on.')
# Loss config
flags.DEFINE_float('beta', 0.5, 'KL weighting.')
flags.DEFINE_boolean('beta_warmup', True, 'Warm up beta.')
flags.DEFINE_boolean('geco', True, 'Use GECO objective.')
flags.DEFINE_float('g_goal', 0.5655, 'GECO recon goal.')
flags.DEFINE_float('g_lr', 1e-5, 'GECO learning rate.')
flags.DEFINE_float('g_alpha', 0.99, 'GECO momentum for error.')
flags.DEFINE_float('g_init', 1.0, 'GECO inital Lagrange factor.')
flags.DEFINE_float('g_min', 1e-10, 'GECO min Lagrange factor.')
flags.DEFINE_float('g_speedup', 10., 'Scale GECO lr if delta positive.')
# Other
flags.DEFINE_boolean('gpu', True, 'Use GPU if available.')
flags.DEFINE_boolean('multi_gpu', False, 'Use multiple GPUs if available.')
flags.DEFINE_boolean('debug', False, 'Debug flag.')
flags.DEFINE_integer('seed', 0, 'Seed for random number generators.')
flags.DEFINE_string('model_config', 'configs/mnist_mlp.py',
'Path to a model config file.')
flags.DEFINE_string('results_dir', 'checkpoints',
'Top directory for all experimental results.')
flags.DEFINE_string('run_name', 'mnist',
'Name of this job and name of results folder.')
flags.DEFINE_boolean('resume', False, 'Tries to resume a job if True.')
# Logging config
flags.DEFINE_integer('report_loss_every', 100,
'Number of iterations between reporting minibatch loss.')
flags.DEFINE_integer('train_epochs', 20, 'Maximum number of training epochs.')
# Experiment config
flags.DEFINE_integer('batch_size', 32, 'Mini-batch size.')
flags.DEFINE_float('learning_rate', 1e-5, 'SGD learning rate.')
# Parse flags
config = forge.config()
# Prepare enviornment
logdir = osp.join(config.results_dir, config.run_name)
logdir, resume_checkpoint = fet.init_checkpoint(logdir, config.data_config,
config.model_config,
config.resume)
checkpoint_name = osp.join(logdir, 'model.ckpt')
# Load data
train_loader = fet.load(config.data_config, config)
# Load model
model = fet.load(config.model_config, config)
import torch.nn as nn
from torch.distributions.normal import Normal
from forge import flags
from modules.blocks import Flatten
from modules.decoders import BroadcastDecoder
from third_party.sylvester.VAE import VAE
# GatedConvVAE
flags.DEFINE_integer('latent_dimension', 64, 'Latent channels.')
flags.DEFINE_boolean('broadcast_decoder', False,
'Use broadcast decoder instead of deconv.')
# Losses
flags.DEFINE_boolean('pixel_bound', True, 'Bound pixel values to [0, 1].')
flags.DEFINE_float('pixel_std', 0.7, 'StdDev of reconstructed pixels.')
def load(cfg):
return BaselineVAE(cfg)
class BaselineVAE(nn.Module):
def __init__(self, cfg):
super(BaselineVAE, self).__init__()
cfg.K_steps = None
# Configuration
self.ldim = cfg.latent_dimension
self.pixel_std = cfg.pixel_std
self.pixel_bound = cfg.pixel_bound
self.debug = cfg.debug
# job config
flags.DEFINE_string('data_config', 'configs/mnist_data.py', 'Path to a data config file.')
flags.DEFINE_string('model_config', 'configs/mnist_mlp.py', 'Path to a model config file.')
flags.DEFINE_string('results_dir', '../checkpoints', 'Top directory for all experimental results.')
flags.DEFINE_string('run_name', 'test_run', 'Name of this job. Results will be stored in a corresponding folder.')
flags.DEFINE_boolean('resume', False, 'Tries to resume a job if True.')
# logging config
flags.DEFINE_integer('report_loss_every', int(1e3), 'Number of iterations between reporting minibatch loss - hearbeat.')
flags.DEFINE_integer('save_itr', int(1e4), 'Number of iterations between snapshotting the model.')
flags.DEFINE_integer('train_itr', int(2e6), 'Maximum number of training iterations.')
# experiment config
flags.DEFINE_integer('batch_size', 32, '')
flags.DEFINE_float('learning_rate', 1e-5, 'Initial values of the learning rate')
# gpu
flags.DEFINE_string('gpu', '0', 'Id of the gpu to use for this job.')
# Parse flags
config = forge.config()
# sets visible gpus to config.gpu
fet.set_gpu(config.gpu)
# Prepare enviornment
logdir = osp.join(config.results_dir, config.run_name)
logdir, resume_checkpoint = fet.init_checkpoint(logdir, config.data_config, config.model_config, config.resume)
checkpoint_name = osp.join(logdir, 'model.ckpt')
flags.DEFINE_integer(
"nbhd_size", 25, "The number of samples to use for Monte Carlo estimation")
flags.DEFINE_string("activation_function", "swish",
"Activation function to use in the network")
flags.DEFINE_boolean("batch_norm", True, "Use batch norm in the layers")
flags.DEFINE_bool(
"mean_pooling",
True,
"Use mean pooling insteave of sum pooling in the invariant layer",
)
flags.DEFINE_integer("num_layers", 6, "Number of ResNet layers to use")
flags.DEFINE_string("group", "SE3", "Group to be invariant to")
flags.DEFINE_integer("channels", 1536, "Number of channels in the conv layers")
flags.DEFINE_float(
"fill",
1.0,
"specifies the fraction of the input which is included in local neighborhood. (can be array to specify a different value for each layer",
)
flags.DEFINE_integer(
"lift_samples", 4,
"Number of coset lift samples to use for non-trivial stabilisers")
flags.DEFINE_integer("model_seed", 0, "Model rng seed")
flags.DEFINE_string(
"lie_algebra_nonlinearity",
None,
"Nonlinearity to apply to the norm of the lie algebra elements. Supported are None/tanh",
)
def load(config, **unused_kwargs):
# Logging
flags.DEFINE_integer(
"report_loss_every", 500, "Number of iterations between reporting minibatch loss."
)
flags.DEFINE_integer(
"evaluate_every", 10000, "Number of iterations between reporting validation loss."
)
flags.DEFINE_integer(
"save_check_points",
10,
"frequency with which to save checkpoints, in number of epochs.",
)
flags.DEFINE_boolean("log_train_values", True, "Logs train values if True.")
flags.DEFINE_float(
"ema_alpha", 0.99, "Alpha coefficient for exponential moving average of train logs."
)
# Optimization
flags.DEFINE_integer("train_epochs", 500, "Maximum number of training epochs.")
flags.DEFINE_integer("batch_size", 90, "Mini-batch size.")
flags.DEFINE_float("learning_rate", 1e-5, "SGD learning rate.")
flags.DEFINE_float("beta1", 0.5, "Adam Beta 1 parameter")
flags.DEFINE_float("beta2", 0.9, "Adam Beta 2 parameter")
flags.DEFINE_string(
"lr_schedule",
"none",
"What learning rate schedule to use. Options: cosine, none",
)
flags.DEFINE_boolean(
"parameter_count", False, "If True, print model parameter count and exit"
flags.DEFINE_integer("report_loss_every", 10,
"Number of iterations between reporting minibatch loss.")
flags.DEFINE_integer(
"evaluate_every", 10000,
"Number of iterations between reporting validation loss.")
flags.DEFINE_integer(
"save_check_points",
50,
"frequency with which to save checkpoints, in number of epoches.",
)
flags.DEFINE_boolean("log_train_values", True, "Logs train values if True.")
# Optimization
flags.DEFINE_integer("train_epochs", 200, "Maximum number of training epochs.")
flags.DEFINE_integer("batch_size", 100, "Mini-batch size.")
flags.DEFINE_float("learning_rate", 1e-3, "Adam learning rate.")
flags.DEFINE_float("beta1", 0.9, "Adam Beta 1 parameter")
flags.DEFINE_float("beta2", 0.999, "Adam Beta 2 parameter")
flags.DEFINE_string("lr_schedule", "cosine_annealing",
"Learning rate schedule.")
flags.DEFINE_boolean("clip_grad_norm", False,
"Clip norm of the gradient at max_grad_norm.")
flags.DEFINE_integer("max_grad_norm", 100,
"Maximum norm of gradient when clip_grad_norm is True.")
# GPU device
flags.DEFINE_integer("device", 0, "GPU to use.")
# Debug mode tracks more stuff
flags.DEFINE_boolean("debug", False,
"Track and show on tensorboard more metrics.")
import os
from torch.utils.data import DataLoader
from oil.utils.utils import FixedNumpySeed
from oil.datasetup.datasets import split_dataset
from lie_conv.datasets import QM9datasets
from corm_data.collate import collate_fn
import forge
from forge import flags
flags.DEFINE_float(
"subsample_trainset",
1.0,
"Proportion or number of samples of the full trainset to use",
)
flags.DEFINE_string(
"task",
"h**o",
"Which task in the QM9 dataset to train on. Pass as a comma separated string",
)
flags.DEFINE_boolean("recenter", False,
"Recenter the positions of atoms with charge > 0")
flags.DEFINE_integer("batch_fit", 0, "number of samples to fit to")
flags.DEFINE_integer("data_seed", 0, "seed to pick data with")
def load(config, **unused_kwargs):
"mlp",
"Selects the type of attention kernel to use. mlp/relative_position/dot_product are valid",
)
flags.DEFINE_integer("kernel_dim", 16,
"Hidden layer size to use in kernel MLPs")
flags.DEFINE_integer("num_layers", 6, "Number of ResNet layers to use")
flags.DEFINE_string("group", "SE3", "Group to be invariant to")
flags.DEFINE_integer(
"lift_samples", 1,
"Number of coset lift samples to use for non-trivial stabilisers")
flags.DEFINE_integer(
"mc_samples",
0,
"Number of samples to use for estimating attention. 0 sets to use all points",
)
flags.DEFINE_float("fill", 1.0,
"Select mc_samples from K nearest mc_samples/fill points")
flags.DEFINE_integer("model_seed", 0, "Model rng seed")
flags.DEFINE_string("architecture", "model_1",
"The model architecture to use. model_1/lieconv")
flags.DEFINE_string("attention_fn", "softmax",
"Type of attention function to use. softmax/dot_product")
flags.DEFINE_integer(
"feature_embed_dim",
None,
"Dimensionality of the embedding of the features for each head. Only used by some kernels",
)
flags.DEFINE_float(
"max_sample_norm",
None,
"Maximum sample norm to allow through the lifting stage to prevent numerical issues.",
)
import functools
import numpy as np
import tensorflow as tf
import torch
from torchvision import transforms
# import ipdb
from forge import flags
import pickle
import os
import json
flags.DEFINE_integer("train_size", 10000, "Number of training examples per epoch.")
flags.DEFINE_integer("test_size", 1000, "Number of testing examples per epoch.")
flags.DEFINE_integer("naug", 2, "Number of augmentation.")
flags.DEFINE_float("corner_noise", 0.1, "See `create_constellations`.")
flags.DEFINE_boolean("shuffle_corners", True, "See `create_constellations`.")
flags.DEFINE_float("pattern_upscale", 0.0, "See `create_constellations`.")
flags.DEFINE_float("max_rotation", 0.33, "See `create_constellations`.")
flags.DEFINE_float("global_rotation_angle", 0.0, "See `create_constellations`.")
flags.DEFINE_float("global_translation", 0.0, "See `create_constellations`.")
flags.DEFINE_float("pattern_drop_prob", 0.5, "See `create_constellations`.")
flags.DEFINE_integer("patterns_reps", 2, "See `create_constellations`.")
flags.DEFINE_integer("data_seed", 0, "Seed for data generation.")
def roots_of_unity(n):
x_coors = np.cos(2 * np.pi / n * np.arange(n)[..., np.newaxis])
y_coors = np.sin(2 * np.pi / n * np.arange(n)[..., np.newaxis])
flags.DEFINE_boolean("resume", False, "Tries to resume a job if True.")
# Logging
flags.DEFINE_integer("report_loss_every", 500,
"Number of iterations between reporting minibatch loss.")
flags.DEFINE_integer(
"evaluate_every", 10000,
"Number of iterations between reporting validation loss.")
flags.DEFINE_integer(
"save_check_points",
50,
"frequency with which to save checkpoints, in number of epoches.",
)
flags.DEFINE_boolean("log_train_values", True, "Logs train values if True.")
flags.DEFINE_float(
"ema_alpha", 0.99,
"Alpha coefficient for exponential moving average of train logs.")
flags.DEFINE_boolean("train_aug_t2", False, "T2 Training augmentation.")
flags.DEFINE_boolean("train_aug_se2", True, "SE2 Training augmentation.")
# Optimization
flags.DEFINE_integer("train_epochs", 200, "Maximum number of training epochs.")
flags.DEFINE_integer("batch_size", 90, "Mini-batch size.")
flags.DEFINE_float("learning_rate", 1e-5, "SGD learning rate.")
flags.DEFINE_float("beta1", 0.5, "Adam Beta 1 parameter")
flags.DEFINE_float("beta2", 0.9, "Adam Beta 2 parameter")
# GPU device
flags.DEFINE_integer("device", 0, "GPU to use.")
