def bad_input_colocation_module_fn():
u = tf.add(42, 69, name="u")
with tf.colocate_with(u):
# Inputs must not reference other nodes for colocation.
x = tf.placeholder(tf.float32, name="x")
y = x + 1.0
hub.add_signature(inputs=x, outputs=y)
def unused_input_module_fn():
x = tf.placeholder(tf.int64)
y = tf.placeholder(tf.int64)
result = x*x
hub.add_signature(
inputs={"x": x, "unused": y},
outputs=result)
def table_lookup_module_fn(): x = tf.placeholder(dtype=tf.int64, name="x") keys = tf.constant([0, 1, 2], dtype=tf.int64) values = tf.constant(["index0", "hello", "world"]) tbl_init = tf.contrib.lookup.KeyValueTensorInitializer(keys, values) table = tf.contrib.lookup.HashTable(tbl_init, "UNK") hub.add_signature(inputs=x, outputs=table.lookup(x))
def while_module_fn():
"""Compute x^n with while_loop."""
x = tf.placeholder(dtype=tf.float32, name="x", shape=[])
n = tf.placeholder(dtype=tf.int32, name="n")
_, pow_x = tf.while_loop(
lambda i, ix: i < n, lambda i, ix: [tf.add(i, 1), ix * x],
[tf.constant(0), tf.constant(1.0)])
hub.add_signature(inputs={"x": x, "n": n}, outputs=pow_x)
def text_module_fn():
weights = tf.get_variable(
"weights", dtype=tf.float32, shape=[100, 10])
# initializer=tf.random_uniform_initializer())
text = tf.placeholder(tf.string, shape=[None])
hash_buckets = tf.string_to_hash_bucket_fast(text, weights.get_shape()[0])
embeddings = tf.gather(weights, hash_buckets)
hub.add_signature(inputs=text, outputs=embeddings)
def bad_state_colocation_module_fn():
u = tf.add(42, 69, name="u")
with tf.colocate_with(u):
# State-holding nodes must not reference other nodes for colocation.
v = tf.Variable(1.0, name="v")
x = tf.placeholder(dtype=tf.float32)
y = x + v
hub.add_signature(inputs=x, outputs=y)
def module_fn():
graph_def = tf.GraphDef()
with tf.gfile.GFile(frozen_graph_file, 'rb') as f:
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
graph = tf.get_default_graph()
x = graph.get_tensor_by_name('x:0')
y = graph.get_tensor_by_name('y:0')
hub.add_signature('default', inputs={'x': x}, outputs={'y': y})
def module_fn():
"""A module summing one normal and one partitioned variable."""
partitioner = tf.fixed_size_partitioner(partitions)
var_1 = tf.get_variable(
initializer=tf.ones(shape),
name="partitioned_variable",
partitioner=partitioner)
var_2 = tf.get_variable(
initializer=tf.ones(shape), name="normal_variable")
hub.add_signature(outputs=var_1 + var_2)
def stateful_module_fn_with_colocation():
v = tf.get_variable(
"var123", shape=[],
initializer=tf.constant_initializer(1.0),
use_resource=False)
v_value = v.value()
x = tf.placeholder(dtype=tf.float32, name="x")
with tf.colocate_with(v), tf.colocate_with(x):
y = tf.add(v_value, x, name="y")
hub.add_signature(inputs=x, outputs=y)
def multiple_signature_module_fn():
"""Stateful module with multiple signatures."""
weight = tf.Variable([3.0])
x_input = tf.placeholder(dtype=tf.float32)
x_output = tf.multiply(x_input, weight)
hub.add_signature("mul", inputs=x_input, outputs=x_output)
y_input = tf.placeholder(dtype=tf.float32)
y_output = tf.divide(y_input, weight)
hub.add_signature("div", inputs=y_input, outputs=y_output)
def module_with_variables():
tf.get_variable(
name="weights",
shape=[3],
initializer=tf.zeros_initializer())
tf.get_variable(
name="partition",
shape=[4],
initializer=tf.zeros_initializer(),
partitioner=tf.fixed_size_partitioner(3))
hub.add_signature(outputs=tf.constant(1.0))
def good_colocation_module_fn():
w = tf.Variable(42 + 69, name="w")
with tf.colocate_with(w):
# Colocation references among state nodes is ok.
v = tf.Variable(1.0, name="v")
assert v.op.colocation_groups() == [tf.compat.as_bytes("loc:@w")]
x = tf.placeholder(dtype=tf.float32, name="x")
with tf.colocate_with(x):
# Colocation references from other nodes to state nodes is ok.
y = x + v
assert sorted(y.op.colocation_groups()) == [tf.compat.as_bytes("loc:@x")]
hub.add_signature(inputs=x, outputs=y)
def nested_control_flow_module_fn():
"""Compute the sum of elements greater than 'a' with nested control flow."""
elems = tf.placeholder(
dtype=tf.float32, name="elems", shape=[None])
a = tf.placeholder(dtype=tf.float32, name="a")
def sum_above_a(acc, x):
return acc + tf.cond(x > a, lambda: x, lambda: 0.0)
hub.add_signature(
inputs={"elems": elems, "a": a},
outputs=tf.foldl(sum_above_a, elems, initializer=tf.constant(0.0)))
def layers_module_fn():
"""Module that exercises the use of layers."""
# This is a plain linear map Mx+b regularized by the sum of the squares
# of the coefficients in M and b.
x = tf.placeholder(dtype=tf.float32, shape=[None, 2], name="x")
# Cancels internal factor 0.5.
l2_reg = tf.contrib.layers.l2_regularizer(scale=2.0)
h = tf.layers.dense(
x, 2,
activation=None,
kernel_regularizer=l2_reg,
bias_regularizer=l2_reg)
hub.add_signature(inputs=x, outputs=h)
def module_with_duplicate_asset():
vocabulary_file = self.create_vocab_file("tokens2.txt", ["1", "2", "3"])
indices1 = tf.placeholder(dtype=tf.int64, name="indices1")
indices2 = tf.placeholder(dtype=tf.int64, name="indices2")
hub.add_signature(
inputs={
"indices_1": indices1,
"indices_2": indices2,
},
outputs={
"x": do_table_lookup(indices1, vocabulary_file),
"y": do_table_lookup(indices2, vocabulary_file),
})
def nested_cond_module_fn():
"""Computes relu(x) with nested conditionals."""
x = tf.placeholder(dtype=tf.float32, name="x", shape=[])
# pylint: disable=g-long-lambda
result = tf.cond(
0 < x,
lambda: tf.cond(3 < x,
lambda: tf.identity(x),
lambda: tf.multiply(x, 1.0)),
lambda: tf.cond(x < -3,
lambda: tf.constant(0.0),
lambda: tf.multiply(0.0, 1.0)))
# pylint: enable=g-long-lambda
hub.add_signature(inputs=x, outputs=result)
def batch_norm_module_fn(is_training):
"""Module that exercises batch normalization, incl. UPDATE_OPS."""
x = tf.placeholder(dtype=tf.float32, shape=[None, 1], name="x")
y = tf.contrib.layers.batch_norm(
x,
center=True,
scale=True,
is_training=is_training,
# Let the moving_mean (aggregated for normalization at serving time)
# decay quickly for more accurate values after few iterations.
decay=0.6,
# TODO(b/38416827): re-enable after the tests are updated.
fused=False)
hub.add_signature(inputs=x, outputs=y)
def hub_module_fn():
"""Creates the TF graph for the hub module."""
model_fn = t2t_model.T2TModel.make_estimator_model_fn(
model_name,
hparams,
decode_hparams=decode_hparams,
use_tpu=FLAGS.use_tpu)
features = problem.serving_input_fn(hparams).features
# we must do a copy of the features, as the model_fn can add additional
# entries there (like hyperparameter settings etc).
original_features = features.copy()
spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)
hub.add_signature(
inputs=original_features,
outputs=spec.export_outputs["serving_default"].outputs)
def module_fn():
"""Spec function for a token embedding module."""
tokens = tf.placeholder(shape=[None], dtype=tf.string, name="tokens")
embeddings_var = tf.get_variable(
initializer=tf.zeros([vocab_size + num_oov_buckets, embeddings_dim]),
name=EMBEDDINGS_VAR_NAME,
dtype=tf.float32)
lookup_table = tf.contrib.lookup.index_table_from_file(
vocabulary_file=vocabulary_file,
num_oov_buckets=num_oov_buckets,
)
ids = lookup_table.lookup(tokens)
combined_embedding = tf.nn.embedding_lookup(params=embeddings_var, ids=ids)
hub.add_signature("default", {"tokens": tokens},
{"default": combined_embedding})
def module_fn_with_preprocessing():
"""Spec function for a full-text embedding module with preprocessing."""
sentences = tf.placeholder(shape=[None], dtype=tf.string, name="sentences")
# Perform a minimalistic text preprocessing by removing punctuation and
# splitting on spaces.
normalized_sentences = tf.regex_replace(
input=sentences, pattern=r"\pP", rewrite="")
tokens = tf.string_split(normalized_sentences, " ")
# In case some of the input sentences are empty before or after
# normalization, we will end up with empty rows. We do however want to
# return embedding for every row, so we have to fill in the empty rows with
# a default.
tokens, _ = tf.sparse_fill_empty_rows(tokens, "")
# In case all of the input sentences are empty before or after
# normalization, we will end up with a SparseTensor with shape [?, 0]. After
# filling in the empty rows we must ensure the shape is set properly to
# [?, 1].
tokens = tf.sparse_reset_shape(tokens)
embeddings_var = tf.get_variable(
initializer=tf.zeros([vocab_size + num_oov_buckets, embeddings_dim]),
name=EMBEDDINGS_VAR_NAME,
dtype=tf.float32)
lookup_table = tf.contrib.lookup.index_table_from_file(
vocabulary_file=vocabulary_file,
num_oov_buckets=num_oov_buckets,
)
sparse_ids = tf.SparseTensor(
indices=tokens.indices,
values=lookup_table.lookup(tokens.values),
dense_shape=tokens.dense_shape)
combined_embedding = tf.nn.embedding_lookup_sparse(
params=embeddings_var,
sp_ids=sparse_ids,
sp_weights=None,
combiner="sqrtn")
hub.add_signature("default", {"sentences": sentences},
{"default": combined_embedding})
def image_module_fn():
"""Maps 1x2 images to sums of each color channel."""
images = tf.placeholder(dtype=tf.float32, shape=[None, 1, 2, 3])
sum_channels = tf.reduce_sum(images, axis=[1, 2])
hub.add_signature(inputs={"images": images}, outputs=sum_channels)
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 stateful_rv_module_fn():
r = tf.get_variable(
"rv_var123", shape=[],
initializer=tf.constant_initializer(10.0),
use_resource=True)
hub.add_signature(outputs=r.value())
def module_fn(training):
"""Hub module definition."""
image_input = tf.placeholder(
shape=[None] + self._observations_shape[1:], dtype=tf.float32)
if self._extra_shape:
inputs_extra = tf.placeholder(
shape=[None, self._extra_shape[-1]], dtype=tf.float32)
else:
inputs_extra = None
noise_shape = ([None] * (len(self._latent_shape) - 1)
+ [self._latent_shape[-1]])
noise_input = tf.placeholder(shape=noise_shape, dtype=tf.float32)
with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
latents, endpoints_encoder = self._encoder(image_input, inputs_extra,
training=training)
latent_means = latents[Ellipsis, :self._latent_shape[-1]]
latent_log_sigmas = latents[Ellipsis, self._latent_shape[-1]:]
with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
outputs, endpoints_decoder = self._decoder(latent_means, inputs_extra,
training=training)
outputs_channels = outputs.get_shape().as_list()[-1] // 2
if self._output_distribution == "gaussian":
reconstructions = outputs[Ellipsis, :outputs_channels]
reconstruction_log_sigmas = outputs[Ellipsis, outputs_channels:]
elif self._output_distribution == "bernoulli":
reconstructions = outputs
reconstruction_log_sigmas = outputs
recon_inputs = {"image": image_input}
if self._extra_shape:
recon_inputs.update({"inputs_extra": inputs_extra})
hub.add_signature(inputs=recon_inputs, outputs=reconstructions,
name="reconstruct")
outputs = {"reconstructions": reconstructions,
"reconstruction_log_sigmas": reconstruction_log_sigmas,
"latent_means": latent_means,
"latent_log_sigmas": latent_log_sigmas}
outputs.update(endpoints_encoder)
outputs.update(endpoints_decoder)
hub.add_signature(inputs=recon_inputs, outputs=outputs,
name="reconstruct_with_extras")
with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
outputs, endpoints_decoder = self._decoder(noise_input, inputs_extra,
training=training)
if self._output_distribution == "gaussian":
image_samples = outputs[Ellipsis, :outputs_channels]
image_samples_log_sigmas = outputs[Ellipsis, outputs_channels:]
elif self._output_distribution == "bernoulli":
image_samples = outputs
image_samples_log_sigmas = outputs
sample_inputs = {"noise": noise_input}
if self._extra_shape:
sample_inputs.update({"inputs_extra": inputs_extra})
hub.add_signature(inputs=sample_inputs, outputs=image_samples,
name="sample")
outputs = {"image_samples": image_samples,
"image_samples_log_sigmas": image_samples_log_sigmas}
outputs.update(endpoints_decoder)
hub.add_signature(inputs=sample_inputs, outputs=outputs,
name="sample_with_extras")
def module_fn(is_training):
"""Module function."""
input_ids = tf.placeholder(tf.int32, [None, None], "input_ids")
input_mask = tf.placeholder(tf.int32, [None, None], "input_mask")
segment_ids = tf.placeholder(tf.int32, [None, None], "segment_ids")
mlm_positions = tf.placeholder(tf.int32, [None, None], "mlm_positions")
albert_config_path = os.path.join(FLAGS.albert_directory,
"albert_config.json")
albert_config = modeling.AlbertConfig.from_json_file(albert_config_path)
model = modeling.AlbertModel(config=albert_config,
is_training=is_training,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=segment_ids,
use_one_hot_embeddings=False)
mlm_logits = get_mlm_logits(model, albert_config, mlm_positions)
vocab_model_path = os.path.join(FLAGS.albert_directory, "30k-clean.model")
vocab_file_path = os.path.join(FLAGS.albert_directory, "30k-clean.vocab")
config_file = tf.constant(value=albert_config_path,
dtype=tf.string,
name="config_file")
vocab_model = tf.constant(value=vocab_model_path,
dtype=tf.string,
name="vocab_model")
# This is only for visualization purpose.
vocab_file = tf.constant(value=vocab_file_path,
dtype=tf.string,
name="vocab_file")
# By adding `config_file, vocab_model and vocab_file`
# to the ASSET_FILEPATHS collection, TF-Hub will
# rewrite this tensor so that this asset is portable.
tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, config_file)
tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, vocab_model)
tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, vocab_file)
hub.add_signature(name="tokens",
inputs=dict(input_ids=input_ids,
input_mask=input_mask,
segment_ids=segment_ids),
outputs=dict(sequence_output=model.get_sequence_output(),
pooled_output=model.get_pooled_output()))
hub.add_signature(name="mlm",
inputs=dict(input_ids=input_ids,
input_mask=input_mask,
segment_ids=segment_ids,
mlm_positions=mlm_positions),
outputs=dict(sequence_output=model.get_sequence_output(),
pooled_output=model.get_pooled_output(),
mlm_logits=mlm_logits))
hub.add_signature(name="tokenization_info",
inputs={},
outputs=dict(vocab_file=vocab_model,
do_lower_case=tf.constant(
FLAGS.do_lower_case)))
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]))
def plus_one():
x = tf.compat.v1.sparse.placeholder(dtype=tf.float32, name='x')
y = tf.identity(tf.SparseTensor(x.indices, x.values + 1,
x.dense_shape))
hub.add_signature(inputs=x, outputs=y)
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, 0.001) * 1000
#
# Builds the neural network
if pad == 0:
net = slim.conv3d(feature_layer,
16,
5,
activation_fn=tf.nn.leaky_relu,
padding='same')
elif pad == 2:
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)
if distribution == 'logistic':
net = slim.conv3d(net, 32, 3, activation_fn=tf.nn.tanh)
else:
net = slim.conv3d(net, 32, 3, activation_fn=tf.nn.leaky_relu)
#Predicted mask
masknet = slim.conv3d(net, 8, 1, activation_fn=tf.nn.leaky_relu)
out_mask = slim.conv3d(masknet, 1, 1, activation_fn=None)
pred_mask = tf.nn.sigmoid(out_mask)
# 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)
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())
rawsample = mixture_dist.sample()
sample = rawsample * pred_mask
rawloglik = mixture_dist.log_prob(obs_layer)
print(rawloglik)
print(out_mask)
print(mask_layer)
#loss1 = - rawloglik* out_mask #This can be constant mask as well if we use mask_layer instead
if masktype == 'constant': loss1 = -rawloglik * mask_layer
elif masktype == 'vary': loss1 = -rawloglik * pred_mask
loss2 = tf.nn.sigmoid_cross_entropy_with_logits(logits=out_mask,
labels=mask_layer)
loglik = -(loss1 + loss2)
hub.add_signature(inputs={
'features': feature_layer,
'labels': obs_layer
},
outputs={
'sample': sample,
'loglikelihood': loglik,
'loc': loc,
'scale': scale,
'logits': logits,
'rawsample': rawsample,
'pred_mask': pred_mask,
'out_mask': out_mask,
'rawloglik': rawloglik,
'loss1': loss1,
'loss2': loss2
})
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
})
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
})
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
})
def stateful_non_rv_module_fn():
v = tf.get_variable(
"var123", shape=[],
initializer=tf.constant_initializer(10.0),
use_resource=False)
hub.add_signature(outputs=v.value())
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)
def double_module_fn(): w = tf.Variable([2.0, 4.0]) x = tf.compat.v1.placeholder(dtype=tf.float32) hub.add_signature(inputs=x, outputs=x*w)
def module_fn():
inputs = tf.placeholder(dtype=tf.float32, shape=[None, 50])
layer1 = tf.layers.dense(inputs, 200)
layer2 = tf.layers.dense(layer1, 100)
outputs = dict(default=layer2, hidden_activations=layer1)
hub.add_signature(inputs=inputs, outputs=outputs)
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})
def module_fn():
"""Spec function for a token embedding module."""
# init
_bos_id = 256
_eos_id = 257
_bow_id = 258
_eow_id = 259
_pad_id = 260
_max_word_length = 50
_parallel_iterations = 10
_max_batch_size = 1024
id_dtype = tf.int32
id_nptype = np.int32
max_word_length = tf.constant(_max_word_length,
dtype=id_dtype,
name='max_word_length')
version = tf.constant('from_dp_1', dtype=tf.string, name='version')
# the charcter representation of the begin/end of sentence characters
def _make_bos_eos(c):
r = np.zeros([_max_word_length], dtype=id_nptype)
r[:] = _pad_id
r[0] = _bow_id
r[1] = c
r[2] = _eow_id
return tf.constant(r, dtype=id_dtype)
bos_ids = _make_bos_eos(_bos_id)
eos_ids = _make_bos_eos(_eos_id)
def token2ids(token):
with tf.name_scope("token2ids_preprocessor"):
char_ids = tf.decode_raw(token,
tf.uint8,
name='decode_raw2get_char_ids')
char_ids = tf.cast(char_ids, tf.int32, name='cast2int_token')
char_ids = tf.strided_slice(char_ids, [0],
[max_word_length - 2], [1],
name='slice2resized_token')
ids_num = tf.shape(char_ids)[0]
fill_ids_num = (_max_word_length - 2) - ids_num
pads = tf.fill([fill_ids_num], _pad_id)
bow_token_eow_pads = tf.concat(
[[_bow_id], char_ids, [_eow_id], pads],
0,
name='concat2bow_token_eow_pads')
return bow_token_eow_pads
def sentence_tagging_and_padding(sen_dim):
with tf.name_scope("sentence_tagging_and_padding_preprocessor"):
sen = sen_dim[0]
dim = sen_dim[1]
extra_dim = tf.shape(sen)[0] - dim
sen = tf.slice(sen, [0, 0], [dim, max_word_length],
name='slice2sen')
bos_sen_eos = tf.concat([[bos_ids], sen, [eos_ids]],
0,
name='concat2bos_sen_eos')
bos_sen_eos_plus_one = bos_sen_eos + 1
bos_sen_eos_pads = tf.pad(bos_sen_eos_plus_one,
[[0, extra_dim], [0, 0]],
"CONSTANT",
name='pad2bos_sen_eos_pads')
return bos_sen_eos_pads
# Input placeholders to the biLM.
tokens = tf.placeholder(shape=(None, None),
dtype=tf.string,
name='ph2tokens')
sequence_len = tf.placeholder(shape=(None, ),
dtype=tf.int32,
name='ph2sequence_len')
tok_shape = tf.shape(tokens)
line_tokens = tf.reshape(tokens,
shape=[-1],
name='reshape2line_tokens')
with tf.device('/cpu:0'):
tok_ids = tf.map_fn(token2ids,
line_tokens,
dtype=tf.int32,
back_prop=False,
parallel_iterations=_parallel_iterations,
name='map_fn2get_tok_ids')
tok_ids = tf.reshape(tok_ids, [tok_shape[0], tok_shape[1], -1],
name='reshape2tok_ids')
with tf.device('/cpu:0'):
sen_ids = tf.map_fn(sentence_tagging_and_padding,
(tok_ids, sequence_len),
dtype=tf.int32,
back_prop=False,
parallel_iterations=_parallel_iterations,
name='map_fn2get_sen_ids')
# Build the biLM graph.
bilm = BidirectionalLanguageModel(options,
str(weight_file),
max_batch_size=_max_batch_size)
embeddings_op = bilm(sen_ids)
# Get an op to compute ELMo (weighted average of the internal biLM layers)
elmo_output = weight_layers('elmo_output', embeddings_op, l2_coef=0.0)
weighted_op = elmo_output['weighted_op']
mean_op = elmo_output['mean_op']
word_emb = elmo_output['word_emb']
lstm_outputs1 = elmo_output['lstm_outputs1']
lstm_outputs2 = elmo_output['lstm_outputs2']
hub.add_signature("tokens", {
"tokens": tokens,
"sequence_len": sequence_len
}, {
"elmo": weighted_op,
"default": mean_op,
"word_emb": word_emb,
"lstm_outputs1": lstm_outputs1,
"lstm_outputs2": lstm_outputs2,
"version": version
})
# #########################Next signature############################# #
# Input placeholders to the biLM.
def_strings = tf.placeholder(shape=(None), dtype=tf.string)
def_tokens_sparse = tf.string_split(def_strings)
def_tokens_dense = tf.sparse_to_dense(
sparse_indices=def_tokens_sparse.indices,
output_shape=def_tokens_sparse.dense_shape,
sparse_values=def_tokens_sparse.values,
default_value='')
def_mask = tf.not_equal(def_tokens_dense, '')
def_int_mask = tf.cast(def_mask, dtype=tf.int32)
def_sequence_len = tf.reduce_sum(def_int_mask, axis=-1)
def_tok_shape = tf.shape(def_tokens_dense)
def_line_tokens = tf.reshape(def_tokens_dense,
shape=[-1],
name='reshape2line_tokens')
with tf.device('/cpu:0'):
def_tok_ids = tf.map_fn(token2ids,
def_line_tokens,
dtype=tf.int32,
back_prop=False,
parallel_iterations=_parallel_iterations,
name='map_fn2get_tok_ids')
def_tok_ids = tf.reshape(def_tok_ids,
[def_tok_shape[0], def_tok_shape[1], -1],
name='reshape2tok_ids')
with tf.device('/cpu:0'):
def_sen_ids = tf.map_fn(sentence_tagging_and_padding,
(def_tok_ids, def_sequence_len),
dtype=tf.int32,
back_prop=False,
parallel_iterations=_parallel_iterations,
name='map_fn2get_sen_ids')
# Get ops to compute the LM embeddings.
def_embeddings_op = bilm(def_sen_ids)
# Get an op to compute ELMo (weighted average of the internal biLM layers)
def_elmo_output = weight_layers('elmo_output',
def_embeddings_op,
l2_coef=0.0,
reuse=True)
def_weighted_op = def_elmo_output['weighted_op']
def_mean_op = def_elmo_output['mean_op']
def_word_emb = def_elmo_output['word_emb']
def_lstm_outputs1 = def_elmo_output['lstm_outputs1']
def_lstm_outputs2 = def_elmo_output['lstm_outputs2']
hub.add_signature("default", {"strings": def_strings}, {
"elmo": def_weighted_op,
"default": def_mean_op,
"word_emb": def_word_emb,
"lstm_outputs1": def_lstm_outputs1,
"lstm_outputs2": def_lstm_outputs2,
"version": version
})
def half_plus_two():
a = tf.get_variable("a", shape=[])
b = tf.get_variable("b", shape=[])
x = tf.placeholder(tf.float32)
y = a * x + b
hub.add_signature(inputs=x, outputs=y)
def _stateless_module_fn(self): """Simple module that squares an input.""" x = tf.placeholder(tf.int64) y = x*x hub.add_signature(inputs=x, outputs=y)
def _module_fn(is_training):
"""A module_fn for use with hub.create_module_spec().
Args:
is_training: a boolean, passed to the config.network_fn.
This is meant to control whether batch norm, dropout etc. are built
in training or inference mode for this graph version.
Raises:
ValueError: if network_fn outputs are not as expected.
"""
# Set up the module input, and attach an ImageModuleInfo about it.
with tf.name_scope('hub_input'):
default_size = (hparams.image_size, ) * 2
image_module_info = hub.ImageModuleInfo()
size_info = image_module_info.default_image_size
size_info.height, size_info.width = default_size
# TODO(b/72731449): Support variable input size.
shape = (None, ) + default_size + (3, )
images = tf.placeholder(dtype=tf.float32,
shape=shape,
name='images')
hub.attach_image_module_info(image_module_info)
# The input is expected to have RGB color values in the range [0,1]
# and gets converted for AmoebaNet to the Inception-style range [-1,+1].
scaled_images = tf.multiply(images, 2.0)
scaled_images = tf.subtract(scaled_images, 1.0)
# Build the net.
logits, end_points = model_builder.build_network(
scaled_images, num_classes, is_training, hparams)
with tf.name_scope('hub_output'):
# Extract the feature_vectors output.
try:
feature_vectors = end_points['global_pool']
except KeyError:
tf.logging.error('Valid keys of end_points are:',
', '.join(end_points))
raise
with tf.name_scope('feature_vector'):
if feature_vectors.shape.ndims != 2:
raise ValueError(
'Wrong rank (expected 2 after squeeze) '
'in feature_vectors:', feature_vectors)
# Extract the logits output (if applicable).
if num_classes:
with tf.name_scope('classification'):
if logits.shape.ndims != 2:
raise ValueError('Wrong rank (expected 2) in logits:',
logits)
# Add named signatures.
hub.add_signature('image_feature_vector', dict(images=images),
dict(end_points, default=feature_vectors))
if num_classes:
hub.add_signature('image_classification', dict(images=images),
dict(end_points, default=logits))
# Add the default signature.
if num_classes:
hub.add_signature('default', dict(images=images),
dict(default=logits))
else:
hub.add_signature('default', dict(images=images),
dict(default=feature_vectors))
def _module_fn(self, model, batch_size):
"""Module Function to create a TF Hub module spec.
Args:
model: `tf.estimator.ModeKeys` value.
batch_size: batch size.
"""
if model not in {"gen", "disc"}:
raise ValueError("Model {} not support in module_fn()".format(model))
placeholder_fn = tf.placeholder if batch_size is None else tf.zeros
is_training = False
inputs = {}
y = None
if model == "gen":
inputs["z"] = placeholder_fn(
shape=(batch_size, self._z_dim),
dtype=tf.float32,
name="z_for_eval")
elif model == "disc":
inputs["images"] = placeholder_fn(
shape=[batch_size] + list(self._dataset.image_shape),
dtype=tf.float32,
name="images_for_eval")
if self.conditional:
inputs["labels"] = placeholder_fn(
shape=(batch_size,),
dtype=tf.int32,
name="labels_for_eval")
y = self._get_one_hot_labels(inputs["labels"])
else:
y = None
logging.info("Creating module for model %s with inputs %s and y=%s",
model, inputs, y)
outputs = {}
if model == "disc":
outputs["prediction"], _, _ = self.discriminator(
inputs["images"], y=y, is_training=is_training)
else:
z = inputs["z"]
generated = self.generator(z=z, y=y, is_training=is_training)
if self._g_use_ema and not is_training:
g_vars = [var for var in tf.trainable_variables()
if "generator" in var.name]
ema = tf.train.ExponentialMovingAverage(decay=self._ema_decay)
# Create the variables that will be loaded from the checkpoint.
ema.apply(g_vars)
def ema_getter(getter, name, *args, **kwargs):
var = getter(name, *args, **kwargs)
ema_var = ema.average(var)
if ema_var is None:
var_names_without_ema = {"u_var", "accu_mean", "accu_variance",
"accu_counter", "update_accus"}
if name.split("/")[-1] not in var_names_without_ema:
logging.warning("Could not find EMA variable for %s.", name)
return var
return ema_var
with tf.variable_scope("", values=[z, y], reuse=True,
custom_getter=ema_getter):
generated = self.generator(z, y=y, is_training=is_training)
outputs["generated"] = generated
hub.add_signature(inputs=inputs, outputs=outputs)
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)
def _build_graph(self, batch_dim: Optional[int] = None):
"""Builds the TensorFlow graph for evaluating the functional.
Args:
batch_dim: the batch dimension of the grid to use in the model. Default:
None (determine at runtime). This should only be set if building a model
in order to export and ahead-of-time compile it into a standalone
library.
"""
self._functional = hub.Module(spec=self._model_path)
grid_coords = tf.placeholder(
tf.float32, shape=[batch_dim, 3], name='grid_coords')
grid_weights = tf.placeholder(
tf.float32, shape=[batch_dim], name='grid_weights')
# Density information.
rho_a = tf.placeholder(tf.float32, shape=[6, batch_dim], name='rho_a')
rho_b = tf.placeholder(tf.float32, shape=[6, batch_dim], name='rho_b')
# Split into corresponding terms.
rho_only_a, grad_a_x, grad_a_y, grad_a_z, _, tau_a = tf.unstack(
rho_a, axis=0)
rho_only_b, grad_b_x, grad_b_y, grad_b_z, _, tau_b = tf.unstack(
rho_b, axis=0)
# Evaluate |\del \rho|^2 for each spin density and for the total density.
norm_grad_a = (grad_a_x**2 + grad_a_y**2 + grad_a_z**2)
norm_grad_b = (grad_b_x**2 + grad_b_y**2 + grad_b_z**2)
grad_x = grad_a_x + grad_b_x
grad_y = grad_a_y + grad_b_y
grad_z = grad_a_z + grad_b_z
norm_grad = (grad_x**2 + grad_y**2 + grad_z**2)
# The local Hartree-Fock energy densities at each grid point for the alpha-
# and beta-spin densities for each value of omega.
# Note an omega of 0 indicates no screening of the Coulomb potential.
hfxa = tf.placeholder(
tf.float32, shape=[batch_dim, len(self._omega_values)], name='hfxa')
hfxb = tf.placeholder(
tf.float32, shape=[batch_dim, len(self._omega_values)], name='hfxb')
# Make all features 2D arrays on input for ease of handling inside the
# functional.
features = {
'grid_coords': grid_coords,
'grid_weights': tf.expand_dims(grid_weights, 1),
'rho_a': tf.expand_dims(rho_only_a, 1),
'rho_b': tf.expand_dims(rho_only_b, 1),
'tau_a': tf.expand_dims(tau_a, 1),
'tau_b': tf.expand_dims(tau_b, 1),
'norm_grad_rho_a': tf.expand_dims(norm_grad_a, 1),
'norm_grad_rho_b': tf.expand_dims(norm_grad_b, 1),
'norm_grad_rho': tf.expand_dims(norm_grad, 1),
'hfxa': hfxa,
'hfxb': hfxb,
}
tensor_dict = {f'tensor_dict${k}': v for k, v in features.items()}
predictions = self._functional(tensor_dict, as_dict=True)
local_xc = predictions['grid_contribution']
weighted_local_xc = local_xc * grid_weights
unweighted_xc = tf.reduce_sum(local_xc, axis=0)
xc = tf.reduce_sum(weighted_local_xc, axis=0)
# The potential is the local exchange correlation divided by the
# total density. Add a small constant to deal with zero density.
self._vxc = local_xc / (rho_only_a + rho_only_b + 1E-12)
# The derivatives of the exchange-correlation (XC) energy with respect to
# input features. PySCF weights the (standard) derivatives by the grid
# weights, so we need to compute this with respect to the unweighted sum
# over grid points.
self._vrho = tf.gradients(
unweighted_xc, [features['rho_a'], features['rho_b']],
name='GRAD_RHO',
unconnected_gradients=tf.UnconnectedGradients.ZERO)
self._vsigma = tf.gradients(
unweighted_xc, [
features['norm_grad_rho_a'], features['norm_grad_rho_b'],
features['norm_grad_rho']
],
name='GRAD_SIGMA',
unconnected_gradients=tf.UnconnectedGradients.ZERO)
self._vtau = tf.gradients(
unweighted_xc, [features['tau_a'], features['tau_b']],
name='GRAD_TAU',
unconnected_gradients=tf.UnconnectedGradients.ZERO)
# Standard meta-GGAs do not have a dependency on local HF, so we need to
# compute the contribution to the Fock matrix ourselves. Just use the
# weighted XC energy to avoid having to weight this later.
self._vhf = tf.gradients(
xc, [features['hfxa'], features['hfxb']],
name='GRAD_HFX',
unconnected_gradients=tf.UnconnectedGradients.ZERO)
self._placeholders = FunctionalInputs(
rho_a=rho_a,
rho_b=rho_b,
hfx_a=hfxa,
hfx_b=hfxb,
grid_coords=grid_coords,
grid_weights=grid_weights)
outputs = {
'vxc': self._vxc,
'vrho': tf.stack(self._vrho),
'vsigma': tf.stack(self._vsigma),
'vtau': tf.stack(self._vtau),
'vhf': tf.stack(self._vhf),
}
# Create the signature for TF-Hub, including both the energy and functional
# derivatives.
# This is a no-op if _build_graph is called outside of
# hub.create_module_spec.
hub.add_signature(
inputs=attr.asdict(self._placeholders), outputs=outputs)
def stateful_module_fn():
v = tf.get_variable(
"var123", shape=[3],
initializer=tf.constant_initializer([1.0, 2.0, 3.0]))
hub.add_signature(outputs=v.value())
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)
def assets_module_fn(): indices = tf.placeholder(dtype=tf.int64, name="indices") outputs = do_table_lookup(indices, vocabulary_file) hub.add_signature(inputs=indices, outputs=outputs)
def _stateless_module_fn():
"""Simple module that squares an input."""
x = tf.compat.v1.placeholder(tf.int64)
y = x * x
hub.add_signature(inputs=x, outputs=y)
def consumer_module_fn(): indices = tf.placeholder(dtype=tf.int64, name="indices") inner_module = hub.Module(exported_hub_module) inner_module_output = inner_module(indices) output = tf.identity(inner_module_output) hub.add_signature(inputs=indices, outputs=output)
def module_fn():
"""Spec function for a token embedding module."""
# init
_bos_id = 256
_eos_id = 257
_bow_id = 258
_eow_id = 259
_pad_id = 260
_max_word_length = 50
_parallel_iterations = 10
_max_batch_size = 1024
id_dtype = tf.int32
id_nptype = np.int32
max_word_length = tf.constant(_max_word_length, dtype=id_dtype, name='max_word_length')
version = tf.constant('from_dp_1', dtype=tf.string, name='version')
# the charcter representation of the begin/end of sentence characters
def _make_bos_eos(c):
r = np.zeros([_max_word_length], dtype=id_nptype)
r[:] = _pad_id
r[0] = _bow_id
r[1] = c
r[2] = _eow_id
return tf.constant(r, dtype=id_dtype)
bos_ids = _make_bos_eos(_bos_id)
eos_ids = _make_bos_eos(_eos_id)
def token2ids(token):
with tf.name_scope("token2ids_preprocessor"):
char_ids = tf.decode_raw(token, tf.uint8, name='decode_raw2get_char_ids')
char_ids = tf.cast(char_ids, tf.int32, name='cast2int_token')
char_ids = tf.strided_slice(char_ids, [0], [max_word_length - 2],
[1], name='slice2resized_token')
ids_num = tf.shape(char_ids)[0]
fill_ids_num = (_max_word_length - 2) - ids_num
pads = tf.fill([fill_ids_num], _pad_id)
bow_token_eow_pads = tf.concat([[_bow_id], char_ids, [_eow_id], pads],
0, name='concat2bow_token_eow_pads')
return bow_token_eow_pads
def sentence_tagging_and_padding(sen_dim):
with tf.name_scope("sentence_tagging_and_padding_preprocessor"):
sen = sen_dim[0]
dim = sen_dim[1]
extra_dim = tf.shape(sen)[0] - dim
sen = tf.slice(sen, [0, 0], [dim, max_word_length], name='slice2sen')
bos_sen_eos = tf.concat([[bos_ids], sen, [eos_ids]], 0, name='concat2bos_sen_eos')
bos_sen_eos_plus_one = bos_sen_eos + 1
bos_sen_eos_pads = tf.pad(bos_sen_eos_plus_one, [[0, extra_dim], [0, 0]],
"CONSTANT", name='pad2bos_sen_eos_pads')
return bos_sen_eos_pads
# Input placeholders to the biLM.
tokens = tf.placeholder(shape=(None, None), dtype=tf.string, name='ph2tokens')
sequence_len = tf.placeholder(shape=(None, ), dtype=tf.int32, name='ph2sequence_len')
tok_shape = tf.shape(tokens)
line_tokens = tf.reshape(tokens, shape=[-1], name='reshape2line_tokens')
with tf.device('/cpu:0'):
tok_ids = tf.map_fn(
token2ids,
line_tokens,
dtype=tf.int32, back_prop=False, parallel_iterations=_parallel_iterations,
name='map_fn2get_tok_ids')
tok_ids = tf.reshape(tok_ids, [tok_shape[0], tok_shape[1], -1], name='reshape2tok_ids')
with tf.device('/cpu:0'):
sen_ids = tf.map_fn(
sentence_tagging_and_padding,
(tok_ids, sequence_len),
dtype=tf.int32, back_prop=False, parallel_iterations=_parallel_iterations,
name='map_fn2get_sen_ids')
# Build the biLM graph.
bilm = BidirectionalLanguageModel(options, str(weight_file),
max_batch_size=_max_batch_size)
embeddings_op = bilm(sen_ids)
# Get an op to compute ELMo (weighted average of the internal biLM layers)
elmo_output = weight_layers('elmo_output', embeddings_op, l2_coef=0.0)
weighted_op = elmo_output['weighted_op']
mean_op = elmo_output['mean_op']
word_emb = elmo_output['word_emb']
lstm_outputs1 = elmo_output['lstm_outputs1']
lstm_outputs2 = elmo_output['lstm_outputs2']
hub.add_signature("tokens", {"tokens": tokens, "sequence_len": sequence_len},
{"elmo": weighted_op,
"default": mean_op,
"word_emb": word_emb,
"lstm_outputs1": lstm_outputs1,
"lstm_outputs2": lstm_outputs2,
"version": version})
# #########################Next signature############################# #
# Input placeholders to the biLM.
def_strings = tf.placeholder(shape=(None), dtype=tf.string)
def_tokens_sparse = tf.string_split(def_strings)
def_tokens_dense = tf.sparse_to_dense(sparse_indices=def_tokens_sparse.indices,
output_shape=def_tokens_sparse.dense_shape,
sparse_values=def_tokens_sparse.values,
default_value=''
)
def_mask = tf.not_equal(def_tokens_dense, '')
def_int_mask = tf.cast(def_mask, dtype=tf.int32)
def_sequence_len = tf.reduce_sum(def_int_mask, axis=-1)
def_tok_shape = tf.shape(def_tokens_dense)
def_line_tokens = tf.reshape(def_tokens_dense, shape=[-1], name='reshape2line_tokens')
with tf.device('/cpu:0'):
def_tok_ids = tf.map_fn(
token2ids,
def_line_tokens,
dtype=tf.int32, back_prop=False, parallel_iterations=_parallel_iterations,
name='map_fn2get_tok_ids')
def_tok_ids = tf.reshape(def_tok_ids, [def_tok_shape[0], def_tok_shape[1], -1], name='reshape2tok_ids')
with tf.device('/cpu:0'):
def_sen_ids = tf.map_fn(
sentence_tagging_and_padding,
(def_tok_ids, def_sequence_len),
dtype=tf.int32, back_prop=False, parallel_iterations=_parallel_iterations,
name='map_fn2get_sen_ids')
# Get ops to compute the LM embeddings.
def_embeddings_op = bilm(def_sen_ids)
# Get an op to compute ELMo (weighted average of the internal biLM layers)
def_elmo_output = weight_layers('elmo_output', def_embeddings_op, l2_coef=0.0, reuse=True)
def_weighted_op = def_elmo_output['weighted_op']
def_mean_op = def_elmo_output['mean_op']
def_word_emb = def_elmo_output['word_emb']
def_lstm_outputs1 = def_elmo_output['lstm_outputs1']
def_lstm_outputs2 = def_elmo_output['lstm_outputs2']
hub.add_signature("default", {"strings": def_strings},
{"elmo": def_weighted_op,
"default": def_mean_op,
"word_emb": def_word_emb,
"lstm_outputs1": def_lstm_outputs1,
"lstm_outputs2": def_lstm_outputs2,
"version": version})
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.**3 - 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
})
def invalid_text_module_fn(): text = tf.placeholder(tf.string, shape=[10]) hub.add_signature(inputs=text, outputs=tf.zeros([10, 3]))
def assets_module_fn():
indices = tf.placeholder(dtype=tf.int64, name="indices")
table = tf.contrib.lookup.index_to_string_table_from_file(
vocabulary_file=vocab_filename, default_value="UNKNOWN")
outputs = table.lookup(indices)
hub.add_signature(inputs=indices, outputs=outputs)
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
})
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)
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})
def update_ops_module_fn(): counter = tf.Variable(0, trainable=False) tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, counter.assign_add(1)) hub.add_signature(inputs=None, outputs=counter.value())
def plus_one():
x = tf.compat.v1.placeholder(dtype=tf.float32, name='x')
y = x + 1
hub.add_signature(inputs=x, outputs=y)
def multiple_outputs_module_fn():
x = tf.placeholder(dtype=tf.float32)
v = tf.Variable([3.0])
hub.add_signature(
inputs={"x": x},
outputs={"y": v * x, "z": v * v * x})
def assets_module_fn():
indices = tf_v1.placeholder(dtype=tf.int64, name="indices")
table = index_to_string_table_from_file(
vocabulary_file=vocab_filename, default_value="UNKNOWN")
outputs = table.lookup(indices)
hub.add_signature(inputs=indices, outputs=outputs)
def AutoEncoder_module():
inputs = tf.placeholder(tf.float32, shape=[None, IMAGE_DIM])
output = tf.reshape(inputs, [-1, 1, 96, 96])
# Compression
# 96x96
output = lib.ops.conv2d.Conv2D('AutoEncoder.1',
1,
DIM,
5,
output,
stride=2)
output = tf.nn.relu(output)
#48x48
output = lib.ops.conv2d.Conv2D('AutoEncoder.2',
DIM,
2 * DIM,
5,
output,
stride=2)
output = tf.nn.relu(output)
#24x24
output = lib.ops.conv2d.Conv2D('AutoEncoder.3',
2 * DIM,
4 * DIM,
5,
output,
stride=2)
output = tf.nn.relu(output)
#12x12
output = lib.ops.conv2d.Conv2D('AutoEncoder.4',
4 * DIM,
8 * DIM,
5,
output,
stride=2)
output = tf.nn.relu(output)
#6x6
#output = lib.ops.conv2d.Conv2D('AutoEncoder.5', 8*DIM, int(DIM/64), 5, output)
#output = tf.nn.relu(output)
#output = tf.reshape(tf.layers.flatten(output), (-1, LATENT_DIM))
output = tf.layers.flatten(output)
output_latent = tf.layers.dense(output, LATENT_DIM, activation=None)
#tf.Print(output_latent)
print(output_latent)
hub.add_signature(inputs=inputs, outputs=output_latent, name='latent')
# Decompressioin
#output = tf.reshape(
activation = None
output = tf.layers.dense(output_latent,
DIM * 6 * 6,
use_bias=False,
activation=activation)
output = tf.reshape(output, [-1, DIM, 6, 6])
output = lib.ops.deconv2d.Deconv2D('AutoEncoder.6', DIM, 4 * DIM, 5,
output)
output = tf.nn.relu(output)
output = lib.ops.deconv2d.Deconv2D('AutoEncoder.7', 4 * DIM, 2 * DIM, 5,
output)
output = tf.nn.relu(output)
output = lib.ops.deconv2d.Deconv2D('AutoEncoder.8', 2 * DIM, DIM, 5,
output)
output = tf.nn.relu(output)
output = lib.ops.deconv2d.Deconv2D('AutoEncoder.9', DIM, 1, 5, output)
output = tf.nn.relu(output)
output = tf.reshape(output, [-1, IMAGE_DIM])
hub.add_signature(inputs=output_latent,
outputs=output,
name='decode_latent')
hub.add_signature(inputs=inputs, outputs=output)
return output
