def test_f1_score(tensor_true, tensor_pred):
assert tf.is_tensor(f1_score(tensor_true, tensor_pred))
def wcce(y_true, y_pred):
Kweights = tf.constant(weights)
if not tf.is_tensor(y_pred): y_pred = tf.constant(y_pred)
y_true = tf.cast(y_true, y_pred.dtype)
return tf.keras.backend.categorical_crossentropy(
y_true, y_pred) * tf.keras.backend.sum(y_true * Kweights, axis=-1)
def fake_quantize(inputs,
quant_min=None,
quant_max=None,
num_bits=8,
axis=None,
unsigned=False,
affine=False):
"""Universal tensor fake quantization function
Args:
inputs: A Tensor of dtype float32.
quant_min: Scalar (0-d Tensor), 1-d Tensor or None
quant_max: Scalar (0-d Tensor), 1-d Tensor or None.
num_bits: An integer used to calculate scaling factor, `scale = (2^(num_bits-1) - 1) / max`.
Effectively, it indicates how many integer bits is used to represent the value.
axis: Integer or None. If specified, quant_min and quant_max must be vector and will be broadcasted to inputs.
Default None, indicates per tensor quantization.
unsigned: A boolean. If True, use unsigned int8. Default False.
affine: A boolean. If True, use affine quantization. Default False.
Returns:
outputs: A Tensor with same type as inputs
Raises:
TypeError: Wrong input types.
ValueError: Wrong input values.
"""
if not tf.is_tensor(inputs):
raise TypeError("inputs should be a Tensor")
if not isinstance(num_bits, int):
raise TypeError("num_bits should be an integer")
if num_bits <= 0:
raise ValueError("num_bits should > 0")
if quant_max is None and quant_min is None:
raise NotImplementedError("dynamic quantization is not supported yet")
if quant_min is not None and quant_max is not None:
if not tf.is_tensor(quant_max) or not tf.is_tensor(quant_min):
raise TypeError(
"quant_min and quant_max should be Scalar (0-d Tensor), 1-d Tensor or None"
)
if quant_max.shape != quant_min.shape:
raise ValueError("shape mismatch between quant_min and quant_max")
if len(quant_max.shape) == 0:
if axis is not None:
raise ValueError(
"quan_min/quant_max is a Scalar, support per tensor quantization, axis must be None"
)
elif len(quant_max.shape) == 1:
if axis is None:
raise ValueError(
"quan_min/quant_max is a Tensor, support per axis quantization, axis must be set"
)
if not isinstance(axis, int):
raise TypeError("axis should be an integer")
if not -len(inputs.shape) <= axis < len(inputs.shape):
raise ValueError(
"invalid axis {} for inputs with dimentsion {}".format(
axis, len(inputs.shape)))
else:
raise ValueError(
"quant_min and quant_max should be Scalar (0-d Tensor), 1-d Tensor or None"
)
else:
raise ValueError("one of quant_min and quant_max is None")
# do broadcast obviously for per axis quantization
if axis is not None:
if axis < 0:
axis += len(inputs.shape)
for i in range(len(inputs.shape)):
if i != axis:
quant_min = tf.expand_dims(quant_min, i)
quant_max = tf.expand_dims(quant_max, i)
epsilon = 1. / (1 << 24) # Minimum fp16 representable
@tf.custom_gradient
def fake_quantize_core(inputs, quant_min, quant_max):
def _scaled_fake_quantize(inputs, quant_min, quant_max):
# TODO(haow): Add check for negative values in inputs if unsigned
bound = 2.0**(num_bits - 1 + int(unsigned)) - 1.0
quant_amax = tf.maximum(tf.abs(quant_min), tf.abs(quant_max))
scale = bound / quant_amax
# Treat quant_max smaller than minimum representable of fp16 0.
# Value quantized with quant_amax=0 should all be 0, thus set scale to 1
scale = tf.compat.v2.where(tf.math.less_equal(quant_amax, epsilon),
tf.constant(1.), scale)
quantized = tf.clip_by_value(tf.math.round(inputs * scale), -bound,
bound)
outputs = quantized / scale
return outputs
def _affine_fake_quantize(inputs, quant_min, quant_max):
if unsigned:
min_bound = 0
max_bound = 2.0**num_bits - 1.0
zero_point = 2.0**(num_bits - 1)
else:
min_bound = -2.0**(num_bits - 1)
max_bound = 2.0**(num_bits - 1) - 1.0
zero_point = 0.0
step_size = (quant_max - quant_min) / (2.0**num_bits - 1.0)
# in case step_size is too small, may need improved
quant_zero = tf.compat.v2.where(
tf.math.less_equal(step_size, epsilon),
tf.constant(zero_point),
tf.math.round(quant_min / step_size) - min_bound)
quantized = tf.compat.v2.where(
tf.math.less_equal(step_size, epsilon), quant_zero,
tf.clip_by_value(
tf.math.round(inputs / step_size) - quant_zero, min_bound,
max_bound))
outputs = tf.compat.v2.where(
tf.math.less_equal(step_size, epsilon), quant_max,
(quantized + quant_zero) * step_size)
return outputs
if not affine:
outputs = _scaled_fake_quantize(inputs, quant_min, quant_max)
else:
outputs = _affine_fake_quantize(inputs, quant_min, quant_max)
def grad(grad_outputs):
# Boundary check is exclusive in case we'll need to support inplace
if not affine:
pass_condition = tf.math.less(
tf.abs(inputs),
tf.maximum(tf.abs(quant_min), tf.abs(quant_max)))
else:
pass_condition = tf.math.logical_and(
tf.math.greater(inputs, quant_min),
tf.math.less(inputs, quant_max))
grad_inputs = tf.compat.v2.where(pass_condition, grad_outputs,
tf.constant(0.))
return grad_inputs, None, None
return outputs, grad
return fake_quantize_core(inputs, quant_min, quant_max)
def _compare_s2s(s1, s2, method, s1_shape=None, s2_shape=None):
s1_is_tensor = tf.is_tensor(s1)
s2_is_tensor = tf.is_tensor(s2)
if s1_is_tensor or s2_is_tensor:
if not s1_is_tensor:
s1 = tf.convert_to_tensor(s1, dtype=tf.float32)
if not s2_is_tensor:
s2 = tf.convert_to_tensor(s2, dtype=tf.float32)
s12_is_tensor = True
s1_height = s1_shape[0]
s2_height = s2_shape[0]
s1_width = s1_shape[1]
s2_width = s2_shape[1]
depth = s1_shape[2]
size = depth * s1_width * s1_height
else:
if type(s1) is not np.ndarray:
s1 = np.asarray(s1)
if type(s2) is not np.ndarray:
s2 = np.asarray(s2)
s12_is_tensor = False
s1_height = s1.shape[1]
s2_height = s2.shape[1]
s1_width = s1.shape[2]
s2_width = s2.shape[2]
depth = s1.shape[3]
size = s1.size
areas = []
diffs = []
if s1_width > s2_width:
width_min = s2_width
width_max = s1_width
p1_c = 0
p2_c = 1
else:
width_min = s1_width
width_max = s2_width
p1_c = 1
p2_c = 0
p1w_cr = p1_c
p2w_cr = p2_c
if s1_height > s2_height:
height_min = s2_height
height_max = s1_height
p1_r = 0
p2_r = 1
else:
height_min = s1_height
height_max = s2_height
p1_r = 1
p2_r = 0
p1h_cr = p1_r
p2h_cr = p2_r
ws_min = width_min / width_max
ws_max = width_max / width_min
hs_min = height_min / height_max
hs_max = height_max / height_min
for h_max in range(height_max):
h_min = h_max * hs_min
h_min_max = int(h_min + 1) * hs_max
for w_max in range(width_max):
w_min = w_max * ws_min
w_min_max = int(w_min + 1) * ws_max
if s1_width > s2_width:
p1w = w_max
p2w = int(w_min)
else:
p1w = int(w_min)
p2w = w_max
if s1_height > s2_height:
p1h = h_max
p2h = int(h_min)
else:
p1h = int(h_min)
p2h = h_max
areas.append((min(w_max + 1, w_min_max) - w_max) *
(min(h_max + 1, h_min_max) - h_max))
diffs.append(
pixels_diff(s1[:, p1h, p1w, :], s2[:, p2h, p2w, :], method))
if col_changed(w_min, w_min + ws_min):
areas.append((w_max + 1 - w_min_max) *
(min(h_max + 1, h_min_max) - h_max))
diffs.append(
pixels_diff(s1[:, p1h, p1w + p1_c, :],
s2[:, p2h, p2w + p2_c, :], method))
if row_changed(h_min, h_min + hs_min):
areas.append((min(w_max + 1, w_min_max) - w_max) *
(h_max + 1 - h_min_max))
diffs.append(
pixels_diff(s1[:, p1h + p1_r, p1w, :], s2[:, p2h + p2_r,
p2w, :], method))
if col_row_changed(w_min, w_min + ws_min, h_min, h_min + hs_min):
areas.append((w_max + 1 - w_min_max) * (h_max + 1 - h_min_max))
diffs.append(
pixels_diff(s1[:, p1h + p1h_cr, p1w + p1w_cr, :],
s2[:, p2h + p2h_cr, p2w + p2w_cr, :], method))
if not s12_is_tensor:
areas = np.asarray(areas)
diffs = np.sum(np.asarray(diffs), axis=-1)
result = np.sum(np.multiply(areas, diffs))
else:
areas = tf.convert_to_tensor(areas, dtype=tf.float32)
diffs = tf.reduce_sum(tf.stack(diffs), axis=-1)
result = tf.reduce_sum(tf.multiply(areas, diffs))
result *= min(1, s1_width / s2_width) * min(1, s1_height / s2_height)
if method == _COMPARE_MAE or \
method == _COMPARE_MSE or \
method == _COMPARE_RMSE or \
method == _COMPARE_PSNR :
result /= size
if method == _COMPARE_RMSE:
result = math.sqrt(result)
elif method == _COMPARE_PSNR:
if not s12_is_tensor:
result = 10 * math.log(65025 / result)
else:
result = 10 * tf.math.log(1 / result)
return result
def call(
self,
input_ids=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
labels=None,
training=False,
**kwargs,
):
r"""
labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
Labels for computing the cross entropy classification loss. Indices should be in ``[0, ...,
config.vocab_size - 1]``.
"""
inputs = input_processing(
func=self.call,
config=self.config,
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
labels=labels,
training=training,
kwargs_call=kwargs,
)
transformer_outputs = self.transformer(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
token_type_ids=inputs["token_type_ids"],
position_ids=inputs["position_ids"],
head_mask=inputs["head_mask"],
inputs_embeds=inputs["inputs_embeds"],
output_attentions=inputs["output_attentions"],
output_hidden_states=inputs["output_hidden_states"],
return_dict=inputs["return_dict"],
training=inputs["training"],
)
hidden_states = transformer_outputs[0]
logits = self.score(hidden_states)
logits_shape = shape_list(logits)
in_logits = None
if self.config.pad_token_id is None:
sequence_lengths = -1
else:
if inputs["input_ids"] is not None:
sequence_lengths = (
tf.reduce_sum(
tf.cast(tf.math.not_equal(inputs["input_ids"], self.config.pad_token_id), tf.int32),
-1,
keepdims=False,
)
- 1
)
def get_seq_element(sequence_position, input_batch):
return tf.strided_slice(
input_batch, [sequence_position, 0], [sequence_position + 1, input_batch.shape[-1]], [1, 1]
)
result = tf.map_fn(
fn=lambda t: get_seq_element(t[0], t[1]), elems=[sequence_lengths, logits], dtype="float"
)
in_logits = tf.reshape(result, [logits_shape[0], logits_shape[-1]])
else:
sequence_lengths = -1
logger.warning(
f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
)
loss = None
if inputs["labels"] is not None:
if input_ids is not None:
batch_size, sequence_length = shape_list(inputs["input_ids"])[:2]
else:
batch_size, sequence_length = shape_list(inputs["inputs_embeds"])[:2]
assert (
self.config.pad_token_id is not None or batch_size == 1
), "Cannot handle batch sizes > 1 if no padding token is defined."
if not tf.is_tensor(sequence_lengths):
in_logits = logits[0:batch_size, sequence_lengths]
loss = self.compute_loss(
tf.reshape(inputs["labels"], [-1, 1]), tf.reshape(in_logits, [-1, self.num_labels])
)
pooled_logits = in_logits if in_logits is not None else logits
if not inputs["return_dict"]:
output = (pooled_logits,) + transformer_outputs[1:]
return ((loss,) + output) if loss is not None else output
return TFSequenceClassifierOutput(
loss=loss,
logits=pooled_logits,
hidden_states=transformer_outputs.hidden_states,
attentions=transformer_outputs.attentions,
)
def mssim_loss(self, y_true, y_pred):
if not tf.is_tensor(y_true):
y_true = K.variable(y_true)
y_pred = K.variable(y_pred)
return tf.reduce_mean(tf.image.ssim_multiscale(y_true, y_pred, 2.0))
def to_representation_for_type(value,
tf_function_cache,
type_spec=None,
device=None):
"""Verifies or converts the `value` to an eager object matching `type_spec`.
WARNING: This function is only partially implemented. It does not support
data sets at this point.
The output of this function is always an eager tensor, eager dataset, a
representation of a TensorFlow computation, or a nested structure of those
that matches `type_spec`, and when `device` has been specified, everything
is placed on that device on a best-effort basis.
TensorFlow computations are represented here as zero- or one-argument Python
callables that accept their entire argument bundle as a single Python object.
Args:
value: The raw representation of a value to compare against `type_spec` and
potentially to be converted.
tf_function_cache: A cache obeying `dict` semantics that can be used to look
up previously embedded TensorFlow functions.
type_spec: An instance of `tff.Type`, can be `None` for values that derive
from `typed_object.TypedObject`.
device: The optional device to place the value on (for tensor-level values).
Returns:
Either `value` itself, or a modified version of it.
Raises:
TypeError: If the `value` is not compatible with `type_spec`.
"""
type_spec = type_utils.reconcile_value_with_type_spec(value, type_spec)
if isinstance(value, computation_base.Computation):
return to_representation_for_type(
computation_impl.ComputationImpl.get_proto(value), tf_function_cache,
type_spec, device)
elif isinstance(value, pb.Computation):
key = (value.SerializeToString(), str(type_spec), device)
cached_fn = tf_function_cache.get(key)
if cached_fn:
return cached_fn
embedded_fn = embed_tensorflow_computation(value, type_spec, device)
tf_function_cache[key] = embedded_fn
return embedded_fn
elif isinstance(type_spec, computation_types.NamedTupleType):
type_elem = anonymous_tuple.to_elements(type_spec)
value_elem = (
anonymous_tuple.to_elements(anonymous_tuple.from_container(value)))
result_elem = []
if len(type_elem) != len(value_elem):
raise TypeError('Expected a {}-element tuple, found {} elements.'.format(
len(type_elem), len(value_elem)))
for (t_name, el_type), (v_name, el_val) in zip(type_elem, value_elem):
if t_name != v_name:
raise TypeError(
'Mismatching element names in type vs. value: {} vs. {}.'.format(
t_name, v_name))
el_repr = to_representation_for_type(el_val, tf_function_cache, el_type,
device)
result_elem.append((t_name, el_repr))
return anonymous_tuple.AnonymousTuple(result_elem)
elif device is not None:
py_typecheck.check_type(device, str)
with tf.device(device):
return to_representation_for_type(
value, tf_function_cache, type_spec=type_spec, device=None)
elif isinstance(value, EagerValue):
return value.internal_representation
elif isinstance(value, executor_value_base.ExecutorValue):
raise TypeError(
'Cannot accept a value embedded within a non-eager executor.')
elif isinstance(type_spec, computation_types.TensorType):
if not tf.is_tensor(value):
value = tf.convert_to_tensor(value, dtype=type_spec.dtype)
elif hasattr(value, 'read_value'):
# a tf.Variable-like result, get a proper tensor.
value = value.read_value()
value_type = (
computation_types.TensorType(value.dtype.base_dtype, value.shape))
if not type_analysis.is_assignable_from(type_spec, value_type):
raise TypeError(
'The apparent type {} of a tensor {} does not match the expected '
'type {}.'.format(value_type, value, type_spec))
return value
elif isinstance(type_spec, computation_types.SequenceType):
if isinstance(value, list):
value = tensorflow_utils.make_data_set_from_elements(
None, value, type_spec.element)
py_typecheck.check_type(value,
type_conversions.TF_DATASET_REPRESENTATION_TYPES)
element_type = computation_types.to_type(value.element_spec)
value_type = computation_types.SequenceType(element_type)
type_analysis.check_assignable_from(type_spec, value_type)
return value
else:
raise TypeError('Unexpected type {}.'.format(type_spec))
def serialize_scene(canvas_width,
canvas_height,
shapes,
shape_groups,
filter=pydiffvg.PixelFilter(type=diffvg.FilterType.box,
radius=tf.constant(0.5)),
output_type=OutputType.color,
use_prefiltering=False):
"""
Given a list of shapes, convert them to a linear list of argument,
so that we can use it in TF.
"""
with tf.device('/device:cpu:' + str(pydiffvg.get_cpu_device_id())):
num_shapes = len(shapes)
num_shape_groups = len(shape_groups)
args = []
args.append(tf.constant(canvas_width))
args.append(tf.constant(canvas_height))
args.append(tf.constant(num_shapes))
args.append(tf.constant(num_shape_groups))
args.append(tf.constant(output_type))
args.append(tf.constant(use_prefiltering))
for shape in shapes:
if isinstance(shape, pydiffvg.Circle):
args.append(ShapeType.asTensor(diffvg.ShapeType.circle))
args.append(tf.identity(shape.radius))
args.append(tf.identity(shape.center))
elif isinstance(shape, pydiffvg.Ellipse):
args.append(ShapeType.asTensor(diffvg.ShapeType.ellipse))
args.append(tf.identity(shape.radius))
args.append(tf.identity(shape.center))
elif isinstance(shape, pydiffvg.Path):
assert (shape.points.shape[1] == 2)
args.append(ShapeType.asTensor(diffvg.ShapeType.path))
args.append(
tf.identity(shape.num_control_points, type=tf.int32))
args.append(tf.identity(shape.points))
args.append(tf.constant(shape.is_closed))
elif isinstance(shape, pydiffvg.Polygon):
assert (shape.points.shape[1] == 2)
args.append(ShapeType.asTensor(diffvg.ShapeType.path))
if shape.is_closed:
args.append(tf.zeros(shape.points.shape[0],
dtype=tf.int32))
else:
args.append(
tf.zeros(shape.points.shape[0] - 1, dtype=tf.int32))
args.append(tf.identity(shape.points))
args.append(tf.constant(shape.is_closed))
elif isinstance(shape, pydiffvg.Rect):
args.append(ShapeType.asTensor(diffvg.ShapeType.rect))
args.append(tf.identity(shape.p_min))
args.append(tf.identity(shape.p_max))
else:
assert (False)
args.append(tf.identity(shape.stroke_width))
for shape_group in shape_groups:
args.append(tf.identity(shape_group.shape_ids))
# Fill color
if shape_group.fill_color is None:
args.append(__EMPTY_TENSOR)
elif tf.is_tensor(shape_group.fill_color):
args.append(ColorType.asTensor(diffvg.ColorType.constant))
args.append(tf.identity(shape_group.fill_color))
elif isinstance(shape_group.fill_color, pydiffvg.LinearGradient):
args.append(
ColorType.asTensor(diffvg.ColorType.linear_gradient))
args.append(tf.identity(shape_group.fill_color.begin))
args.append(tf.identity(shape_group.fill_color.end))
args.append(tf.identity(shape_group.fill_color.offsets))
args.append(tf.identity(shape_group.fill_color.stop_colors))
elif isinstance(shape_group.fill_color, pydiffvg.RadialGradient):
args.append(
ColorType.asTensor(diffvg.ColorType.radial_gradient))
args.append(tf.identity(shape_group.fill_color.center))
args.append(tf.identity(shape_group.fill_color.radius))
args.append(tf.identity(shape_group.fill_color.offsets))
args.append(tf.identity(shape_group.fill_color.stop_colors))
if shape_group.fill_color is not None:
# go through the underlying shapes and check if they are all closed
for shape_id in shape_group.shape_ids:
if isinstance(shapes[shape_id], pydiffvg.Path):
if not shapes[shape_id].is_closed:
warnings.warn(
"Detected non-closed paths with fill color. This might causes unexpected results.",
Warning)
# Stroke color
if shape_group.stroke_color is None:
args.append(__EMPTY_TENSOR)
elif tf.is_tensor(shape_group.stroke_color):
args.append(tf.constant(0))
args.append(tf.identity(shape_group.stroke_color))
elif isinstance(shape_group.stroke_color, pydiffvg.LinearGradient):
args.append(
ColorType.asTensor(diffvg.ColorType.linear_gradient))
args.append(tf.identity(shape_group.stroke_color.begin))
args.append(tf.identity(shape_group.stroke_color.end))
args.append(tf.identity(shape_group.stroke_color.offsets))
args.append(tf.identity(shape_group.stroke_color.stop_colors))
elif isinstance(shape_group.stroke_color, pydiffvg.RadialGradient):
args.append(
ColorType.asTensor(diffvg.ColorType.radial_gradient))
args.append(tf.identity(shape_group.stroke_color.center))
args.append(tf.identity(shape_group.stroke_color.radius))
args.append(tf.identity(shape_group.stroke_color.offsets))
args.append(tf.identity(shape_group.stroke_color.stop_colors))
args.append(tf.constant(shape_group.use_even_odd_rule))
# Transformation
args.append(tf.identity(shape_group.shape_to_canvas))
args.append(FilterType.asTensor(filter.type))
args.append(tf.constant(filter.radius))
return args
def calc_tau_int_from_dir(
input_path: str,
hmc: bool = False,
px_cutoff: float = None,
therm_frac: float = 0.2,
num_pts: int = 50,
nstart: int = 100,
make_plot: bool = True,
save_data: bool = True,
keep_charges: bool = False,
):
"""
NOTE: `load_charges_from_dir` returns `output_arr`:
`output_arr` is a list of dicts, each of which look like:
output = {
'beta': beta,
'lf': lf,
'eps': eps,
'traj_len': lf * eps,
'qarr': charges,
'run_params': params,
'run_dir': run_dir,
}
"""
output_arr = load_charges_from_dir(input_path, hmc=hmc)
if output_arr is None:
logger.info(', '.join([
'WARNING: Skipping entry!',
f'\t unable to load charge data from {input_path}',
]))
return None
tint_data = {}
for output in output_arr:
run_dir = output['run_dir']
beta = output['beta']
lf = output['lf']
# eps = output['eps']
# traj_len = output['traj_len']
run_params = output['run_params']
if hmc:
data_dir = os.path.join(input_path, 'run_data')
plot_dir = os.path.join(input_path, 'plots', 'tau_int_plots')
else:
run_dir = output['run_params']['run_dir']
data_dir = os.path.join(run_dir, 'run_data')
plot_dir = os.path.join(run_dir, 'plots')
outfile = os.path.join(data_dir, 'tau_int_data.z')
outfile1 = os.path.join(data_dir, 'tint_data.z')
fdraws = os.path.join(plot_dir, 'tau_int_vs_draws.pdf')
ftlen = os.path.join(plot_dir, 'tau_int_vs_traj_len.pdf')
c1 = os.path.isfile(outfile)
c11 = os.path.isfile(outfile1)
c2 = os.path.isfile(fdraws)
c3 = os.path.isfile(ftlen)
if c1 or c11 or c2 or c3:
loaded = io.loadz(outfile)
output.update(loaded)
logger.info(', '.join([
'WARNING: Loading existing data'
f'\t Found existing data at: {outfile}.',
]))
loaded = io.loadz(outfile)
n = loaded.get('draws', loaded.get('narr', None))
tint = loaded.get('tau_int', loaded.get('tint', None))
output.update(loaded)
xeps_check = 'xeps' in output['run_params'].keys()
veps_check = 'veps' in output['run_params'].keys()
if xeps_check and veps_check:
xeps = tf.reduce_mean(output['run_params']['xeps'])
veps = tf.reduce_mean(output['run_params']['veps'])
eps = tf.reduce_mean([xeps, veps]).numpy()
else:
eps = output['eps']
if isinstance(eps, list):
eps = tf.reduce_mean(eps)
elif tf.is_tensor(eps):
try:
eps = eps.numpy()
except AttributeError:
eps = tf.reduce_mean(eps)
traj_len = lf * eps
qarr, _ = therm_arr(output['qarr'], therm_frac=therm_frac)
n, tint = calc_autocorr(qarr.T, num_pts=num_pts, nstart=nstart)
# output.update({
# 'draws': n,
# 'tau_int': tint,
# 'qarr.shape': qarr.shape,
# })
tint_data[run_dir] = {
'run_dir': run_dir,
'run_params': run_params,
'lf': lf,
'eps': eps,
'traj_len': traj_len,
'narr': n,
'tint': tint,
'qarr.shape': qarr.shape,
}
if save_data:
io.savez(tint_data, outfile, 'tint_data')
# io.savez(output, outfile, name='tint_data')
if make_plot:
# fbeta = os.path.join(plot_dir, 'tau_int_vs_beta.pdf')
io.check_else_make_dir(plot_dir)
prefix = 'HMC' if hmc else 'L2HMC'
xlabel = 'draws'
ylabel = r'$\tau_{\mathrm{int}}$ (estimate)'
title = (f'{prefix}, '
+ r'$\beta=$' + f'{beta}, '
+ r'$N_{\mathrm{lf}}=$' + f'{lf}, '
+ r'$\varepsilon=$' + f'{eps:.2g}, '
+ r'$\lambda=$' + f'{traj_len:.2g}')
_, ax = plt.subplots(constrained_layout=True)
best = []
for t in tint.T:
_ = ax.plot(n, t, marker='.', color='k')
best.append(t[-1])
_ = ax.set_ylabel(ylabel)
_ = ax.set_xlabel(xlabel)
_ = ax.set_title(title)
_ = ax.set_xscale('log')
_ = ax.set_yscale('log')
_ = ax.grid(alpha=0.4)
logger.info(f'Saving figure to: {fdraws}')
_ = plt.savefig(fdraws, dpi=400, bbox_inches='tight')
plt.close('all')
_, ax = plt.subplots()
for b in best:
_ = ax.plot(traj_len, b, marker='.', color='k')
_ = ax.set_ylabel(ylabel)
_ = ax.set_xlabel(r'trajectory length, $\lambda$')
_ = ax.set_title(title)
_ = ax.set_yscale('log')
_ = ax.grid(True, alpha=0.4)
logger.info(f'Saving figure to: {ftlen}')
_ = plt.savefig(ftlen, dpi=400, bbox_inches='tight')
plt.close('all')
return tint_data
def assemble_result_from_graph(type_spec, binding, output_map):
"""Assembles a result stamped into a `tf.Graph` given type signature/binding.
This method does roughly the opposite of `capture_result_from_graph`, in that
whereas `capture_result_from_graph` starts with a single structured object
made up of tensors and computes its type and bindings, this method starts
with the type/bindings and constructs a structured object made up of tensors.
Args:
type_spec: The type signature of the result to assemble, an instance of
`types.Type` or something convertible to it.
binding: The binding that relates the type signature to names of tensors in
the graph, an instance of `pb.TensorFlow.Binding`.
output_map: The mapping from tensor names that appear in the binding to
actual stamped tensors (possibly renamed during import).
Returns:
The assembled result, a Python object that is composed of tensors, possibly
nested within Python structures such as anonymous tuples.
Raises:
TypeError: If the argument or any of its parts are of an uexpected type.
ValueError: If the arguments are invalid or inconsistent witch other, e.g.,
the type and binding don't match, or the tensor is not found in the map.
"""
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.Type)
py_typecheck.check_type(binding, pb.TensorFlow.Binding)
py_typecheck.check_type(output_map, dict)
for k, v in output_map.items():
py_typecheck.check_type(k, str)
if not tf.is_tensor(v):
raise TypeError(
'Element with key {} in the output map is {}, not a tensor.'.
format(k, py_typecheck.type_string(type(v))))
binding_oneof = binding.WhichOneof('binding')
if type_spec.is_tensor():
if binding_oneof != 'tensor':
raise ValueError(
'Expected a tensor binding, found {}.'.format(binding_oneof))
elif binding.tensor.tensor_name not in output_map:
raise ValueError(
'Tensor named {} not found in the output map.'.format(
binding.tensor.tensor_name))
else:
return output_map[binding.tensor.tensor_name]
elif type_spec.is_struct():
if binding_oneof != 'struct':
raise ValueError(
'Expected a struct binding, found {}.'.format(binding_oneof))
else:
type_elements = structure.to_elements(type_spec)
if len(binding.struct.element) != len(type_elements):
raise ValueError(
'Mismatching tuple sizes in type ({}) and binding ({}).'.
format(len(type_elements), len(binding.struct.element)))
result_elements = []
for (element_name,
element_type), element_binding in zip(type_elements,
binding.struct.element):
element_object = assemble_result_from_graph(
element_type, element_binding, output_map)
result_elements.append((element_name, element_object))
if type_spec.python_container is None:
return structure.Struct(result_elements)
container_type = type_spec.python_container
if (py_typecheck.is_named_tuple(container_type)
or py_typecheck.is_attrs(container_type)):
return container_type(**dict(result_elements))
return container_type(result_elements)
elif type_spec.is_sequence():
if binding_oneof != 'sequence':
raise ValueError(
'Expected a sequence binding, found {}.'.format(binding_oneof))
else:
sequence_oneof = binding.sequence.WhichOneof('binding')
if sequence_oneof == 'variant_tensor_name':
variant_tensor = output_map[
binding.sequence.variant_tensor_name]
return make_dataset_from_variant_tensor(
variant_tensor, type_spec.element)
else:
raise ValueError('Unsupported sequence binding \'{}\'.'.format(
sequence_oneof))
else:
raise ValueError('Unsupported type \'{}\'.'.format(type_spec))
def fetch_value_in_session(sess, value):
"""Fetches `value` in `session`.
Args:
sess: The session in which to perform the fetch (as a single run).
value: A Python object of a form analogous to that constructed by the
function `assemble_result_from_graph`, made of tensors and anononymous
tuples, or a `tf.data.Dataset`.
Returns:
A Python object with structure similar to `value`, but with tensors
replaced with their values, and data sets replaced with lists of their
elements, all fetched with a single call `session.run()`.
Raises:
ValueError: If `value` is not a `tf.data.Dataset` or not a structure of
tensors and anonoymous tuples.
"""
py_typecheck.check_type(sess, tf.compat.v1.Session)
# TODO(b/113123634): Investigate handling `list`s and `tuple`s of
# `tf.data.Dataset`s and what the API would look like to support this.
if isinstance(value, type_conversions.TF_DATASET_REPRESENTATION_TYPES):
with sess.graph.as_default():
iterator = tf.compat.v1.data.make_one_shot_iterator(value)
next_element = iterator.get_next()
elements = []
while True:
try:
elements.append(sess.run(next_element))
except tf.errors.OutOfRangeError:
break
return elements
else:
flattened_value = structure.flatten(value)
dataset_results = {}
flat_tensors = []
for idx, v in enumerate(flattened_value):
if isinstance(v, type_conversions.TF_DATASET_REPRESENTATION_TYPES):
dataset_tensors = fetch_value_in_session(sess, v)
if not dataset_tensors:
# An empty list has been returned; we must pack the shape information
# back in or the result won't typecheck.
element_structure = v.element_spec
dummy_elem = make_dummy_element_for_type_spec(
element_structure)
dataset_tensors = [dummy_elem]
dataset_results[idx] = dataset_tensors
elif tf.is_tensor(v):
flat_tensors.append(v)
else:
raise ValueError('Unsupported value type {}.'.format(v))
# Note that `flat_tensors` could be an empty tuple, but it could also be a
# list of empty tuples.
if flat_tensors or any(x for x in flat_tensors):
flat_computed_tensors = sess.run(flat_tensors)
else:
flat_computed_tensors = flat_tensors
flattened_results = _interleave_dataset_results_and_tensors(
dataset_results, flat_computed_tensors)
def _to_unicode(v):
if isinstance(v, bytes):
return v.decode('utf-8')
return v
if tf.is_tensor(value) and value.dtype == tf.string:
flattened_results = [
_to_unicode(result) for result in flattened_results
]
return structure.pack_sequence_as(value, flattened_results)
def capture_result_from_graph(result, graph):
"""Captures a result stamped into a tf.Graph as a type signature and binding.
Args:
result: The result to capture, a Python object that is composed of tensors,
possibly nested within Python structures such as dictionaries, lists,
tuples, or named tuples.
graph: The instance of tf.Graph to use.
Returns:
A tuple (type_spec, binding), where 'type_spec' is an instance of
computation_types.Type that describes the type of the result, and 'binding'
is an instance of TensorFlow.Binding that indicates how parts of the result
type relate to the tensors and ops that appear in the result.
Raises:
TypeError: If the argument or any of its parts are of an uexpected type.
"""
def _get_bindings_for_elements(name_value_pairs, graph, type_fn):
"""Build `(type_spec, binding)` tuple for name value pairs."""
element_name_type_binding_triples = [
((k, ) + capture_result_from_graph(v, graph))
for k, v in name_value_pairs
]
type_spec = type_fn([((e[0], e[1]) if e[0] else e[1])
for e in element_name_type_binding_triples])
binding = pb.TensorFlow.Binding(struct=pb.TensorFlow.StructBinding(
element=[e[2] for e in element_name_type_binding_triples]))
return type_spec, binding
# TODO(b/113112885): The emerging extensions for serializing SavedModels may
# end up introducing similar concepts of bindings, etc., we should look here
# into the possibility of reusing some of that code when it's available.
if isinstance(result, TENSOR_REPRESENTATION_TYPES):
with graph.as_default():
result = tf.constant(result)
if tf.is_tensor(result):
if hasattr(result, 'read_value'):
# We have a tf.Variable-like result, get a proper tensor to fetch.
with graph.as_default():
result = result.read_value()
return (computation_types.TensorType(result.dtype.base_dtype,
result.shape),
pb.TensorFlow.Binding(tensor=pb.TensorFlow.TensorBinding(
tensor_name=result.name)))
elif py_typecheck.is_named_tuple(result):
# Special handling needed for collections.namedtuples since they do not have
# anything in the way of a shared base class. Note we don't want to rely on
# the fact that collections.namedtuples inherit from 'tuple' because we'd be
# failing to retain the information about naming of tuple members.
# pylint: disable=protected-access
name_value_pairs = result._asdict().items()
# pylint: enable=protected-access
return _get_bindings_for_elements(
name_value_pairs, graph,
functools.partial(computation_types.StructWithPythonType,
container_type=type(result)))
elif py_typecheck.is_attrs(result):
name_value_pairs = attr.asdict(result,
dict_factory=collections.OrderedDict,
recurse=False)
return _get_bindings_for_elements(
name_value_pairs.items(), graph,
functools.partial(computation_types.StructWithPythonType,
container_type=type(result)))
elif isinstance(result, structure.Struct):
return _get_bindings_for_elements(structure.to_elements(result), graph,
computation_types.StructType)
elif isinstance(result, collections.Mapping):
if isinstance(result, collections.OrderedDict):
name_value_pairs = result.items()
else:
name_value_pairs = sorted(result.items())
return _get_bindings_for_elements(
name_value_pairs, graph,
functools.partial(computation_types.StructWithPythonType,
container_type=type(result)))
elif isinstance(result, (list, tuple)):
element_type_binding_pairs = [
capture_result_from_graph(e, graph) for e in result
]
return (computation_types.StructWithPythonType(
[e[0] for e in element_type_binding_pairs], type(result)),
pb.TensorFlow.Binding(struct=pb.TensorFlow.StructBinding(
element=[e[1] for e in element_type_binding_pairs])))
elif isinstance(result, type_conversions.TF_DATASET_REPRESENTATION_TYPES):
variant_tensor = tf.data.experimental.to_variant(result)
element_structure = result.element_spec
try:
element_type = computation_types.to_type(element_structure)
except TypeError as e:
raise TypeError(
'TFF does not support Datasets that yield elements of structure {!s}'
.format(element_structure)) from e
return (computation_types.SequenceType(element_type),
pb.TensorFlow.Binding(sequence=pb.TensorFlow.SequenceBinding(
variant_tensor_name=variant_tensor.name)))
else:
raise TypeError(
'Cannot capture a result of an unsupported type {}.'.format(
py_typecheck.type_string(type(result))))
def test_combined_bce_rmse(tensor_true, tensor_pred):
assert tf.is_tensor(combined_bce_rmse(tensor_true, tensor_pred))
def test_f1_loss(tensor_true, tensor_pred):
assert tf.is_tensor(f1_loss(tensor_true, tensor_pred))
with pytest.raises(ValueError):
f1_loss(tensor_true, tf.constant([[1, 0, 0], [1, 0, 0]]))
def _check_settings(self, data, data_type):
"""Check settings and return necessary parameters for trafo and inverse
trafo method.
Parameters
----------
data : numpy.ndarray or tf.Tensor
The data that will be transformed.
data_type : str
Specifies what kind of data this is. This must be one of:
'ic78', 'deepcore', 'label', 'misc'
Returns
-------
type(data)
The transformed data
Raises
------
ValueError
If DataTransformer object has not created or loaded a trafo model.
If provided data_type is unkown.
"""
dtype = data.dtype
data_type = data_type.lower()
if not self._setup_complete:
raise ValueError('DataTransformer needs to create or load a trafo'
'model prior to transform call.')
if data_type not in ['ic78', 'deepcore', 'label', 'misc']:
raise ValueError('data_type {!r} is unknown!'.format(data_type))
# check if shape of data matches expected shape
if data_type == 'ic78':
shape = [10, 10, 60, self.trafo_model['num_bins']]
elif data_type == 'deepcore':
shape = [8, 60, self.trafo_model['num_bins']]
else:
shape = self.trafo_model['{}_shape'.format(data_type)]
if list(data.shape[1:]) != shape:
raise ValueError('Shape of data {!r} does'.format(data.shape[1:]) +
' not match expected shape {!r}'.format(shape))
if data_type in ['ic78', 'deepcore']:
log_name = 'log_dom_bins'
normalize_name = 'normalize_dom_data'
else:
log_name = 'log_{}_bins'.format(data_type)
normalize_name = 'normalize_{}_data'.format(data_type)
is_tf = tf.is_tensor(data)
if is_tf:
if dtype != self._tf_float_dtype:
data = tf.cast(data, dtype=self._tf_float_dtype)
else:
data = np.array(data, dtype=self._np_float_dtype)
# choose numpy or tensorflow log function
if is_tf:
log_func = tf.math.log
exp_func = tf.exp
else:
log_func = np.log
exp_func = np.exp
return data, log_name, normalize_name, log_func, exp_func, is_tf, dtype
def _mapping_fn(x):
if not tf.is_tensor(x):
x = tf.convert_to_tensor(x)
return computation_types.TensorType(x.dtype.base_dtype, x.shape)
def cross_attention(self,
query,
key,
value,
mask=None,
bidirectional=False,
use_cache=False,
position_bias=None,
position_embeddings=None,
kv_position_embeddings=None,
query_length=None,
past_key_value_state=None,
training=False):
bs, qlen, _ = get_shape(query)
klen = get_shape(key)[1]
if past_key_value_state is not None:
error_message = ("past_key_value_state should have 2 past states: "
"keys and values. Got {} past states".format(
len(past_key_value_state)))
assert (len(past_key_value_state) == 2), error_message
real_qlen = (qlen + get_shape(past_key_value_state[0])[2]
if query_length is None else query_length)
else:
real_qlen = qlen
qlen = real_qlen
q = self.q(query) # (bs, qlen, inner_dim)
q = self._split_heads(q, bs) # (bs, n_heads, qlen, dim_per_head)
if past_key_value_state is None:
k = self.k(key) # (bs, klen, inner_dim)
v = self.v(value) # (bs, klen, inner_dim)
k = self._split_heads(k, bs) # (bs, n_heads, klen, dim_per_head)
v = self._split_heads(v, bs) # (bs, n_heads, vlen, dim_per_head)
else:
k, v = past_key_value_state
if tf.is_tensor(use_cache):
if hasattr(use_cache, "numpy"):
use_cache = bool(use_cache.numpy())
else:
use_cache = True
if use_cache:
present_key_value_state = (k, v)
else:
present_key_value_state = None
if position_bias is None:
if self.num_relative_buckets is None:
raise ValueError(
"No position_bias provided and no weights to compute position_bias"
)
position_bias = self.compute_bias(qlen=qlen,
klen=klen,
bidirectional=bidirectional)
# if key and values are already calculated
# we want only the last query position bias
if past_key_value_state is not None:
position_bias = position_bias[:, :, -1:, :]
if mask is not None:
position_bias = position_bias + mask # (bs, n_heads, seq_len, seq_len)
# (bs, n_heads, seq_len, seq_len)
scores = tf.einsum("bnqd,bnkd->bnqk", q, k)
scores += position_bias
# (bs, n_heads, seq_len, seq_len)
attention_weights = tf.nn.softmax(scores, axis=-1)
# (bs, n_heads, seq_len, seq_len)
attention_weights = self.dropout(attention_weights, training=training)
# (bs, n_heads, seq_len, dim_per_head)
hidden_states = tf.matmul(attention_weights, v)
hidden_states = self._join_heads(hidden_states,
bs) # (bs, seq_len, dim)
hidden_states = self.o(hidden_states) # (bs, seq_len, out_dim)
outputs = {
"hidden_states": hidden_states,
"key_value_state": present_key_value_state,
"attention_weights": attention_weights,
"position_bias": position_bias,
}
return outputs
def infer_type(arg: Any) -> Optional[computation_types.Type]:
"""Infers the TFF type of the argument (a `computation_types.Type` instance).
WARNING: This function is only partially implemented.
The kinds of arguments that are currently correctly recognized:
- tensors, variables, and data sets,
- things that are convertible to tensors (including numpy arrays, builtin
types, as well as lists and tuples of any of the above, etc.),
- nested lists, tuples, namedtuples, anonymous tuples, dict, and OrderedDicts.
Args:
arg: The argument, the TFF type of which to infer.
Returns:
Either an instance of `computation_types.Type`, or `None` if the argument is
`None`.
"""
# TODO(b/113112885): Implement the remaining cases here on the need basis.
if arg is None:
return None
elif isinstance(arg, typed_object.TypedObject):
return arg.type_signature
elif tf.is_tensor(arg):
return computation_types.TensorType(arg.dtype.base_dtype, arg.shape)
elif isinstance(arg, TF_DATASET_REPRESENTATION_TYPES):
element_type = computation_types.to_type(arg.element_spec)
return computation_types.SequenceType(element_type)
elif isinstance(arg, anonymous_tuple.AnonymousTuple):
return computation_types.NamedTupleType([
(k, infer_type(v)) if k else infer_type(v)
for k, v in anonymous_tuple.iter_elements(arg)
])
elif py_typecheck.is_attrs(arg):
items = attr.asdict(arg,
dict_factory=collections.OrderedDict,
recurse=False)
return computation_types.NamedTupleTypeWithPyContainerType(
[(k, infer_type(v)) for k, v in items.items()], type(arg))
elif py_typecheck.is_named_tuple(arg):
items = arg._asdict()
return computation_types.NamedTupleTypeWithPyContainerType(
[(k, infer_type(v)) for k, v in items.items()], type(arg))
elif isinstance(arg, dict):
if isinstance(arg, collections.OrderedDict):
items = arg.items()
else:
items = sorted(arg.items())
return computation_types.NamedTupleTypeWithPyContainerType(
[(k, infer_type(v)) for k, v in items], type(arg))
elif isinstance(arg, (tuple, list)):
elements = []
all_elements_named = True
for element in arg:
all_elements_named &= py_typecheck.is_name_value_pair(element)
elements.append(infer_type(element))
# If this is a tuple of (name, value) pairs, the caller most likely intended
# this to be a NamedTupleType, so we avoid storing the Python container.
if all_elements_named:
return computation_types.NamedTupleType(elements)
else:
return computation_types.NamedTupleTypeWithPyContainerType(
elements, type(arg))
elif isinstance(arg, str):
return computation_types.TensorType(tf.string)
elif isinstance(arg, (np.generic, np.ndarray)):
return computation_types.TensorType(tf.dtypes.as_dtype(arg.dtype),
arg.shape)
else:
dtype = {
bool: tf.bool,
int: tf.int32,
float: tf.float32
}.get(type(arg))
if dtype:
return computation_types.TensorType(dtype)
else:
# Now fall back onto the heavier-weight processing, as all else failed.
# Use make_tensor_proto() to make sure to handle it consistently with
# how TensorFlow is handling values (e.g., recognizing int as int32, as
# opposed to int64 as in NumPy).
try:
# TODO(b/113112885): Find something more lightweight we could use here.
tensor_proto = tf.make_tensor_proto(arg)
return computation_types.TensorType(
tf.dtypes.as_dtype(tensor_proto.dtype),
tf.TensorShape(tensor_proto.tensor_shape))
except TypeError as err:
raise TypeError(
'Could not infer the TFF type of {}: {}'.format(
py_typecheck.type_string(type(arg)), err))
def sample_shape(self):
"""Sample shape of random variable as a `TensorShape`."""
if tf.is_tensor(self._sample_shape):
return tf.TensorShape(tf.get_static_value(self._sample_shape))
return tf.TensorShape(self._sample_shape)
def has_tensors(*x): return np.any([tf.is_tensor(t) for t in tf.nest.flatten(x)])
def call(
self,
input_ids=None,
past=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
use_cache=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
labels=None,
training=False,
**kwargs,
):
r"""
labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
Labels for computing the cross entropy classification loss. Indices should be in ``[0, ...,
config.vocab_size - 1]``.
"""
inputs = input_processing(
func=self.call,
config=self.config,
input_ids=input_ids,
past=past,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
labels=labels,
training=training,
kwargs_call=kwargs,
)
transformer_outputs = self.transformer(
input_ids=inputs["input_ids"],
past=inputs["past"],
attention_mask=inputs["attention_mask"],
token_type_ids=inputs["token_type_ids"],
position_ids=inputs["position_ids"],
head_mask=inputs["head_mask"],
inputs_embeds=inputs["inputs_embeds"],
use_cache=inputs["use_cache"],
output_attentions=inputs["output_attentions"],
output_hidden_states=inputs["output_hidden_states"],
return_dict=inputs["return_dict"],
training=inputs["training"],
)
hidden_states = transformer_outputs[0]
logits = self.score(hidden_states)
logits_shape = shape_list(logits)
in_logits = None
if self.config.pad_token_id is None:
sequence_lengths = -1
else:
if inputs["input_ids"] is not None:
sequence_lengths = (
tf.reduce_sum(
tf.cast(tf.math.not_equal(inputs["input_ids"], self.config.pad_token_id), tf.int32),
-1,
keepdims=False,
)
- 1
)
in_logits = tf.gather(logits, sequence_lengths, batch_dims=1, axis=1)
else:
sequence_lengths = -1
logger.warning(
f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
)
loss = None
if inputs["labels"] is not None:
assert (
self.config.pad_token_id is not None or logits_shape[0] == 1
), "Cannot handle batch sizes > 1 if no padding token is defined."
if not tf.is_tensor(sequence_lengths):
in_logits = logits[0 : logits_shape[0], sequence_lengths]
loss = self.compute_loss(tf.reshape(inputs["labels"], [-1]), tf.reshape(in_logits, [-1, self.num_labels]))
pooled_logits = in_logits if in_logits is not None else logits
if not inputs["return_dict"]:
output = (pooled_logits,) + transformer_outputs[1:]
return ((loss,) + output) if loss is not None else output
return TFSequenceClassifierOutputWithPast(
loss=loss,
logits=pooled_logits,
past_key_values=transformer_outputs.past_key_values,
hidden_states=transformer_outputs.hidden_states,
attentions=transformer_outputs.attentions,
)
def convert_to_tensor(elem):
if tf.is_tensor(elem):
return elem
elif isinstance(elem, np.ndarray):
return tf.convert_to_tensor(elem)
return tf.convert_to_tensor(elem, dtype=default_float())
def _compare_s2h(s, h, method, s_shape=None, h_shape=None):
s_is_tensor = tf.is_tensor(s)
h_is_tensor = tf.is_tensor(h)
if s_is_tensor or h_is_tensor:
if not s_is_tensor:
s = tf.convert_to_tensor(s, dtype=tf.float32)
if not h_is_tensor:
h = tf.convert_to_tensor(h, dtype=tf.float32)
sh_is_tensor = True
s_height = s_shape[0]
h_height = h_shape[0]
s_width = s_shape[1]
h_width = h_shape[1]
depth = s_shape[2]
s_size = depth * s_width * s_height
h_size = depth * h_width * h_height
else:
if type(s) is not np.ndarray:
s = np.asarray(s)
if type(h) is not np.ndarray:
h = np.asarray(h)
sh_is_tensor = False
s_height = s.shape[1]
h_height = h.shape[1]
s_width = s.shape[2]
h_width = h.shape[2]
depth = s.shape[3]
s_size = s.size
h_size = h.size
areas = []
diffs = []
width_sq = s_width
height_sq = s_height
if h_height > 1:
h_rad_i_w = (s_width / (h_width - 0.5)) / 2
h_rad_i_h = ((s_height / (h_height - 0.5)) / 2) / (math.sqrt(3) / 2)
else:
h_rad_i_w = (s_width / h_width) / 2
h_rad_i_h = (s_height / 2) / (math.sqrt(3) / 2)
h_rad_i = max(h_rad_i_w, h_rad_i_h)
h_rad_o = h_rad_i / (math.sqrt(3) / 2)
h_dia_o = 2 * h_rad_o
h_dia_i = 2 * h_rad_i
h_dist_w = h_dia_i
h_dist_h = 1.5 * h_rad_o
if h_height > 1:
width_hex = h_width * h_dia_i + h_rad_i
height_hex = h_dia_o + (h_height - 1) * h_dist_h
else:
width_hex = h_width * h_dia_i
height_hex = h_dia_o
wb = (width_hex - width_sq) / 2
hb = (height_hex - height_sq) / 2
ws_s = 1.0
ws_h = h_rad_i
hs_s = 1.0
hs_h = h_rad_o / 2
sr_h = hs_h / ws_h
wi = 0.0
hi = 0.0
while hi < height_sq:
hh = hb + hi
hhm = hh % hs_h
hhmr = hs_h - hhm
hm = hi % hs_s
hmr = hs_s - hm
hmr_min = min(hmr, hhmr)
hd = hh % h_dist_h
hn = max(hd + hmr_min, np.nextafter(hd, hd + 1))
hu = int(hi / hs_s)
hhu = int(hh / h_dist_h)
while wi < width_sq:
wh = wb + wi if not hhu % 2 else wb - h_rad_i + wi
whm = wh % ws_h
whmr = ws_h - whm
wm = wi % ws_s
wmr = ws_s - wm
wmr_min = min(wmr, whmr)
wd = wh % h_dist_w
wn = max(wd + wmr_min, np.nextafter(wd, wd + 1))
wu = int(wi / ws_s)
whu = int(wh / h_dist_w)
if hd < hs_h - sr_h * wd:
w_min = max(wd, (hs_h - hn) / sr_h)
w_max = min((hs_h - hd) / sr_h, wn)
h_min = hs_h - sr_h * w_max
h_max = hs_h - sr_h * w_min
whuc = whu if hhu % 2 else whu - 1
hhuc = hhu - 1
ps = depth * (hu * s_width + wu)
ph = depth * (hhu * h_width + whu)
phc = depth * (hhuc * h_width + whuc)
if h_min < hn:
area = wmr_min * (h_min - hd) + ((w_max - w_min) *
(h_max - h_min)) / 2
if phc >= 0:
areas.append(area)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhuc, whuc, :],
method))
areas.append(wmr_min * hmr_min - area)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhu, whu, :],
method))
else:
if phc >= 0:
areas.append(wmr_min * hmr_min)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhuc, whuc, :],
method))
elif hd < hs_h - sr_h * (h_dist_w - wn):
w_min = max(wd, (hd + hs_h) / sr_h)
w_max = min((hn + hs_h) / sr_h, wn)
h_min = hs_h - sr_h * (h_dist_w - w_min)
h_max = hs_h - sr_h * (h_dist_w - w_max)
whuc = whu + 1 if hhu % 2 else whu
hhuc = hhu - 1
ps = depth * (hu * s_width + wu)
ph = depth * (hhu * h_width + whu)
phc = depth * (hhuc * h_width + whuc)
if h_min < hn:
area = wmr_min * (h_min - hd) + ((w_max - w_min) *
(h_max - h_min)) / 2
if phc >= 0 and phc < h_size:
areas.append(area)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhuc, whuc, :],
method))
areas.append(wmr_min * hmr_min - area)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhu, whu, :],
method))
else:
if phc >= 0 and phc < h_size:
areas.append(wmr_min * hmr_min)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhuc, whuc, :],
method))
else:
ps = depth * (hu * s_width + wu)
ph = depth * (hhu * h_width + whu)
areas.append(wmr_min * hmr_min)
diffs.append(
pixels_diff(s[:, hu, wu, :], h[:, hhu, whu, :], method))
wi = max(wi + wmr_min, np.nextafter(wi + wmr_min,
wi + wmr_min + 1))
wi = 0.0
hi = max(hi + hmr_min, np.nextafter(hi + hmr_min, hi + hmr_min + 1))
if not sh_is_tensor:
areas = np.asarray(areas)
diffs = np.sum(np.asarray(diffs), axis=-1)
result = np.sum(np.multiply(areas, diffs), axis=-1)
else:
areas = tf.convert_to_tensor(areas, dtype=tf.float32)
diffs = tf.reduce_sum(tf.stack(diffs), axis=-1)
result = tf.reduce_sum(tf.multiply(areas, diffs), axis=-1)
if method == _COMPARE_MAE or \
method == _COMPARE_MSE or \
method == _COMPARE_RMSE or \
method == _COMPARE_PSNR :
result /= s_size
if method == _COMPARE_RMSE:
result = math.sqrt(result)
elif method == _COMPARE_PSNR:
if not sh_is_tensor:
result = 10 * math.log(65025 / result)
else:
result = 10 * tf.math.log(1 / result)
return result
def convert_to_tensor(t):
return tf.convert_to_tensor(t) if not tf.is_tensor(t) else t
def has_tensors(*x):
return np.any(
[tf.is_tensor(t) for t in tf.nest.flatten(x, expand_composites=True)])
def assert_matching_dtypes_and_inner_shapes(tensors,
specs,
caller,
tensors_name,
specs_name,
allow_extra_fields=False):
"""Returns `True` if tensors and specs have matching dtypes and inner shapes.
Args:
tensors: A nest of tensor objects.
specs: A nest of `tf.TypeSpec` objects.
caller: The object calling `assert...`.
tensors_name: (str) Name to use for the tensors in case of an error.
specs_name: (str) Name to use for the specs in case of an error.
allow_extra_fields: If `True`, then `tensors` may contain more keys or list
fields than strictly required by `specs`.
Raises:
ValueError: If the tensors do not match the specs' dtypes or their inner
shapes do not match the specs' shapes.
"""
if allow_extra_fields:
tensors = prune_extra_keys(specs, tensors)
assert_same_structure(
tensors,
specs,
message=('{}: {} and {} do not have matching structures'.format(
caller, tensors_name, specs_name)))
flat_tensors = nest.flatten(tensors)
flat_specs = tf.nest.flatten(specs)
flat_tensors = [
tf.convert_to_tensor(t, dtype_hint=s.dtype)
if not tf.is_tensor(t) else t
for (t, s) in zip(flat_tensors, flat_specs)
]
tensor_shapes = [t.shape for t in flat_tensors]
tensor_dtypes = [t.dtype for t in flat_tensors]
spec_shapes = [spec_shape(s) for s in flat_specs]
spec_dtypes = [t.dtype for t in flat_specs]
compatible = True
if any(s_dtype != t_dtype
for s_dtype, t_dtype in zip(spec_dtypes, tensor_dtypes)):
compatible = False
else:
for s_shape, t_shape in zip(spec_shapes, tensor_shapes):
if s_shape.ndims in (0, None) or t_shape.ndims is None:
continue
if s_shape.ndims > t_shape.ndims:
compatible = False
break
if not s_shape.is_compatible_with(t_shape[-s_shape.ndims:]):
compatible = False
break
if not compatible:
get_dtypes = lambda v: tf.nest.map_structure(lambda x: x.dtype, v)
get_shapes = lambda v: tf.nest.map_structure(spec_shape, v)
raise ValueError(
'{}: Inconsistent dtypes or shapes between {} and {}.\n'
'dtypes:\n{}\nvs.\n{}.\n'
'shapes:\n{}\nvs.\n{}.'.format(caller, tensors_name, specs_name,
get_dtypes(tensors),
get_dtypes(specs),
get_shapes(tensors),
get_shapes(specs)))
def call(
self,
input_ids=None,
past_key_values=None,
attention_mask=None,
token_type_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
use_cache=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
labels=None,
training=False,
**kwargs,
):
r"""
labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,
config.vocab_size - 1]`.
"""
transformer_outputs = self.transformer(
input_ids=input_ids,
past_key_values=past_key_values,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
training=training,
)
hidden_states = transformer_outputs[0]
logits = self.classifier(hidden_states)
in_logits = None
if self.config.pad_token_id is None:
sequence_lengths = -1
else:
if input_ids is not None:
sequence_lengths = (tf.reduce_sum(
tf.cast(
tf.math.not_equal(input_ids, self.config.pad_token_id),
dtype=input_ids.dtype,
),
-1,
keepdims=False,
) - 1)
in_logits = tf.gather(logits,
sequence_lengths,
batch_dims=1,
axis=1)
else:
sequence_lengths = -1
logger.warning(
f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
)
loss = None
if labels is not None:
if input_ids is not None:
batch_size, sequence_length = shape_list(input_ids)[:2]
else:
batch_size, sequence_length = shape_list(inputs_embeds)[:2]
assert (
self.config.pad_token_id is not None or batch_size == 1
), "Cannot handle batch sizes > 1 if no padding token is defined."
if not tf.is_tensor(sequence_lengths):
in_logits = logits[0:batch_size, sequence_lengths]
loss = self.hf_compute_loss(
tf.reshape(labels, [-1, 1]),
tf.reshape(in_logits, [-1, self.num_labels]))
pooled_logits = in_logits if in_logits is not None else logits
if not return_dict:
output = (pooled_logits, ) + transformer_outputs[1:]
return ((loss, ) + output) if loss is not None else output
return TFSequenceClassifierOutput(
loss=loss,
logits=pooled_logits,
hidden_states=transformer_outputs.hidden_states,
attentions=transformer_outputs.attentions,
)
def sum(x, axis=None, keepdims=False, name=None):
if not tf.is_tensor(x):
x = tf.convert_to_tensor(x)
if x.dtype == bool:
x = cast(x, int32)
return tf.reduce_sum(x, axis, keepdims, name)
def as_numpy(x):
if tf.is_tensor(x):
return x.numpy()
else:
return x
def test_dice_loss(tensor_true, tensor_pred):
assert tf.is_tensor(dice_loss(tensor_true, tensor_pred))
