def make_encoder(params, is_training):
network_type = params['network_type']
if network_type == 'fully_connected':
encoder_ = fully_connected_encoder(params, is_training)
elif network_type == 'conv':
encoder_ = conv_encoder(params, is_training)
elif network_type == 'infoGAN':
encoder_ = infoGAN_encoder(params, is_training)
elif network_type == 'resnet':
encoder_ = resnet_encoder(params, is_training)
else:
raise NotImplementedError("Network type not implemented.")
def encoder_spec():
x = tf.compat.v1.placeholder(tf.float32, shape=params['full_size'])
z = encoder_(x)
hub.add_signature(inputs={'x': x}, outputs={'z': z})
enc_spec = hub.create_module_spec(encoder_spec)
encoder = hub.Module(enc_spec, name='encoder', trainable=True)
hub.register_module_for_export(encoder, "encoder")
return encoder
def make_decoder(params, is_training):
network_type = params['network_type']
if network_type == 'fully_connected':
decoder_ = fully_connected_decoder(params, is_training)
elif network_type == 'conv':
decoder_ = conv_decoder(params, is_training)
elif network_type == 'infoGAN':
decoder_ = infoGAN_decoder(params, is_training)
elif network_type == 'resnet':
decoder_ = resnet_decoder(params, is_training)
else:
raise NotImplementedError("Network type not implemented.")
def decoder_spec():
z = tf.compat.v1.placeholder(tf.float32,
shape=[None, params['latent_size']])
x = decoder_(z)
hub.add_signature(inputs={'z': z}, outputs={'x': x})
dec_spec = hub.create_module_spec(decoder_spec)
decoder = hub.Module(dec_spec, name='decoder', trainable=True)
hub.register_module_for_export(decoder, "decoder")
return decoder
def model_fn(features, labels, mode, params):
training = mode == tf.estimator.ModeKeys.TRAIN
spec = tfhub.create_module_spec(module_fn,
tags_and_args=[
({'train'}, {
'training': True,
'params': params
}),
(set(), {
'training': False,
'params': params
}),
])
tags = {'train'} if training else None
module = tfhub.Module(spec, trainable=training, tags=tags)
tfhub.register_module_for_export(module, 'doodle')
image = features['image']
image.shape.assert_is_compatible_with([None, 28, 28, 1])
predictions = module(image, as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
export_outputs={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
tf.estimator.export.PredictOutput(predictions),
})
with tf.variable_scope('losses'):
cross_entropy_loss = tf.losses.sparse_softmax_cross_entropy(
labels=labels, logits=predictions['logits'])
total_loss = tf.losses.get_total_loss()
tf.summary.image('image', image)
tf.summary.scalar('total_loss', total_loss)
metric_ops = metrics.calculate(labels, predictions['classes'],
params['num_classes'])
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
eval_metric_ops=metric_ops)
if mode == tf.estimator.ModeKeys.TRAIN:
global_step = tf.train.get_or_create_global_step()
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.variable_scope('optimizer'):
optimizer = tf.train.AdamOptimizer(params['learning_rate'])
with tf.control_dependencies(update_ops):
fit = optimizer.minimize(total_loss, global_step)
return tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=fit,
eval_metric_ops=metric_ops)
def body(self, features):
hparams = self.hparams
model_opt = {
'nr_resnet': hparams.nr_resnet,
'nr_filters': hparams.hidden_size,
'nr_logistic_mix': 1,
'resnet_nonlinearity': 'concat_elu',
'energy_distance': False,
'dropout_p': hparams.dropout
}
def pixel_cnn_fn(input_layer):
model = tf.make_template('model', model_spec)
out = model(input_layer, None, ema=None, **model_opt)
out = tf.layers.dense(out, 2, activation=None)
loc, scale = tf.split(out, num_or_size_splits=2, axis=-1)
scale = tf.nn.softplus(scale) + 1e-4
distribution = tfp.distributions.Independent(
tfp.distributions.Normal(loc=loc, scale=scale))
sample = distribution.sample()
log_prob = distribution.log_prob(input_layer)
grads = tf.gradients(log_prob, input_layer)[0]
print(grads)
output = {
'sample': sample,
'log_prob': log_prob,
'logits': out,
'mu': loc,
'scale': scale,
'grads': grads
}
return output
# During predict, we use a tf hub module, otherwise we stick to normal
# behavior to allow for multi gpu training
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
def make_model_spec():
input_layer = tf.placeholder(
tf.float32, shape=features["inputs"].get_shape())
outputs = pixel_cnn_fn(input_layer)
hub.add_signature(inputs=input_layer, outputs=outputs)
spec = hub.create_module_spec(make_model_spec,
drop_collections=['checkpoints'])
pixelcnn = hub.Module(spec, name="pixelcnn_module", trainable=True)
hub.register_module_for_export(pixelcnn, "pixelcnn")
output = pixelcnn(features["inputs"], as_dict=True)
else:
with tf.variable_scope('pixelcnn_module'):
output = pixel_cnn_fn(features["inputs"])
self.image_summary("inputs", features["inputs"])
self.image_summary("samples", output["sample"])
return output["logits"], {"training": -output["log_prob"]}
def flow_model_fn(features, labels, mode, params, config):
"""
Model function to create a VAE estimator
"""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
if mode == tf.estimator.ModeKeys.PREDICT:
y = features
def flow_module_spec():
cond_layer = tf.placeholder(tf.float32, shape=[None, n_cond])
flow = params['flow_fn'](cond_layer, is_training)
hub.add_signature(inputs=cond_layer,
outputs=flow.sample(tf.shape(cond_layer)[0]))
flow_spec = hub.create_module_spec(flow_module_spec)
flow = hub.Module(flow_spec, name='flow_module')
hub.register_module_for_export(flow, "code_sampler")
predictions = {'code': flow(y)}
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
x = features['x']
y = features['y']
# Loads the encoding function to work on the images
code = x
with tf.variable_scope("flow_module"):
cond_layer = y
flow = params['flow_fn'](cond_layer, is_training)
loglikelihood = flow.log_prob(code)
# This is the loglikelihood of a batch of images
tf.summary.scalar('loglikelihood', tf.reduce_mean(loglikelihood))
loss = -tf.reduce_mean(loglikelihood)
# Training of the model
global_step = tf.train.get_or_create_global_step()
learning_rate = tf.train.cosine_decay(params["learning_rate"], global_step,
params["max_steps"])
tf.summary.scalar("learning_rate", learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
eval_metric_ops = {
"loglikelihood": tf.metrics.mean(tf.reduce_mean(loglikelihood))
}
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def create_ict_module(params, mode):
"""Create hub module."""
tags_and_args = []
for is_training in (True, False):
tags = set()
if is_training:
tags.add("train")
tags_and_args.append((tags, dict(is_training=is_training)))
ict_module_spec = hub.create_module_spec(functools.partial(module_fn,
params=params),
tags_and_args=tags_and_args)
ict_module = hub.Module(
ict_module_spec,
tags={"train"} if mode == tf.estimator.ModeKeys.TRAIN else {},
trainable=True)
hub.register_module_for_export(ict_module, "ict")
return ict_module
def _model_fn(features, labels, mode):
"""A model_fn that uses a mock TF-Hub module."""
del labels
spec = hub.create_module_spec(text_module_fn)
embedding = hub.Module(spec)
if register_module:
hub.register_module_for_export(embedding, _EXPORT_MODULE_NAME)
predictions = embedding(features[_TEXT_FEATURE_NAME])
loss = tf.constant(0.0)
global_step = tf.train.get_global_step()
train_op = tf.assign_add(global_step, 1)
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op)
def make_hub_predictor(model_fn):
"""Creates a tf.hub module for export (in PREDICT mode).
Args:
model_fn: This function is called with the placeholder inputs and
is_training as arguments and should call the model, returning both the
end_points collection and the tensor that should become the hub module's
default prediction (for the default signature).
Returns:
The tf.hub module.
"""
# This defines a function called by the hub module to create the model's
# graph in a new/empty tf.Graph, hence it creates the placeholder etc.
def create_model_fn(is_training): # pylint: disable=missing-docstring
input_shape = utils.str2intlist(FLAGS.serving_input_shape)
img = tf.placeholder(shape=input_shape, dtype=tf.float32)
# This is an example of calling `apply_model_semi` with only one of the
# inputs provided. The outputs will simply use the given names:
end_points, predictions = model_fn(img, is_training)
# Register both the class output and all endpoints to the hub module.
hub.add_signature(inputs={'image': img}, outputs=predictions)
hub.add_signature(inputs={'image': img},
outputs=end_points,
name='representation')
tf_hub_module_spec = hub.create_module_spec(
create_model_fn,
[(["is_training"], {
'is_training': True
}), (set(), {
'is_training': False
})],
# For some not understood reason, this is necessary when the model uses
# cross_replica_batch_norm. We verified that moving averages are still
# being stored in the hub module just fine.
drop_collections=[tf.GraphKeys.MOVING_AVERAGE_VARIABLES])
tf_hub_module = hub.Module(tf_hub_module_spec, trainable=False, tags=set())
hub.register_module_for_export(tf_hub_module, export_name='module')
return tf_hub_module
def _model_fn(features, labels, mode):
"""A model_fn that uses a mock TF-Hub module."""
del labels
spec = hub.create_module_spec(text_module_fn)
embedding = hub.Module(spec)
if register_module:
hub.register_module_for_export(embedding, _EXPORT_MODULE_NAME)
predictions = embedding(features[_TEXT_FEATURE_NAME])
loss = tf.constant(0.0)
global_step = tf.train.get_global_step()
train_op = tf.assign_add(global_step, 1)
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op)
def model_fn(data, mode):
"""Produces a loss for the rotation task.
Args:
data: Dict of inputs containing, among others, "image" and "label."
mode: model's mode: training, eval or prediction
Returns:
EstimatorSpec
"""
num_angles = 4
images = data['image']
if mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL]:
images = tf.reshape(images, [-1] + images.get_shape().as_list()[-3:])
with tf.variable_scope('module'):
image_fn = lambda: images
logits = apply_model(
image_fn=image_fn,
is_training=(mode == tf.estimator.ModeKeys.TRAIN),
num_outputs=num_angles,
make_signature=False)
else:
input_shape = utils.str2intlist(
FLAGS.get_flag_value('serving_input_shape', 'None,None,None,3'))
image_fn = lambda: tf.placeholder(
shape=input_shape, # pylint: disable=g-long-lambda
dtype=tf.float32)
apply_model_function = functools.partial(apply_model,
image_fn=image_fn,
num_outputs=num_angles,
make_signature=True)
tf_hub_module_spec = hub.create_module_spec(apply_model_function,
[(utils.TAGS_IS_TRAINING, {
'is_training': True
}),
(set(), {
'is_training': False
})])
tf_hub_module = hub.Module(tf_hub_module_spec,
trainable=False,
tags=set())
hub.register_module_for_export(tf_hub_module, export_name='module')
logits = tf_hub_module(images)
return trainer.make_estimator(mode, predictions=logits)
labels = tf.reshape(data['label'], [-1])
# build loss and accuracy
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,
logits=logits)
loss = tf.reduce_mean(loss)
eval_metrics = (
lambda labels, logits: { # pylint: disable=g-long-lambda
'accuracy':
tf.metrics.accuracy(labels=labels,
predictions=tf.argmax(logits, axis=-1))
},
[labels, logits])
return trainer.make_estimator(mode, loss, eval_metrics, logits)
def _mdn_inv_model_fn(features,
labels,
nchannels,
n_y,
n_mixture,
dropout,
optimizer,
mode,
logoffset=1e-3,
log=True):
'''Train inverse model i.e. go from the halo field to matter overdensity
'''
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
# Builds the neural network
net = slim.conv3d(feature_layer,
16,
5,
activation_fn=tf.nn.leaky_relu,
padding='same')
#net = wide_resnet(feature_layer, 8, activation_fn=tf.nn.leaky_relu, is_training=is_training)
net = wide_resnet(net,
16,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = wide_resnet(net,
32,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = wide_resnet(net,
32,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = slim.conv3d(net, 32, 3, activation_fn=tf.nn.leaky_relu)
# Define the probabilistic layer
net = slim.conv3d(net,
3 * n_mixture * nchannels,
1,
activation_fn=None)
cube_size = tf.shape(obs_layer)[1]
net = tf.reshape(
net,
[-1, cube_size, cube_size, cube_size, nchannels, n_mixture * 3])
logits, loc, unconstrained_scale = tf.split(net,
num_or_size_splits=3,
axis=-1)
print('\nloc :\n', loc)
scale = tf.nn.softplus(unconstrained_scale[...]) + 1e-3
distribution = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits[...]),
#components_distribution=tfd.MultivariateNormalDiag(loc=loc[...,0], scale_diag=scale))
components_distribution=tfd.Normal(loc=loc[...], scale=scale))
print('\ngmm\n', distribution)
# Define a function for sampling, and a function for estimating the log likelihood
if log:
print('Logged it')
sample = tf.exp(distribution.sample()) - logoffset
print('\ninf dist sample :\n', distribution.sample())
logfeature = tf.log(tf.add(logoffset, obs_layer), 'logfeature')
print('\nlogfeature :\n', logfeature)
prob = distribution.prob(logfeature[...])
loglik = distribution.log_prob(logfeature[...])
else:
print('UnLogged it')
sample = distribution.sample()
print('\ninf dist sample :\n', distribution.sample())
loglik = distribution.log_prob(obs_layer[...])
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik,
'sigma': scale,
'mean': loc,
'logits': logits
})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
features_ = features['features']
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
features_ = features
samples = predictions['sample']
print('\nsamples :\n', samples)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
loglik = predictions['loglikelihood']
print('\nloglik :\n', loglik)
####Compute and register loss function
neg_log_likelihood = -tf.reduce_sum(loglik, axis=-1)
neg_log_likelihood = tf.reduce_mean(neg_log_likelihood)
tf.losses.add_loss(neg_log_likelihood)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = list(np.array([1e4, 2e4, 4e4, 5e4]).astype(int))
values = [1e-3, 5e-4, 1e-4, 5e-5, 1e-5]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('rate', learning_rate)
tf.summary.scalar('total_loss', total_loss)
tf.summary.scalar('loss', neg_log_likelihood)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def _mdn_model_fn(features, labels, n_y, n_mixture, dropout, optimizer, mode):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
# Builds the neural network
net = slim.conv3d(feature_layer,
16,
5,
activation_fn=tf.nn.leaky_relu,
padding='valid')
#net = wide_resnet(feature_layer, 8, activation_fn=tf.nn.leaky_relu, is_training=is_training)
net = wide_resnet(net,
16,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = wide_resnet(net,
32,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = wide_resnet(net,
32,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = slim.conv3d(net, 32, 3, activation_fn=tf.nn.tanh)
# Define the probabilistic layer
net = slim.conv3d(net, n_mixture * 3 * n_y, 1, activation_fn=None)
cube_size = tf.shape(obs_layer)[1]
net = tf.reshape(
net, [-1, cube_size, cube_size, cube_size, n_y, n_mixture * 3])
# net = tf.reshape(net, [None, None, None, None, n_y, n_mixture*3])
loc, unconstrained_scale, logits = tf.split(net,
num_or_size_splits=3,
axis=-1)
scale = tf.nn.softplus(unconstrained_scale)
# Form mixture of discretized logistic distributions. Note we shift the
# logistic distribution by -0.5. This lets the quantization capture "rounding"
# intervals, `(x-0.5, x+0.5]`, and not "ceiling" intervals, `(x-1, x]`.
discretized_logistic_dist = tfd.QuantizedDistribution(
distribution=tfd.TransformedDistribution(
distribution=tfd.Logistic(loc=loc, scale=scale),
bijector=tfb.AffineScalar(shift=-0.5)),
low=0.,
high=2.**4 - 1)
mixture_dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits),
components_distribution=discretized_logistic_dist)
# Define a function for sampling, and a function for estimating the log likelihood
sample = tf.squeeze(mixture_dist.sample())
loglik = mixture_dist.log_prob(obs_layer)
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik
})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
loglik = predictions['loglikelihood']
# Compute and register loss function
neg_log_likelihood = -tf.reduce_sum(loglik, axis=-1)
neg_log_likelihood = tf.reduce_mean(neg_log_likelihood)
tf.losses.add_loss(neg_log_likelihood)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = [15000, 30000, 45000, 60000]
values = [0.00001, 0.000005, 0.000001, 0.0000005, 0.0000001]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('loss', neg_log_likelihood)
tf.summary.scalar('rate', learning_rate)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def _mdn_nozero_model_fn(features,
labels,
nchannels,
n_y,
n_mixture,
dropout,
optimizer,
mode,
pad,
lr0=1e-3,
dsitribution='logistic'):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
mask_layer = tf.clip_by_value(obs_layer, 0, 1)
#
# Builds the neural network
net = slim.conv3d(feature_layer,
16,
3,
activation_fn=tf.nn.leaky_relu,
padding='valid')
subnet = tf.identity(net[:, 3:-3, 3:-3, 3:-3, :])
net = valid_resnet(net,
16,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = slim.conv3d(net,
16,
3,
activation_fn=tf.nn.leaky_relu,
padding='valid')
net = net + subnet
# Define the probabilistic layer
likenet = slim.conv3d(net, 64, 1, activation_fn=tf.nn.leaky_relu)
net = slim.conv3d(likenet, n_mixture * 3 * n_y, 1, activation_fn=None)
# Define the probabilistic layer
cube_size = tf.shape(obs_layer)[1]
net = tf.reshape(
net, [-1, cube_size, cube_size, cube_size, n_y, n_mixture * 3])
# net = tf.reshape(net, [None, None, None, None, n_y, n_mixture*3])
loc, unconstrained_scale, logits = tf.split(net,
num_or_size_splits=3,
axis=-1)
scale = tf.nn.softplus(unconstrained_scale) + 1e-3
# Form mixture of discretized logistic distributions. Note we shift the
# logistic distribution by -0.5. This lets the quantization capture "rounding"
# intervals, `(x-0.5, x+0.5]`, and not "ceiling" intervals, `(x-1, x]`.
if distribution == 'logistic':
discretized_logistic_dist = tfd.QuantizedDistribution(
distribution=tfd.TransformedDistribution(
distribution=tfd.Logistic(loc=loc, scale=scale),
bijector=tfb.AffineScalar(shift=-0.5)),
low=0.,
high=2.**3 - 1)
mixture_dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits),
components_distribution=discretized_logistic_dist)
elif distribution == 'normal':
mixture_dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits),
components_distribution=tfd.Normal(loc=loc, scale=scale))
# Define a function for sampling, and a function for estimating the log likelihood
#sample = tf.squeeze(mixture_dist.sample())
sample = mixture_dist.sample()
loglik = mixture_dist.log_prob(obs_layer)
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik,
'loc': loc,
'scale': scale,
'logits': logits
})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
loss = -predictions['loglikelihood']
# Compute and register loss function
loss = tf.reduce_mean(loss)
tf.losses.add_loss(loss)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = list(np.array([1e4, 2e4, 4e4, 5e4, 6e4]).astype(int))
values = [lr0, lr0 / 2, lr0 / 10, lr0 / 20, lr0 / 100, lr0 / 1000]
#values = [1e-3, 5e-4, 1e-4, 5e-5, 1e-5, 1e-6]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('rate', learning_rate)
logging_hook = tf.train.LoggingTensorHook(
{
"iter": global_step,
"loss": loss
}, every_n_iter=50)
tf.summary.scalar('loss', loss)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
training_hooks=[logging_hook])
def _mdn_unetmodel_fn(features,
labels,
nchannels,
n_y,
n_mixture,
dropout,
optimizer,
mode,
pad,
fsize=8,
nsub=2,
distribution='logistic'):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
# Builds the neural network
if pad == 0:
d00 = slim.conv3d(feature_layer,
fsize,
5,
activation_fn=tf.nn.leaky_relu,
padding='same')
elif pad == 2:
d00 = slim.conv3d(feature_layer,
fsize,
5,
activation_fn=tf.nn.leaky_relu,
padding='valid')
if pad == 4:
d00 = slim.conv3d(feature_layer,
fsize,
5,
activation_fn=tf.nn.leaky_relu,
padding='valid')
d00 = slim.conv3d(d00,
fsize * 2,
5,
activation_fn=tf.nn.leaky_relu,
padding='valid')
## #downsample
## dd = [[d00]]
## cfsize = fsize
## for i in range(nsub):
## d0 = dd[-1][-1]
## d1 = wide_resnet(d0, cfsize, activation_fn=tf.nn.leaky_relu)
## d2 = wide_resnet(d1, cfsize, activation_fn=tf.nn.leaky_relu)
## dsub = slim.max_pool3d(d2, kernel_size=3, stride=2, padding='SAME')
## dd.append([d1, d2, dsub])
## cfsize *= 2
##
## #lower layer
## d0 = dd[-1][-1]
## d1 = wide_resnet(d0, cfsize, activation_fn=tf.nn.leaky_relu)
## d2 = wide_resnet(d1, cfsize, activation_fn=tf.nn.leaky_relu)
##
## up = [[d1, d2]]
## #upsample
## for i in range(nsub):
## cfsize = cfsize // 2
## usub = up[-1][-1]
## dup = dd.pop()
## u0 = dynamic_deconv3d('up%d'%i, usub, shape=[3,3,3,cfsize], activation=tf.nn.leaky_relu)
## #u0 = slim.conv3d_transpose(usub, fsize, kernel_size=3, stride=2)
## uc = tf.concat([u0, dup[1]], axis=-1)
## u1 = wide_resnet(uc, cfsize, activation_fn=tf.nn.leaky_relu)
## u2 = wide_resnet(u1, cfsize, activation_fn=tf.nn.leaky_relu)
## up.append([u0, u1, u1c, u2])
##
## u0 = up[-1][-1]
## net = slim.conv3d(u0, 1, 3, activation_fn=tf.nn.tanh)
##
#downsample #restructure code while doubling filter size
cfsize = fsize
d1 = wide_resnet(d00, cfsize, activation_fn=tf.nn.leaky_relu)
d2 = wide_resnet(d1, cfsize, activation_fn=tf.nn.leaky_relu)
dd = [d2]
for i in range(nsub):
cfsize *= 2
print(i, cfsize)
dsub = slim.max_pool3d(dd[-1],
kernel_size=3,
stride=2,
padding='SAME')
d1 = wide_resnet(dsub, cfsize, activation_fn=tf.nn.leaky_relu)
d2 = wide_resnet(d1, cfsize, activation_fn=tf.nn.leaky_relu)
dd.append(d2)
print(len(dd))
#upsample
usub = dd.pop()
for i in range(nsub):
u0 = dynamic_deconv3d('up%d' % i,
usub,
shape=[3, 3, 3, cfsize],
activation=tf.identity)
cfsize = cfsize // 2
print(i, cfsize)
u0 = slim.conv3d(u0,
cfsize,
1,
activation_fn=tf.identity,
padding='same')
#u0 = slim.conv3d_transpose(usub, fsize, kernel_size=3, stride=2)
uc = tf.concat([u0, dd.pop()], axis=-1)
u1 = wide_resnet(uc, cfsize, activation_fn=tf.nn.leaky_relu)
u2 = wide_resnet(u1, cfsize, activation_fn=tf.nn.leaky_relu)
usub = u2
print(len(dd))
net = slim.conv3d(usub, 1, 3, activation_fn=tf.nn.tanh)
# Define the probabilistic layer
net = slim.conv3d(net, n_mixture * 3 * n_y, 1, activation_fn=None)
cube_size = tf.shape(obs_layer)[1]
net = tf.reshape(
net, [-1, cube_size, cube_size, cube_size, n_y, n_mixture * 3])
# net = tf.reshape(net, [None, None, None, None, n_y, n_mixture*3])
loc, unconstrained_scale, logits = tf.split(net,
num_or_size_splits=3,
axis=-1)
scale = tf.nn.softplus(unconstrained_scale) + 1e-3
# Form mixture of discretized logistic distributions. Note we shift the
# logistic distribution by -0.5. This lets the quantization capture "rounding"
# intervals, `(x-0.5, x+0.5]`, and not "ceiling" intervals, `(x-1, x]`.
if distribution == 'logistic':
discretized_logistic_dist = tfd.QuantizedDistribution(
distribution=tfd.TransformedDistribution(
distribution=tfd.Logistic(loc=loc, scale=scale),
bijector=tfb.AffineScalar(shift=-0.5)),
low=0.,
high=2.**3 - 1)
mixture_dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits),
components_distribution=discretized_logistic_dist)
elif distribution == 'normal':
mixture_dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits),
components_distribution=tfd.Normal(loc=loc, scale=scale))
# Define a function for sampling, and a function for estimating the log likelihood
#sample = tf.squeeze(mixture_dist.sample())
sample = mixture_dist.sample()
loglik = mixture_dist.log_prob(obs_layer)
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik,
'loc': loc,
'scale': scale,
'logits': logits
})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
loglik = predictions['loglikelihood']
# Compute and register loss function
neg_log_likelihood = -tf.reduce_sum(loglik, axis=-1)
neg_log_likelihood = tf.reduce_mean(neg_log_likelihood)
tf.losses.add_loss(neg_log_likelihood)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = list(np.array([1e4, 2e4, 4e4, 5e4, 6e4]).astype(int))
values = [1e-3, 5e-4, 1e-4, 5e-5, 1e-5, 1e-6]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('rate', learning_rate)
tf.summary.scalar('loss', neg_log_likelihood)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def _mdn_pixmodel_fn(features,
labels,
nchannels,
n_y,
n_mixture,
dropout,
optimizer,
mode,
pad,
cfilter_size=None,
f_map=8):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
conditional_im = wide_resnet(feature_layer,
16,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
conditional_im = wide_resnet(conditional_im,
16,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
conditional_im = wide_resnet(conditional_im,
1,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
conditional_im = tf.concat((feature_layer, conditional_im), -1)
# Builds the neural network
ul = [[obs_layer]]
for i in range(10):
ul.append(
PixelCNN3Dlayer(i,
ul[i],
f_map=f_map,
full_horizontal=True,
h=None,
conditional_im=conditional_im,
cfilter_size=cfilter_size,
gatedact='sigmoid'))
h_stack_in = ul[-1][-1]
with tf.variable_scope("fc_1"):
fc1 = GatedCNN([1, 1, 1, 1],
h_stack_in,
orientation=None,
gated=False,
mask='b').output()
with tf.variable_scope("fc_2"):
fc2 = GatedCNN([1, 1, 1, n_mixture * 3 * n_y],
fc1,
orientation=None,
gated=False,
mask='b',
activation=False).output()
cube_size = tf.shape(obs_layer)[1]
net = tf.reshape(
fc2, [-1, cube_size, cube_size, cube_size, n_y, n_mixture * 3])
loc, unconstrained_scale, logits = tf.split(net,
num_or_size_splits=3,
axis=-1)
scale = tf.nn.softplus(unconstrained_scale) + 1e-3
# Form mixture of discretized logistic distributions. Note we shift the
# logistic distribution by -0.5. This lets the quantization capture "rounding"
# intervals, `(x-0.5, x+0.5]`, and not "ceiling" intervals, `(x-1, x]`.
# discretized_logistic_dist = tfd.QuantizedDistribution(
# distribution=tfd.TransformedDistribution(
# distribution=tfd.Logistic(loc=loc, scale=scale),
# bijector=tfb.AffineScalar(shift=-0.5)),
# low=0.,
# high=2.**3-1)
mixture_dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(logits=logits),
components_distribution=tfd.Normal(loc, scale))
# Define a function for sampling, and a function for estimating the log likelihood
#sample = tf.squeeze(mixture_dist.sample())
sample = mixture_dist.sample()
loglik = mixture_dist.log_prob(obs_layer)
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik,
'loc': loc,
'scale': scale,
'logits': logits
})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
loglik = predictions['loglikelihood']
# Compute and register loss function
neg_log_likelihood = -tf.reduce_sum(loglik, axis=-1)
neg_log_likelihood = tf.reduce_mean(neg_log_likelihood)
tf.losses.add_loss(neg_log_likelihood)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = list(
np.array([500, 1e3, 3e3, 1e4, 2e4, 3e4, 4e4]).astype(int))
values = [1e-3, 1e-3, 1e-3, 5e-4, 1e-4, 5e-5, 1e-5, 1e-6]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('rate', learning_rate)
tf.summary.scalar('loss', neg_log_likelihood)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def model_fn(features, labels, mode, params):
"""Model function."""
del labels
# ==============================
# Input features
# ==============================
# [batch_size, query_seq_len]
query_inputs = features["query_inputs"]
# [batch_size, num_candidates, candidate_seq_len]
candidate_inputs = features["candidate_inputs"]
# [batch_size, num_candidates, query_seq_len + candidate_seq_len]
joint_inputs = features["joint_inputs"]
# [batch_size, num_masks]
mlm_targets = features["mlm_targets"]
mlm_positions = features["mlm_positions"]
mlm_mask = features["mlm_mask"]
# ==============================
# Create modules.
# ==============================
bert_module = hub.Module(
spec=params["bert_hub_module_handle"],
name="locbert",
tags={"train"} if mode == tf.estimator.ModeKeys.TRAIN else {},
trainable=True)
hub.register_module_for_export(bert_module, "locbert")
embedder_module = hub.Module(
spec=params["embedder_hub_module_handle"],
name="embedder",
tags={"train"} if mode == tf.estimator.ModeKeys.TRAIN else {},
trainable=True)
hub.register_module_for_export(embedder_module, "embedder")
if params["share_embedders"]:
query_embedder_module = embedder_module
else:
query_embedder_module = hub.Module(
spec=params["embedder_hub_module_handle"],
name="embedder",
tags={"train"} if mode == tf.estimator.ModeKeys.TRAIN else {},
trainable=True)
hub.register_module_for_export(embedder_module, "query_embedder")
# ==============================
# Retrieve.
# ==============================
# [batch_size, projected_size]
query_emb = query_embedder_module(
inputs=dict(
input_ids=query_inputs.token_ids,
input_mask=query_inputs.mask,
segment_ids=query_inputs.segment_ids),
signature="projected")
# [batch_size * num_candidates, candidate_seq_len]
flat_candidate_inputs, unflatten = flatten_bert_inputs(
candidate_inputs)
# [batch_size * num_candidates, projected_size]
flat_candidate_emb = embedder_module(
inputs=dict(
input_ids=flat_candidate_inputs.token_ids,
input_mask=flat_candidate_inputs.mask,
segment_ids=flat_candidate_inputs.segment_ids),
signature="projected")
# [batch_size, num_candidates, projected_size]
unflattened_candidate_emb = unflatten(flat_candidate_emb)
# [batch_size, num_candidates]
retrieval_score = tf.einsum("BD,BND->BN", query_emb,
unflattened_candidate_emb)
# ==============================
# Read.
# ==============================
# [batch_size * num_candidates, query_seq_len + candidate_seq_len]
flat_joint_inputs, unflatten = flatten_bert_inputs(joint_inputs)
# [batch_size * num_candidates, num_masks]
flat_mlm_positions, _ = tensor_utils.flatten(
tf.tile(
tf.expand_dims(mlm_positions, 1), [1, params["num_candidates"], 1]))
batch_size, num_masks = tensor_utils.shape(mlm_targets)
# [batch_size * num_candidates, query_seq_len + candidates_seq_len]
flat_joint_bert_outputs = bert_module(
inputs=dict(
input_ids=flat_joint_inputs.token_ids,
input_mask=flat_joint_inputs.mask,
segment_ids=flat_joint_inputs.segment_ids,
mlm_positions=flat_mlm_positions),
signature="mlm",
as_dict=True)
# [batch_size, num_candidates]
candidate_score = retrieval_score
# [batch_size, num_candidates]
candidate_log_probs = tf.math.log_softmax(candidate_score)
# ==============================
# Compute marginal log-likelihood.
# ==============================
# [batch_size * num_candidates, num_masks]
flat_mlm_logits = flat_joint_bert_outputs["mlm_logits"]
# [batch_size, num_candidates, num_masks, vocab_size]
mlm_logits = tf.reshape(
flat_mlm_logits, [batch_size, params["num_candidates"], num_masks, -1])
mlm_log_probs = tf.math.log_softmax(mlm_logits)
# [batch_size, num_candidates, num_masks]
tiled_mlm_targets = tf.tile(
tf.expand_dims(mlm_targets, 1), [1, params["num_candidates"], 1])
# [batch_size, num_candidates, num_masks, 1]
tiled_mlm_targets = tf.expand_dims(tiled_mlm_targets, -1)
# [batch_size, num_candidates, num_masks, 1]
gold_log_probs = tf.batch_gather(mlm_log_probs, tiled_mlm_targets)
# [batch_size, num_candidates, num_masks]
gold_log_probs = tf.squeeze(gold_log_probs, -1)
# [batch_size, num_candidates, num_masks]
joint_gold_log_probs = (
tf.expand_dims(candidate_log_probs, -1) + gold_log_probs)
# [batch_size, num_masks]
marginal_gold_log_probs = tf.reduce_logsumexp(joint_gold_log_probs, 1)
# [batch_size, num_masks]
float_mlm_mask = tf.cast(mlm_mask, tf.float32)
# []
loss = -tf.div_no_nan(
tf.reduce_sum(marginal_gold_log_probs * float_mlm_mask),
tf.reduce_sum(float_mlm_mask))
# ==============================
# Optimization
# ==============================
num_warmup_steps = min(10000, max(100, int(params["num_train_steps"] / 10)))
train_op = optimization.create_optimizer(
loss=loss,
init_lr=params["learning_rate"],
num_train_steps=params["num_train_steps"],
num_warmup_steps=num_warmup_steps,
use_tpu=params["use_tpu"])
# ==============================
# Evaluation
# ==============================
eval_metric_ops = None if params["use_tpu"] else dict()
if mode != tf.estimator.ModeKeys.PREDICT:
# [batch_size, num_masks]
retrieval_utility = marginal_gold_log_probs - gold_log_probs[:, 0]
retrieval_utility *= tf.cast(features["mlm_mask"], tf.float32)
# []
retrieval_utility = tf.div_no_nan(
tf.reduce_sum(retrieval_utility), tf.reduce_sum(float_mlm_mask))
add_mean_metric("retrieval_utility", retrieval_utility, eval_metric_ops)
has_timestamp = tf.cast(
tf.greater(features["export_timestamp"], 0), tf.float64)
off_policy_delay_secs = (
tf.timestamp() - tf.cast(features["export_timestamp"], tf.float64))
off_policy_delay_mins = off_policy_delay_secs / 60.0
off_policy_delay_mins *= tf.cast(has_timestamp, tf.float64)
add_mean_metric("off_policy_delay_mins", off_policy_delay_mins,
eval_metric_ops)
# Create empty predictions to avoid errors when running in prediction mode.
predictions = dict()
if params["use_tpu"]:
return tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
predictions=predictions)
else:
if eval_metric_ops is not None:
# Make sure the eval metrics are updated during training so that we get
# quick feedback from tensorboard summaries when debugging locally.
with tf.control_dependencies([u for _, u in eval_metric_ops.values()]):
loss = tf.identity(loss)
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
predictions=predictions)
def vae_model_fn(features, labels, mode, params, config):
"""
Model function to create a VAE estimator
"""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
# Extract input images
x = features['x']
# Build model functions
encoder_model = make_encoder_fn(params['encoder_fn'],
params['latent_size'],
params['iaf_size'],
is_training=is_training)
decoder_model = partial(params['decoder_fn'], is_training=is_training)
# Define latent prior
prior = tfd.MultivariateNormalDiag(loc=tf.zeros([params['latent_size']]),
scale_identity_multiplier=1.0)
# In predict mode, we encapsulate the model inside a module for exporting
# This is because of a weird bug that makes training MAF unstable inside a
# module
if mode == tf.estimator.ModeKeys.PREDICT:
image_size = x.shape[-2]
n_channels = x.shape[-1]
latent_size = params['latent_size']
def make_encoder_spec():
input_layer = tf.placeholder(
tf.float32, shape=[None, image_size, image_size, n_channels])
encoding = encoder_model(input_layer)
sample = encoding.sample()
log_prob = encoding.log_prob(sample)
hub.add_signature(inputs=input_layer,
outputs={
'sample': sample,
'log_prob': log_prob
})
encoder_spec = hub.create_module_spec(make_encoder_spec)
encoder = hub.Module(encoder_spec, name="encoder_module")
def make_decoder_spec():
code = tf.placeholder(tf.float32, shape=[None, latent_size])
output = decoder_model(code)
if not tf.contrib.framework.is_tensor(output):
output = output.sample()
hub.add_signature(inputs=code, outputs=output)
decoder_spec = hub.create_module_spec(make_decoder_spec)
decoder = hub.Module(decoder_spec, name="decoder_module")
# Register and export encoder and decoder modules
hub.register_module_for_export(encoder, "encoder")
hub.register_module_for_export(decoder, "decoder")
code = encoder(x, as_dict=True)
recon = decoder(code['sample'])
predictions = {
'code': code['sample'],
'reconstruction': recon,
'log_prob': code['log_prob'],
'input': x
}
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
with tf.variable_scope("encoder_module") as sc:
encoding = encoder_model(x)
code = encoding.sample()
log_prob = encoding.log_prob(code)
with tf.variable_scope("decoder_module") as sc:
decoder_output = decoder_model(code)
if tf.contrib.framework.is_tensor(decoder_output):
recon = decoder_output
loglikelihood = params['loglikelihood_fn'](x, recon, features)
else:
# In this case, the decoder is actually returning a distribution
# which we can use to sample from and estimate the lihelihood function
recon = decoder_output.sample()
loglikelihood = decoder_output.log_prob(x)
image_tile_summary("image", tf.to_float(x[:16]), rows=4, cols=4)
if 'psf' in features.keys():
r = tf.expand_dims(tf.spectral.irfft2d(
tf.spectral.rfft2d(recon[:, :, :, 0]) * features['psf']),
axis=-1)
else:
r = recon
image_tile_summary("recon", tf.to_float(r[:16]), rows=4, cols=4)
image_tile_summary("diff", tf.to_float(x[:16] - r[:16]), rows=4, cols=4)
tf.summary.scalar('loglikelihood', tf.reduce_mean(loglikelihood))
kl = log_prob - prior.log_prob(code)
tf.summary.scalar('kl', tf.reduce_mean(kl))
elbo = loglikelihood - kl * params['kl_weight']
loss = -tf.reduce_mean(elbo)
tf.summary.scalar("elbo", tf.reduce_mean(elbo))
# Training of the model
global_step = tf.train.get_or_create_global_step()
learning_rate = tf.train.cosine_decay(params["learning_rate"], global_step,
params["max_steps"])
tf.summary.scalar("learning_rate", learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate,
epsilon=params['adam_epsilon'])
grads_and_vars = optimizer.compute_gradients(loss)
clipped_grads_and_vars = [
(tf.clip_by_norm(grad, params["gradient_clipping"]), var)
for grad, var in grads_and_vars
]
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.apply_gradients(clipped_grads_and_vars,
global_step=global_step)
eval_metric_ops = {
"elbo": tf.metrics.mean(tf.reduce_mean(elbo)),
"kl": tf.metrics.mean(tf.reduce_mean(kl)),
"loglikelihood": tf.metrics.mean(tf.reduce_mean(loglikelihood))
}
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def autoencoder_body(self, features):
""" Customized body function for autoencoders acting on continuous images.
This is based on tensor2tensor.models.research.AutoencoderBasic.body
and should be compatible with most derived classes.
The original autoencoder class relies on embedding the channels to a discrete
vocabulary and defines the loss on that vocab. It's cool and all, but here we
prefer expressing the reconstruction loss as an actual continuous likelihood
function.
"""
hparams = self.hparams
is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
output_activation = tf.nn.softplus if hparams.output_activation == 'softplus' else None
input_shape = [None, ] + common_layers.shape_list(features["inputs"])[1:]
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
# In predict mode, we also define TensorFlow Hub modules for all pieces of
# the autoencoder
if hparams.encode_psf and 'psf' in features:
psf_shape = [None, ] + common_layers.shape_list(features["psf"])[1:]
# First build encoder spec
def make_model_spec():
input_layer = tf.placeholder(tf.float32, shape=input_shape)
x = self.embed(tf.expand_dims(input_layer, -1))
x, encoder_layers = self.encoder(x)
b, b_loss = self.bottleneck(x)
hub.add_signature(inputs=input_layer, outputs=b)
def make_model_spec_psf():
input_layer = tf.placeholder(tf.float32, shape=input_shape)
psf_layer = tf.placeholder(tf.float32, shape=psf_shape)
x = self.embed(tf.expand_dims(input_layer, -1))
# If we have access to the PSF, we add this information to the encoder
if hparams.encode_psf and 'psf' in features:
psf_image = tf.expand_dims(tf.signal.irfft2d(tf.cast(psf_layer[...,0], tf.complex64)), axis=-1)
# Roll the image to undo the fftshift, assuming x1 zero padding and x2 subsampling
psf_image = tf.roll(psf_image, shift=[input_shape[1], input_shape[2]], axis=[1,2])
psf_image = tf.image.resize_with_crop_or_pad(psf_image, input_shape[1], input_shape[2])
net_psf = tf.layers.conv2d(psf_image,
hparams.hidden_size // 4, 5,
padding='same', name="psf_embed_1")
net_psf = common_layers.layer_norm(net_psf, name="psf_norm")
x, encoder_layers = self.encoder(tf.concat([x, net_psf], axis=-1))
else:
x, encoder_layers = self.encoder(x)
b, b_loss = self.bottleneck(x)
hub.add_signature(inputs={'input':input_layer, 'psf':psf_layer}, outputs=b)
spec = hub.create_module_spec(make_model_spec_psf if hparams.encode_psf else make_model_spec, drop_collections=['checkpoints'])
encoder = hub.Module(spec, name="encoder_module")
hub.register_module_for_export(encoder, "encoder")
if hparams.encode_psf:
code = encoder({'input':features["inputs"], 'psf': features['psf']})
else:
code = encoder(features["inputs"])
b_shape = [None, ] + common_layers.shape_list(code)[1:]
res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
res_size = min(res_size, hparams.max_hidden_size)
# Second build decoder spec
def make_model_spec():
input_layer = tf.placeholder(tf.float32, shape=b_shape)
x = self.unbottleneck(input_layer, res_size)
x = self.decoder(x, None)
reconstr = tf.layers.dense(x, input_shape[-1], name="autoencoder_final",
activation=output_activation)
hub.add_signature(inputs=input_layer, outputs=reconstr)
hub.attach_message("stamp_size", tf.train.Int64List(value=[hparams.problem_hparams.img_len]))
try:
hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale[res] for res in hparams.problem_hparams.resolutions]))
except AttributeError:
hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale]))
spec = hub.create_module_spec(make_model_spec, drop_collections=['checkpoints'])
decoder = hub.Module(spec, name="decoder_module")
hub.register_module_for_export(decoder, "decoder")
reconstr = decoder(code)
return reconstr , {"bottleneck_loss": 0.0}
encoder_layers = None
self.is1d = hparams.sample_width == 1
if (hparams.mode != tf.estimator.ModeKeys.PREDICT
or self._encode_on_predict):
labels = features["targets_raw"]
labels_shape = common_layers.shape_list(labels)
shape = common_layers.shape_list(labels)
with tf.variable_scope('encoder_module'):
x = self.embed(tf.expand_dims(labels, -1))
if shape[2] == 1:
self.is1d = True
# Run encoder.
with tf.variable_scope('encoder_module'):
# If we have access to the PSF, we add this information to the encoder
# Note that we only support single band images so far...
if hparams.encode_psf and 'psf' in features:
psf_image = tf.expand_dims(tf.signal.irfft2d(tf.cast(features['psf'][...,0], tf.complex64)), axis=-1)
# Roll the image to undo the fftshift, assuming x1 zero padding and x2 subsampling
psf_image = tf.roll(psf_image, shift=[input_shape[1], input_shape[2]], axis=[1,2])
psf_image = tf.image.resize_with_crop_or_pad(psf_image, input_shape[1], input_shape[2])
net_psf = tf.layers.conv2d(psf_image,
hparams.hidden_size // 4, 5,
padding='same', name="psf_embed_1")
net_psf = common_layers.layer_norm(net_psf, name="psf_norm")
x, encoder_layers = self.encoder(tf.concat([x, net_psf], axis=-1))
else:
x, encoder_layers = self.encoder(x)
# Bottleneck.
with tf.variable_scope('encoder_module'):
b, b_loss = self.bottleneck(x)
xb_loss = 0.0
b_shape = common_layers.shape_list(b)
self._cur_bottleneck_tensor = b
res_size = common_layers.shape_list(x)[-1]
with tf.variable_scope('decoder_module'):
b = self.unbottleneck(b, res_size)
if not is_training:
x = b
else:
l = 2**hparams.num_hidden_layers
warm_step = int(hparams.bottleneck_warmup_steps * 0.25 * l)
nomix_p = common_layers.inverse_lin_decay(warm_step) + 0.01
if common_layers.should_generate_summaries():
tf.summary.scalar("nomix_p_bottleneck", nomix_p)
rand = tf.random_uniform(common_layers.shape_list(x))
# This is the distance between b and x. Having this as loss helps learn
# the bottleneck function, but if we back-propagated to x it would be
# minimized by just setting x=0 and b=0 -- so we don't want too much
# of the influence of this, and we stop-gradient to not zero-out x.
x_stop = tf.stop_gradient(x)
xb_loss = tf.reduce_mean(tf.reduce_sum(
tf.squared_difference(x_stop, b), axis=-1))
# To prevent this loss from exploding we clip at 1, but anneal clipping.
clip_max = 1.0 / common_layers.inverse_exp_decay(
warm_step, min_value=0.001)
xb_clip = tf.maximum(tf.stop_gradient(xb_loss), clip_max)
xb_loss *= clip_max / xb_clip
x = tf.where(tf.less(rand, nomix_p), b, x)
else:
if self._cur_bottleneck_tensor is None:
b = self.sample()
else:
b = self._cur_bottleneck_tensor
self._cur_bottleneck_tensor = b
res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
res_size = min(res_size, hparams.max_hidden_size)
with tf.variable_scope('decoder_module'):
x = self.unbottleneck(b, res_size)
# Run decoder.
with tf.variable_scope('decoder_module'):
x = self.decoder(x, encoder_layers)
# Cut to the right size and mix before returning.
res = x
if hparams.mode != tf.estimator.ModeKeys.PREDICT:
res = x[:, :shape[1], :shape[2], :]
with tf.variable_scope('decoder_module'):
reconstr = tf.layers.dense(res, shape[-1], name="autoencoder_final",
activation=output_activation)
# We apply an optional apodization of the output before taking the
if hparams.output_apodization > 0:
nx = reconstr.get_shape().as_list()[1]
alpha = 2 * hparams.output_apodization / nx
from scipy.signal.windows import tukey
# Create a tukey window
w = tukey(nx, alpha)
w = np.outer(w,w).reshape((1, nx, nx,1)).astype('float32')
# And penalize non zero things at the border
apo_loss = tf.reduce_mean(tf.reduce_sum(((1.- w)*reconstr)**2, axis=[1,2,3]))
else:
w = 1.0
apo_loss = 0.
# We apply the window
reconstr = reconstr * w
# Optionally regularizes further the output
# Anisotropic TV:
tv = tf.reduce_mean(tf.image.total_variation(reconstr))
# Smoothed Isotropic TV:
#im_dx, im_dy = tf.image.image_gradients(reconstr)
#tv = tf.reduce_sum(tf.sqrt(im_dx**2 + im_dy**2 + 1e-6), axis=[1,2,3])
#tv = tf.reduce_mean(tv)
image_summary("without_psf",tf.reshape(reconstr, labels_shape))
# Apply channel-wise convolution with the PSF if requested
if hparams.apply_psf and 'psf' in features:
output_list = []
for i in range(shape[3]):
output_list.append(tf.squeeze(convolve(tf.expand_dims(reconstr[...,i],-1), tf.cast(features['psf'][...,i], tf.complex64),
zero_padding_factor=1)))
reconstr = tf.stack(output_list,axis=-1)
reconstr = tf.reshape(reconstr,shape)
# Losses.
losses = {
"bottleneck_extra": b_loss,
"bottleneck_l2": hparams.bottleneck_l2_factor * xb_loss,
"total_variation": hparams.total_variation_loss * tv,
"apodization_loss": hparams.apodization_loss * apo_loss,
}
loglik = loglikelihood_fn(labels, reconstr, features, hparams)
targets_loss = tf.reduce_mean(- loglik)
tf.summary.scalar("negloglik", targets_loss)
tf.summary.scalar("bottleneck_loss", b_loss)
# Compute final loss
losses["training"] = targets_loss + b_loss + hparams.bottleneck_l2_factor * xb_loss + hparams.total_variation_loss * tv + hparams.apodization_loss * apo_loss
logits = tf.reshape(reconstr, labels_shape)
image_summary("ae", reconstr)
image_summary("input", labels)
return logits, losses
def estimator_model_fn(cls,
hparams,
features,
labels,
mode,
config=None,
params=None,
decode_hparams=None,
use_tpu=False):
if mode not in [
model_fn_lib.ModeKeys.TRAIN, model_fn_lib.ModeKeys.EVAL,
model_fn_lib.ModeKeys.PREDICT
]:
raise ValueError('Mode not recognized: %s' % mode)
if mode is model_fn_lib.ModeKeys.TRAIN:
is_training = True
else:
is_training = False
hparams = hparams_lib.copy_hparams(hparams)
# Instantiate model
data_parallelism = None
if not use_tpu and config:
data_parallelism = config.data_parallelism
reuse = tf.get_variable_scope().reuse
# Instantiate model
self = cls(hparams,
mode,
data_parallelism=data_parallelism,
decode_hparams=decode_hparams,
_reuse=reuse)
generator_inputs = self.sample_noise()
# rename inputs for clarity
real_data = features['inputs']
img_shape = common_layers.shape_list(real_data)[1:4]
real_data.set_shape([hparams.batch_size] + img_shape)
# To satify the TFGAN API setting real data to none on predict
if mode == tf.estimator.ModeKeys.PREDICT:
real_data = None
optimizers = Optimizers(
tf.compat.v1.train.AdamOptimizer(hparams.generator_lr,
hparams.beta1),
tf.compat.v1.train.AdamOptimizer(hparams.discriminator_lr,
hparams.beta1))
# Creates tfhub modules for both generator and discriminator
def make_discriminator_spec():
input_layer = tf.placeholder(tf.float32, shape=[None] + img_shape)
disc_output = self.discriminator(input_layer, None, mode)
hub.add_signature(inputs=input_layer, outputs=disc_output)
disc_spec = hub.create_module_spec(make_discriminator_spec)
def make_generator_spec():
input_layer = tf.placeholder(
tf.float32,
shape=[None] + common_layers.shape_list(generator_inputs)[1:])
gen_output = self.generator(input_layer, mode)
hub.add_signature(inputs=input_layer, outputs=gen_output)
gen_spec = hub.create_module_spec(make_generator_spec)
# Create the modules
discriminator_module = hub.Module(disc_spec,
name="Discriminator_Module",
trainable=True)
generator_module = hub.Module(gen_spec,
name="Generator_Module",
trainable=True)
# Wraps the modules into functions expected by TF-GAN
def generator(code, mode):
p = hparams
out = generator_module(code)
shape = common_layers.shape_list(out)
# Applying convolution by PSF convolution
if p.apply_psf and 'psf' in features:
out = convolve(out,
tf.cast(features['psf'][..., 0], tf.complex64))
# Adds noise according to the provided power spectrum
noise = tf.spectral.rfft2d(tf.random_normal(out.get_shape()[:3]))
thresholded_ps = tf.where(features['ps'] >= 9,
tf.zeros_like(features['ps']),
tf.sqrt(tf.exp(features['ps'])))
noise = noise * tf.cast(thresholded_ps, tf.complex64)
out = out + tf.expand_dims(tf.spectral.irfft2d(noise), axis=-1)
return out
discriminator = lambda image, conditioning, mode: discriminator_module(
image)
# Make GANModel, which encapsulates the GAN model architectures.
gan_model = get_gan_model(mode,
generator,
discriminator,
real_data,
generator_inputs,
add_summaries=self.summaries)
# Make GANLoss, which encapsulates the losses.
if mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL]:
gan_loss = tfgan_train.gan_loss(gan_model,
self.generator_loss,
self.discriminator_loss,
add_summaries=True)
# Make the EstimatorSpec, which incorporates the GANModel, losses, eval
# metrics, and optimizers (if required).
if mode == tf.estimator.ModeKeys.TRAIN:
get_hooks_fn = tfgan_train.get_sequential_train_hooks(
namedtuples.GANTrainSteps(hparams.gen_steps,
hparams.disc_steps))
estimator_spec = get_train_estimator_spec(gan_model,
gan_loss,
optimizers,
get_hooks_fn,
is_chief=True)
elif mode == tf.estimator.ModeKeys.EVAL:
estimator_spec = get_eval_estimator_spec(gan_model, gan_loss)
else: # tf.estimator.ModeKeys.PREDICT
# Register hub modules for export
hub.register_module_for_export(generator_module, "generator")
hub.register_module_for_export(discriminator_module,
"discriminator")
estimator_spec = get_predict_estimator_spec(gan_model)
return estimator_spec
def _mdn_model_fn(features,
labels,
nchannels,
n_y,
dropout,
optimizer,
mode,
loss,
softplus=False):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
shift = tf.Variable(1., dtype=tf.float32, name='shift')
scale = tf.Variable(1., dtype=tf.float32, name='scale')
# Builds the neural network
# ! ny and nchannel need to be the same
cube_size = tf.shape(feature_layer)[1]
# print(cube_size)
chain = tfb.Chain([
tfp.bijectors.Affine(shift=shift, scale_identity_multiplier=scale),
tfb.Invert(
Squeeze3d(event_shape_in=[
cube_size, cube_size, cube_size, nchannels
])),
iRevNetsimple(name='layer1', h=h),
iRevNetsimple(name='layer1b', h=h),
iRevNetsimple(name='layer2', h=h),
iRevNetsimple(name='layer2b', h=h),
#tfb.Permute(np.arange(8)[::-1],axis=-1),
tfb.Permute(np.arange(8)[::-1], axis=-1),
iRevNetsimple(name='layer3', h=h),
iRevNetsimple(name='layer3b', h=h),
iRevNetsimple(name='layer4', h=h),
iRevNetsimple(name='layer4b', h=h),
tfb.Invert(
Squeeze3d(event_shape_in=[
cube_size // 2, cube_size // 2, cube_size // 2, nchannels *
8
])),
iRevNetsimple(name='layer5', h=h),
iRevNetsimple(name='layer5b', h=h),
iRevNetsimple(name='layer6', h=h),
iRevNetsimple(name='layer6b', h=h),
tfb.Permute(np.arange(64)[::-1], axis=-1),
iRevNetsimple(name='layer7', h=h),
iRevNetsimple(name='layer7b', h=h),
iRevNetsimple(name='layer8', h=h),
iRevNetsimple(name='layer8b', h=h),
tfb.Invert(
Squeeze3d(event_shape_in=[
cube_size // 4, cube_size // 4, cube_size // 4, nchannels *
64
])),
iRevNetsimple(name='layer9', h=h, kernel_size=1),
iRevNetsimple(name='layer9b', h=h, kernel_size=1),
iRevNetsimple(name='layer10', h=h, kernel_size=1),
iRevNetsimple(name='layer10b', h=h, kernel_size=1),
tfb.Permute(np.arange(64 * 8)[::-1], axis=-1),
iRevNetsimple(name='layer11', h=h, kernel_size=1),
iRevNetsimple(name='layer11b', h=h, kernel_size=1),
iRevNetsimple(name='layer12', h=h, kernel_size=1),
iRevNetsimple(name='layer12b', h=h, kernel_size=1),
Squeeze3d(event_shape_in=[
cube_size // 4, cube_size // 4, cube_size // 4, nchannels * 64
]),
iRevNetsimple(name='layer13', h=h),
iRevNetsimple(name='layer13b', h=h),
iRevNetsimple(name='layer14', h=h),
iRevNetsimple(name='layer14b', h=h),
tfb.Permute(np.arange(64)[::-1], axis=-1),
iRevNetsimple(name='layer15', h=h),
iRevNetsimple(name='layer15b', h=h),
iRevNetsimple(name='layer16', h=h),
iRevNetsimple(name='layer16b', h=h),
Squeeze3d(event_shape_in=[
cube_size // 2, cube_size // 2, cube_size // 2, nchannels * 8
]),
iRevNetsimple(name='layer17', h=h),
iRevNetsimple(name='layer17b', h=h),
iRevNetsimple(name='layer18', h=h),
iRevNetsimple(name='layer18b', h=h),
tfb.Permute(np.arange(8)[::-1], axis=-1),
iRevNetsimple(name='layer19', h=h),
iRevNetsimple(name='layer19b', h=h),
iRevNetsimple(name='layer20', h=h),
iRevNetsimple(name='layer20b', h=h),
Squeeze3d(
event_shape_in=[cube_size, cube_size, cube_size, nchannels])
])
bijection = chain
# Define the probabilistic layer
net = bijection.forward(feature_layer, name='lambda')
if softplus:
net = tf.nn.softplus(net, name='lambda')
dist = tfd.Poisson(net + 1e-3)
sample = tf.squeeze(dist.sample())
# loglik = dist.log_prob(obs_layer+1)
loglik = dist.log_prob(obs_layer)
#l2 = tf.losses.mean_squared_error(obs_layer, net)
l2 = (tf.square(tf.subtract(obs_layer, net)))
l1 = (tf.abs(tf.subtract(obs_layer, net)))
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik,
'lambda': net,
'l2': l2,
'l1': l1
})
#,'shift':shift, 'scale':scale})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
if loss == 'loglikelihood':
neg_log_likelihood = -predictions['loglikelihood']
elif loss == 'l2':
neg_log_likelihood = predictions['l2']
elif loss == 'l1':
neg_log_likelihood = predictions['l1']
else:
print('Loss not specified')
neg_log_likelihood = tf.reduce_sum(neg_log_likelihood, axis=-1)
neg_log_likelihood = tf.reduce_mean(neg_log_likelihood)
tf.losses.add_loss(neg_log_likelihood)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = list(np.array([1e4, 2e4, 4e4, 5e4, 6e4]).astype(int))
values = [1e-3, 5e-4, 1e-4, 5e-5, 1e-5, 1e-6]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('rate', learning_rate)
#tf.summary.scalar('shift', predictions['shift'])
#tf.summary.scalar('scale', predictions['scale'])
tf.summary.scalar('loss', neg_log_likelihood)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def body(self, features):
hparams = self.hparams
hparamsp = hparams.problem.get_hparams()
x = features['inputs']
cond = {k: features[k] for k in hparamsp.attributes}
# Load the encoder and decoder modules
encoder = hub.Module(hparams.encoder_module, trainable=False)
latent_shape = encoder.get_output_info_dict()['default'].get_shape(
)[1:]
latent_size = latent_shape[0].value * latent_shape[
1].value * latent_shape[2].value
code_shape = encoder.get_output_info_dict()['default'].get_shape()
code_shape = [
-1, code_shape[1].value, code_shape[2].value, code_shape[3].value
]
def get_flow(inputs, is_training=True):
y = tf.concat([
tf.expand_dims(inputs[k], axis=1) for k in hparamsp.attributes
],
axis=1)
y = tf.layers.batch_normalization(y,
name="y_norm",
training=is_training)
flow = self.normalizing_flow(y, latent_size)
return flow
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
# Export the latent flow alone
def flow_module_spec():
inputs = {
k: tf.placeholder(tf.float32, shape=[None])
for k in hparamsp.attributes
}
random_normal = tf.placeholder(tf.float32,
shape=[None, latent_size])
print(
f'\n \n \n !!!!!!!!!!! {random_normal} !!!!!!!! \n \n \n')
flow = get_flow(inputs, is_training=False)
samples = flow._bijector.forward(random_normal)
samples = tf.reshape(samples, code_shape)
hub.add_signature(inputs={
**inputs, 'random_normal': random_normal
},
outputs=samples)
flow_spec = hub.create_module_spec(flow_module_spec)
flow = hub.Module(flow_spec, name='flow_module')
hub.register_module_for_export(flow, "code_sampler")
cond['random_normal'] = tf.random_normal(
shape=[tf.shape(cond[hparamsp.attributes[0]])[0], latent_size])
samples = flow(cond)
return samples, {'loglikelihood': 0}
# Encode the input image
if hparams.encode_psf and 'psf' in features:
code = encoder({'input': x, 'psf': features['psf']})
else:
code = encoder(x)
with tf.variable_scope("flow_module"):
flow = get_flow(cond)
loglikelihood = flow.log_prob(tf.layers.flatten(code))
# This is the loglikelihood of a batch of images
tf.summary.scalar('loglikelihood', tf.reduce_mean(loglikelihood))
loss = -tf.reduce_mean(loglikelihood)
return code, {'training': loss}
def _mdn_model_fn(features, labels, n_y, n_mixture, dropout, optimizer, mode):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
def _module_fn():
"""
Function building the module
"""
feature_layer = tf.placeholder(
tf.float32,
shape=[None, None, None, None, nchannels],
name='input')
obs_layer = tf.placeholder(tf.float32,
shape=[None, None, None, None, n_y],
name='observations')
# Builds the neural network
net = slim.conv3d(feature_layer,
16,
5,
activation_fn=tf.nn.leaky_relu,
padding='valid')
#net = wide_resnet(feature_layer, 8, activation_fn=tf.nn.leaky_relu, is_training=is_training)
net = wide_resnet(net,
16,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = wide_resnet(net,
32,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = wide_resnet(net,
32,
activation_fn=tf.nn.leaky_relu,
keep_prob=dropout,
is_training=is_training)
net = slim.conv3d(net, 32, 3, activation_fn=tf.nn.tanh)
# Define the probabilistic layer
#out_rate = slim.conv3d(net, 1, 1, activation_fn=tf.nn.relu)
#out_rate = tf.math.add(out_rate, 1e-6, name='rate')
net = slim.conv3d(net, n_mixture * n_y, 1, activation_fn=tf.nn.relu)
cube_size = tf.shape(obs_layer)[1]
out_rate = tf.reshape(net, [-1, cube_size, cube_size, cube_size, n_y])
out_rate = tf.math.add(out_rate, 1e-6, name='rate')
pdf = tfd.Poisson(rate=out_rate)
# Define a function for sampling, and a function for estimating the log likelihood
sample = tf.squeeze(pdf.sample())
loglik = pdf.log_prob(obs_layer)
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik
})
# Create model and register module if necessary
spec = hub.create_module_spec(_module_fn)
module = hub.Module(spec, trainable=True)
if isinstance(features, dict):
predictions = module(features, as_dict=True)
else:
predictions = module({
'features': features,
'labels': labels
},
as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
hub.register_module_for_export(module, "likelihood")
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
loglik = predictions['loglikelihood']
# Compute and register loss function
neg_log_likelihood = -tf.reduce_sum(loglik, axis=-1)
neg_log_likelihood = tf.reduce_mean(neg_log_likelihood)
tf.losses.add_loss(neg_log_likelihood)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
global_step = tf.train.get_global_step()
boundaries = [10000, 20000, 30000, 40000]
values = [0.00001, 0.000005, 0.000001, 0.0000005, 0.0000001]
learning_rate = tf.train.piecewise_constant(
global_step, boundaries, values)
train_op = optimizer(learning_rate=learning_rate).minimize(
loss=total_loss, global_step=global_step)
tf.summary.scalar('loss', neg_log_likelihood)
tf.summary.scalar('rate', learning_rate)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": neg_log_likelihood}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def autoencoder_body(self, features):
""" Customized body function for autoencoders acting on continuous images.
This is based on tensor2tensor.models.research.AutoencoderBasic.body
and should be compatible with most derived classes.
The original autoencoder class relies on embedding the channels to a discrete
vocabulary and defines the loss on that vocab. It's cool and all, but here we
prefer expressing the reconstruction loss as an actual continuous likelihood
function.
"""
hparams = self.hparams
is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
output_activation = tf.nn.softplus if hparams.output_activation == 'softplus' else None
input_shape = [
None,
] + common_layers.shape_list(features["inputs"])[1:]
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
# In predict mode, we also define TensorFlow Hub modules for all pieces of
# the autoencoder
# First build encoder spec
def make_model_spec():
input_layer = tf.placeholder(tf.float32, shape=input_shape)
x = self.embed(tf.expand_dims(input_layer, -1))
x, encoder_layers = self.encoder(x)
b, b_loss = self.bottleneck(x)
hub.add_signature(inputs=input_layer, outputs=b)
def make_model_spec_psf():
input_layer = tf.placeholder(tf.float32, shape=input_shape)
psf_layer = tf.placeholder(tf.float32, shape=input_shape)
x = self.embed(tf.expand_dims(input_layer, -1))
# If we have access to the PSF, we add this information to the encoder
if hparams.encode_psf and 'psf' in features:
net_psf = tf.layers.conv2d(psf_layer,
hparams.hidden_size // 4,
5,
padding='same',
name="psf_embed_1")
net_psf = common_layers.layer_norm(net_psf, name="psf_norm")
x, encoder_layers = self.encoder(
tf.concat([x, net_psf], axis=-1))
else:
x, encoder_layers = self.encoder(x)
b, b_loss = self.bottleneck(x)
hub.add_signature(inputs={
'input': input_layer,
'psf': psf_layer
},
outputs=b)
spec = hub.create_module_spec(
make_model_spec_psf if hparams.encode_psf else make_model_spec,
drop_collections=['checkpoints'])
encoder = hub.Module(spec, name="encoder_module")
hub.register_module_for_export(encoder, "encoder")
if hparams.encode_psf:
code = encoder({
'input': features["inputs"],
'psf': features['psf']
})
else:
code = encoder(features["inputs"])
b_shape = [
None,
] + common_layers.shape_list(code)[1:]
res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
res_size = min(res_size, hparams.max_hidden_size)
# Second build decoder spec
def make_model_spec():
input_layer = tf.placeholder(tf.float32, shape=b_shape)
x = self.unbottleneck(input_layer, res_size)
x = self.decoder(x, None)
reconstr = tf.layers.dense(x,
self.num_channels,
name="autoencoder_final",
activation=output_activation)
hub.add_signature(inputs=input_layer, outputs=reconstr)
hub.attach_message(
"stamp_size",
tf.train.Int64List(value=[hparams.problem_hparams.img_len]))
hub.attach_message(
"pixel_size",
tf.train.FloatList(
value=[hparams.problem_hparams.pixel_scale]))
spec = hub.create_module_spec(make_model_spec,
drop_collections=['checkpoints'])
decoder = hub.Module(spec, name="decoder_module")
hub.register_module_for_export(decoder, "decoder")
reconstr = decoder(code)
return reconstr, {"bottleneck_loss": 0.0}
encoder_layers = None
self.is1d = hparams.sample_width == 1
if (hparams.mode != tf.estimator.ModeKeys.PREDICT
or self._encode_on_predict):
labels = features["targets_raw"]
labels_shape = common_layers.shape_list(labels)
shape = common_layers.shape_list(labels)
with tf.variable_scope('encoder_module'):
x = self.embed(tf.expand_dims(labels, -1))
if shape[2] == 1:
self.is1d = True
# Run encoder.
with tf.variable_scope('encoder_module'):
# If we have access to the PSF, we add this information to the encoder
if hparams.encode_psf and 'psf' in features:
net_psf = tf.layers.conv2d(features['psf'],
hparams.hidden_size // 4,
5,
padding='same',
name="psf_embed_1")
net_psf = common_layers.layer_norm(net_psf, name="psf_norm")
x, encoder_layers = self.encoder(
tf.concat([x, net_psf], axis=-1))
else:
x, encoder_layers = self.encoder(x)
# Bottleneck.
with tf.variable_scope('encoder_module'):
b, b_loss = self.bottleneck(x)
xb_loss = 0.0
b_shape = common_layers.shape_list(b)
self._cur_bottleneck_tensor = b
res_size = common_layers.shape_list(x)[-1]
with tf.variable_scope('decoder_module'):
b = self.unbottleneck(b, res_size)
if not is_training:
x = b
else:
l = 2**hparams.num_hidden_layers
warm_step = int(hparams.bottleneck_warmup_steps * 0.25 * l)
nomix_p = common_layers.inverse_lin_decay(warm_step) + 0.01
if common_layers.should_generate_summaries():
tf.summary.scalar("nomix_p_bottleneck", nomix_p)
rand = tf.random_uniform(common_layers.shape_list(x))
# This is the distance between b and x. Having this as loss helps learn
# the bottleneck function, but if we back-propagated to x it would be
# minimized by just setting x=0 and b=0 -- so we don't want too much
# of the influence of this, and we stop-gradient to not zero-out x.
x_stop = tf.stop_gradient(x)
xb_loss = tf.reduce_mean(
tf.reduce_sum(tf.squared_difference(x_stop, b), axis=-1))
# To prevent this loss from exploding we clip at 1, but anneal clipping.
clip_max = 1.0 / common_layers.inverse_exp_decay(warm_step,
min_value=0.001)
xb_clip = tf.maximum(tf.stop_gradient(xb_loss), clip_max)
xb_loss *= clip_max / xb_clip
x = tf.where(tf.less(rand, nomix_p), b, x)
else:
if self._cur_bottleneck_tensor is None:
b = self.sample()
else:
b = self._cur_bottleneck_tensor
self._cur_bottleneck_tensor = b
res_size = self.hparams.hidden_size * 2**self.hparams.num_hidden_layers
res_size = min(res_size, hparams.max_hidden_size)
with tf.variable_scope('decoder_module'):
x = self.unbottleneck(b, res_size)
# Run decoder.
with tf.variable_scope('decoder_module'):
x = self.decoder(x, encoder_layers)
# Cut to the right size and mix before returning.
res = x
if hparams.mode != tf.estimator.ModeKeys.PREDICT:
res = x[:, :shape[1], :shape[2], :]
with tf.variable_scope('decoder_module'):
reconstr = tf.layers.dense(res,
self.num_channels,
name="autoencoder_final",
activation=output_activation)
# Apply channel-wise convolution with the PSF if requested
# TODO: Handle multiple bands
if hparams.apply_psf and 'psf' in features:
if self.num_channels > 1:
raise NotImplementedError
rec_padded = tf.pad(
reconstr[:, :, :, 0],
[[0, 0], [0, int(hparams.psf_convolution_pad_factor * shape[1])],
[0, int(hparams.psf_convolution_pad_factor * shape[2])]])
psf_padded = tf.pad(
features['psf'][..., 0],
[[0, 0], [0, int(hparams.psf_convolution_pad_factor * shape[1])],
[0, int(hparams.psf_convolution_pad_factor * shape[2])]])
reconstr = tf.expand_dims(tf.spectral.irfft2d(
tf.spectral.rfft2d(rec_padded) *
tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded)), tf.complex64)),
axis=-1)
reconstr = reconstr[:, :shape[1], :shape[2], :]
# Losses.
losses = {
"bottleneck_extra": b_loss,
"bottleneck_l2": hparams.bottleneck_l2_factor * xb_loss
}
loglik = loglikelihood_fn(labels, reconstr, features, hparams)
targets_loss = tf.reduce_mean(-loglik)
tf.summary.scalar("negloglik", targets_loss)
tf.summary.scalar("bottleneck_loss", b_loss)
losses["training"] = targets_loss
logits = tf.reshape(reconstr, labels_shape)
image_summary("ae", reconstr)
image_summary("input", labels)
return logits, losses
def creates_estimator_model(images, labels, perms, num_classes, mode):
"""Creates EstimatorSpec for the patch based self supervised models.
Args:
images: images
labels: self supervised labels (class indices)
perms: patch permutations
num_classes: number of different permutations
mode: model's mode: training, eval or prediction
Returns:
EstimatorSpec
"""
print(' +++ Mode: %s, images: %s, labels: %s' % (mode, images, labels))
images = tf.reshape(images, shape=[-1] + images.get_shape().as_list()[-3:])
if mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL]:
with tf.variable_scope('module'):
image_fn = lambda: images
logits = apply_model(
image_fn=image_fn,
is_training=(mode == tf.estimator.ModeKeys.TRAIN),
num_outputs=num_classes,
perms=perms,
make_signature=False)
else:
input_shape = utils.str2intlist(
FLAGS.get_flag_value('serving_input_shape', 'None,None,None,3'))
image_fn = lambda: tf.placeholder( # pylint: disable=g-long-lambda
shape=input_shape,
dtype=tf.float32)
apply_model_function = functools.partial(apply_model,
image_fn=image_fn,
num_outputs=num_classes,
perms=perms,
make_signature=True)
tf_hub_module_spec = hub.create_module_spec(
apply_model_function, [(utils.TAGS_IS_TRAINING, {
'is_training': True
}), (set(), {
'is_training': False
})],
drop_collections=['summaries'])
tf_hub_module = hub.Module(tf_hub_module_spec,
trainable=False,
tags=set())
hub.register_module_for_export(tf_hub_module, export_name='module')
logits = tf_hub_module(images)
return make_estimator(mode, predictions=logits)
# build loss and accuracy
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,
logits=logits)
loss = tf.reduce_mean(loss)
eval_metrics = (
lambda labels, logits: { # pylint: disable=g-long-lambda
'accuracy':
tf.metrics.accuracy(labels=labels,
predictions=tf.argmax(logits, axis=-1))
},
[labels, logits])
return make_estimator(mode, loss, eval_metrics, logits)
