def load_ae(path, target_dataset, batch, all_aes, return_dataset=False):
r_param = re.compile(
'(?P<name>[a-zA-Z][a-z_]*)(?P<value>(T)|(F)|(\d+(\.\d+)?(,\d+)*))')
folders = [x for x in os.path.abspath(path).split('/') if x]
dataset = folders[-2]
if dataset != target_dataset:
tf.logging.log(
tf.logging.WARN,
'Mismatched datasets between classfier and AE (%s, %s)',
target_dataset, dataset)
print(folders)
class_name, argpairs = folders[-1].split('_', 1)
params = {}
for x in r_param.findall(argpairs):
name, value = x[:2]
if ',' in value:
pass
elif value in ('T', 'F'):
__dict__ = dict(T=True, F=False)
value = __dict__[value]
elif '.' in value:
value = float(value)
else:
value = int(value)
params[name] = value
class_ = all_aes[class_name]
dataset = data.get_dataset(dataset, dict(batch_size=batch))
ae = class_(dataset, '/' + os.path.join(*(folders[:-2])), **params)
if return_dataset:
return ae, dataset
else:
return ae, folders[-1]
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
model = XFullyConnected(dataset,
FLAGS.train_dir,
smoothing=FLAGS.smoothing)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
model = XFullyConnected(
dataset,
FLAGS.train_dir,
smoothing=FLAGS.smoothing)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AEBaseline(dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
ae, ae_model = utils.load_ae(FLAGS.ae_dir, FLAGS.dataset, FLAGS.batch,
all_aes.ALL_AES)
with utils.HookReport.disable():
ae.eval_mode()
model = XLatentAE(dataset, FLAGS.train_dir)
model.train_dir = os.path.join(model.train_dir, ae_model)
model.ae = ae
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AEBaseline(
dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AAE(dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
adversary_lr=FLAGS.adversary_lr,
disc_layer_sizes=FLAGS.disc_layer_sizes)
model.train()
def main(argv):
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = VAE(
dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
beta=FLAGS.beta)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AAE(
dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
adversary_lr=FLAGS.adversary_lr,
disc_layer_sizes=FLAGS.disc_layer_sizes)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AEVQVAE(dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
z_log_size=FLAGS.z_log_size,
beta=FLAGS.beta,
num_latents=FLAGS.num_latents)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = ACAI(dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
advweight=FLAGS.advweight,
advdepth=FLAGS.advdepth or FLAGS.depth,
reg=FLAGS.reg)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
ae, ae_model = utils.load_ae(FLAGS.ae_dir, FLAGS.dataset, FLAGS.batch,
all_aes.ALL_AES)
with utils.HookReport.disable():
ae.eval_mode()
model = XLatentAE(
dataset,
FLAGS.train_dir)
model.train_dir = os.path.join(model.train_dir, ae_model)
model.ae = ae
model.train()
def main(argv):
print(FLAGS.flag_values_dict())
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
# latent: [?, 4, 4, 16]
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AEBaseline(
dataset,
FLAGS.train_dir,
latent=FLAGS.latent, # channels of latent
depth=FLAGS.depth, # channels of first convolution
scales=scales)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = AEVQVAE(
dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
z_log_size=FLAGS.z_log_size,
beta=FLAGS.beta,
num_latents=FLAGS.num_latents)
model.train()
def load_ae(path, target_dataset, batch, all_aes, return_dataset=False):
"""
:param path: dataset path
:param target_dataset: mnist32/lines32/...
:param batch: batchsz
:param all_aes: all autoencoders dict
:param return_dataset: return with dataset or not
:return:
"""
r_param = re.compile(
'(?P<name>[a-zA-Z][a-z_]*)(?P<value>(True)|(False)|(\d+(\.\d+)?(,\d+)*))'
)
folders = [
x for x in os.path.abspath(path).split('/') if x
] # [..., 'logs', 'mnist32', 'AEBaseline_depth16_latent2_scales3']
dataset = folders[-2]
if dataset != target_dataset:
tf.logging.log(
tf.logging.WARN,
'Mismatched datasets between classfier and AE (%s, %s)',
target_dataset, dataset)
class_name, argpairs = folders[-1].split(
'_', 1) # 'AEBaseline', depth16_latent2_scales3
params = {} # {depth:16, scales:3, latent:2}
for x in r_param.findall(argpairs):
name, value = x[:2]
if ',' in value:
pass
elif value in ('True', 'False'):
value = value == True
elif '.' in value:
value = float(value)
else:
value = int(value)
params[name] = value
class_ = all_aes[class_name] # get clss by name
dataset = data.get_dataset(
dataset, dict(batch_size=batch)
) # get dataset class, including train/test PrefetchDataset
ae = class_(dataset, '/' + os.path.join(*(folders[:-2])), **params)
if return_dataset:
return ae, dataset
else:
return ae, folders[-1]
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
print("scales: ", scales)
model = MIDRAE(dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
advweight=FLAGS.advweight,
advdepth=FLAGS.advdepth or FLAGS.depth,
reg=FLAGS.reg,
advnoise=FLAGS.advnoise,
advfake=FLAGS.advfake,
wgt_mmd=FLAGS.wgt_mmd)
model.train()
def main(argv):
del argv # Unused.
batch = FLAGS.batch
dataset = data.get_dataset(FLAGS.dataset, dict(batch_size=batch))
scales = int(round(math.log(dataset.width // FLAGS.latent_width, 2)))
model = GIDRAE2(dataset,
FLAGS.train_dir,
latent=FLAGS.latent,
depth=FLAGS.depth,
scales=scales,
advweight=FLAGS.advweight,
advdepth=FLAGS.advdepth or FLAGS.depth,
reg=FLAGS.reg,
int_hidden_layers=FLAGS.int_hidden_layers,
int_hidden_units=FLAGS.int_hidden_units,
disc_hidden_layers=FLAGS.disc_hidden_layers,
beta=FLAGS.beta,
wgt_mmd=FLAGS.wgt_mmd,
use_ema=FLAGS.use_ema,
n_scale=FLAGS.n_scale,
wgt_noise=FLAGS.wgt_noise,
wgt_fake=FLAGS.wgt_fake)
model.train(summary_kimg=FLAGS.summary_kimg)
use_cuda = torch.cuda.is_available()
if use_cuda:
torch.backends.cudnn.benchmark = True
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
if args.seed is not None:
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
print('=> Preparing data..')
train_loader, val_loader, n_class = get_dataset(args.dataset,
args.batch_size,
args.n_worker,
data_root=args.data_root)
net = get_model() # for measure
IMAGE_SIZE = 224 if args.dataset == 'imagenet' else 32
n_flops, n_params = measure_model(net, IMAGE_SIZE, IMAGE_SIZE)
print('=> Model Parameter: {:.3f} M, FLOPs: {:.3f}M'.format(
n_params / 1e6, n_flops / 1e6))
del net
net = get_model() # real training
if args.ckpt_path is not None: # assigned checkpoint path to resume from
print('=> Resuming from checkpoint..')
checkpoint = torch.load(args.ckpt_path)
sd = checkpoint[
envs = [args.domain]
elif args.domain == 'both':
envs = ['MIMIC', 'CXP']
else:
raise NotImplementedError
common_args = {
'envs': envs,
'only_frontal': True,
'subset_label': args.target,
'output_type': args.data_type,
'imagenet_norm': not args.data_type == 'fft',
'augment': int(args.augment),
}
train_data = get_dataset(split='train', **common_args)
val_data = get_dataset(split='val', **common_args)
test_data = get_dataset(split='test', **common_args)
test_data_aug = get_dataset(
augment=1,
split='test',
**{i: common_args[i]
for i in common_args if i != 'augment'})
if args.debug:
val_data = Subset(val_data, list(range(1024)))
test_data = Subset(test_data, list(range(1024)))
test_data_aug = Subset(test_data_aug, list(range(1024)))
else:
val_data = Subset(
val_data,
def train_and_evaluate(config, workdir):
"""Runs a training and evaluation loop.
Args:
config: Configuration to use.
workdir: Working directory for checkpoints and TF summaries. If this
contains checkpoint, training will be resumed from the latest checkpoint.
Returns:
Training state.
"""
rng = jax.random.PRNGKey(config.seed)
rng, data_rng = jax.random.split(rng)
# Make sure config defines num_epochs and num_train_steps appropriately.
utils.check_epochs_and_steps(config)
train_preprocessing_fn, eval_preprocessing_fn = data.parse_preprocessing_strings(
config.get("train_preprocess_str", ""),
config.get("eval_preprocess_str", ""))
assert config.batch_size % jax.local_device_count() == 0, (
f"Batch size ({config.batch_size}) should be divisible by number of "
f"devices ({jax.local_device_count()}).")
per_device_batch_size = config.batch_size // jax.local_device_count()
train_ds, eval_ds, num_classes = data.get_dataset(
config.dataset,
per_device_batch_size,
data_rng,
train_preprocessing_fn=train_preprocessing_fn,
eval_preprocessing_fn=eval_preprocessing_fn,
**config.get("data", {}))
module = ClassificationModule.partial(config=config,
num_classes=num_classes)
optimizer = create_optimizer(config)
# Enables relevant statistics aggregator.
stats_aggregators = []
train_metrics_dict = {
"train_loss": classification_lib.cross_entropy,
"train_accuracy": classification_lib.accuracy
}
eval_metrics_dict = {
"eval_loss": classification_lib.cross_entropy,
"eval_accuracy": classification_lib.accuracy
}
loss_fn = classification_lib.cross_entropy
def loss_from_stats(field, multiplier):
return lambda logits, labels, stats: multiplier * stats[field]
# Add some regularizer to the loss if needed.
if (config.model == "patchnet" and config.selection_method not in [
sample_patches.SelectionMethod.HARD_TOPK,
sample_patches.SelectionMethod.RANDOM
]):
entropy_regularizer = config.get("entropy_regularizer", 0.)
entropy_before_normalization = config.get(
"entropy_before_normalization", False)
stat_field = "entropy"
if entropy_before_normalization:
stat_field = "entropy_before_normalization"
if entropy_regularizer != 0.:
logging.info(
"Add entropy regularizer %s normalization to the loss %f.",
"before" if entropy_before_normalization else "after",
entropy_regularizer)
loss_fn = [
loss_fn,
loss_from_stats(stat_field, entropy_regularizer)
]
def entropy_aggregator(stats):
return {
stat_field: stats[stat_field],
}
stats_aggregators.append(entropy_aggregator)
def add_image_prefix(image_aggregator):
def aggregator(stats):
d = image_aggregator(stats)
return {f"image_{k}": v for k, v in d.items()}
return aggregator
if config.model == "patchnet" and config.get("log_images", True):
@add_image_prefix
def plot_patches(stats):
keys = ["extracted_patches", "x", "scores"]
return {k: stats[k] for k in keys if k in stats}
stats_aggregators.append(plot_patches)
state = classification_lib.training_loop(
module=module,
rng=rng,
train_ds=train_ds,
eval_ds=eval_ds,
loss_fn=loss_fn,
optimizer=optimizer,
train_metrics_dict=train_metrics_dict,
eval_metrics_dict=eval_metrics_dict,
stats_aggregators=stats_aggregators,
config=config,
workdir=workdir)
return state
type=str,
default="Gaussian",
help="type of prior",
choices=["Gaussian", "MixtureGaussians"])
args = parser.parse_args()
if args.device == 'cpu' or (not torch.cuda.is_available()):
device = 'cpu'
else:
device = 'cuda:' + str(args.device)
# Target Logistic regression, and synthetic data
data_X, data_Y = get_dataset(m=args.data_size,
d=args.dim,
type_regression="logistic",
type_data=args.type_data,
data_dir="./data/")
data_X = data_X.to(device=device)
data_Y = data_Y.to(device=device)
# path numerical results
dim = data_X.shape[1]
data_size = data_X.shape[0]
path_results = "./numerical_results/mlmc_subsampling/logistic/prior_{}/{}_d{}_m{}".format(
args.prior, args.type_data, dim, data_size)
if not os.path.exists(path_results):
os.makedirs(path_results)
with open(os.path.join(path_results, 'commandline_args.json'), 'w') as f:
json.dump(vars(args), f)
parser.add_argument('--T', type=int, default=10, help='horizon time')
parser.add_argument('--n_steps', type=int, default=10000, help='number of steps in time discretisation')
parser.add_argument('--seed', type=int, default=1, help='seed')
parser.add_argument('--type_data', type=str, default="synthetic", help="type of data")
parser.add_argument('--prior', type=str, default="Gaussian", help="type of prior")
args = parser.parse_args()
if args.device=='cpu' or (not torch.cuda.is_available()):
device='cpu'
else:
device = 'cuda:'+str(args.device)
# Target Logistic regression, and synthetic data
data_X, data_Y = get_dataset(m=args.data_size, d=args.dim,
type_regression="MixtureGaussians", type_data=args.type_data, data_dir="./data/")
data_X = data_X.to(device=device)
try:
data_Y = data_Y.to(device=device)
except:
pass
# path numerical results
dim = data_X.shape[1]
data_size = data_X.shape[0]
path_results = "./numerical_results/mlmc_subsampling/mixture/prior_{}/{}_d{}_m{}".format(args.prior, args.type_data, dim, data_size)
if not os.path.exists(path_results):
os.makedirs(path_results)
with open(os.path.join(path_results, 'commandline_args.json'), 'w') as f:
json.dump(vars(args), f)
