def __call__(self, step): starting_iteration = self.steps_per_epoch * self.start_epoch starting_iteration = tf.cast(starting_iteration, self.dtype) global_step = tf.cast(step, self.dtype) recomp_iteration = global_step - starting_iteration + 1. decayed_coeff = self.coeff_scheduler(recomp_iteration) # This is an autograph-friendly alternative to checking Tensorflow booleans # in eager mode. scale = tf.minimum( tf.maximum(tf.cast(recomp_iteration, self.dtype), 0.), 1.) return scale * decayed_coeff
def while_loop_body(iteration, distribution, inactive, old_inactive): """Performs one iteration of the projection.""" del old_inactive # Needed by the condition, but not the body (for lint). iteration += 1 scale = (1.0 - tf.reduce_sum( distribution, axis=0, keepdims=True)) / tf.maximum( 1.0, tf.reduce_sum(inactive, axis=0, keepdims=True)) distribution = distribution + (scale * inactive) new_inactive = tf.cast(distribution > 0, distribution.dtype) distribution = distribution * new_inactive return (iteration, distribution, new_inactive, inactive)
def parameterized_infection_dynamics(_, previous_state, parameters): new_infections = tfd.Poisson( parameters['infection_rate'] * previous_state['infected'] * previous_state['susceptible'] / population_size) new_recoveries = tfd.Poisson(previous_state['infected'] * parameters['recovery_rate']) return tfd.JointDistributionNamed({ 'new_infections': new_infections, 'new_recoveries': new_recoveries, 'susceptible': lambda new_infections: tfd.Deterministic( # pylint: disable=g-long-lambda tf.maximum(0., previous_state['susceptible'] - new_infections)), 'infected': lambda new_infections, new_recoveries: tfd.Deterministic( # pylint: disable=g-long-lambda tf.maximum(0., (previous_state['infected'] + new_infections - new_recoveries))) })
def random_brightness(image, max_delta, impl='simclrv2'): """A multiplicative vs additive change of brightness.""" if impl == 'simclrv2': factor = tf.random.uniform([], tf.maximum(1.0 - max_delta, 0), 1.0 + max_delta) image = image * factor elif impl == 'simclrv1': image = tf.image.random_brightness(image, max_delta=max_delta) else: raise ValueError('Unknown impl {} for random brightness.'.format(impl)) return image
def _log_unnormalized_prob(self, x): # The log-probability at negative points is always -inf. # Catch such x's and set the output value accordingly. safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x), 0.) y = safe_x * self.log_rate - tf.math.lgamma(1. + safe_x) is_supported = tf.broadcast_to(tf.equal(x, safe_x), tf.shape(input=y)) neg_inf = tf.fill(tf.shape(input=y), value=np.array(-np.inf, dtype=y.dtype.as_numpy_dtype)) return tf.where(is_supported, y, neg_inf)
def _sample_dataset(self, seed): dataset = dict( train_locations=self._train_locations, train_extents=self._train_extents, ) prior_samples = self.prior_distribution().sample(seed=seed) observation_noise_dist = self._observation_noise_fn( prior_samples['log_intensity']) counts = observation_noise_dist.sample(seed=seed) dataset['train_counts'] = tf.maximum(1, tf.cast(counts, tf.int32)) return dataset
def _get_reinterpreted_batch_ndims(self, distribution_batch_shape_tensor=None): if self._static_reinterpreted_batch_ndims is not None: return self._static_reinterpreted_batch_ndims if self._reinterpreted_batch_ndims is not None: return tf.convert_to_tensor(self._reinterpreted_batch_ndims) if distribution_batch_shape_tensor is None: distribution_batch_shape_tensor = self.distribution.batch_shape_tensor( ) return tf.maximum(0, tf.size(distribution_batch_shape_tensor) - 1)
def create_masks(inputs, target): """function that creates all masks for training/validation""" inputs = tf.cast(tf.math.equal(inputs, 0), tf.float32) target = tf.cast(tf.math.equal(target, 0), tf.float32) encoder_mask = inputs[:, tf.newaxis, tf.newaxis, :] decoder_mask = inputs[:, tf.newaxis, tf.newaxis, :] decoder_target_mask = target[:, tf.newaxis, tf.newaxis, :] look_ahead_mask = 1 - tf.linalg.band_part( tf.ones((target.shape[0], 1, target.shape[1], target.shape[1])), -1, 0) combined_mask = tf.maximum(decoder_target_mask, look_ahead_mask) return encoder_mask, combined_mask, decoder_mask
def valid_transition_fn(_, particles): return tfd.JointDistributionNamedAutoBatched( { 'sales': tfd.Poisson(10. * tf.ones_like(particles['inventory'])), 'inventory': lambda sales: tfd.Deterministic( # pylint: disable=g-long-lambda tf.maximum(0., particles['inventory'] - sales)) }, batch_ndims=1, validate_args=True)
def update_step(self, gradient, variable): """Update step given gradient and the associated model variable.""" beta_1_power = None beta_2_power = None lr = tf.cast(self.learning_rate, variable.dtype) local_step = tf.cast(self.iterations + 1, variable.dtype) beta_1_power = tf.pow(tf.cast(self.beta_1, variable.dtype), local_step) beta_2_power = tf.pow(tf.cast(self.beta_2, variable.dtype), local_step) var_key = self._var_key(variable) m = self._momentums[self._index_dict[var_key]] v = self._velocities[self._index_dict[var_key]] alpha = lr * tf.sqrt(1 - beta_2_power) / (1 - beta_1_power) if isinstance(gradient, tf.IndexedSlices): # Sparse gradients. m.assign_add(-m * (1 - self.beta_1)) m.scatter_add( tf.IndexedSlices(gradient.values * (1 - self.beta_1), gradient.indices)) v.assign_add(-v * (1 - self.beta_2)) v.scatter_add( tf.IndexedSlices( tf.square(gradient.values) * (1 - self.beta_2), gradient.indices, )) if self.amsgrad: v_hat = self._velocity_hats[self._index_dict[var_key]] v_hat.assign(tf.maximum(v_hat, v)) v = v_hat variable.assign_sub((m * alpha) / (tf.sqrt(v) + self.epsilon)) else: # Dense gradients. m.assign_add((gradient - m) * (1 - self.beta_1)) v.assign_add((tf.square(gradient) - v) * (1 - self.beta_2)) if self.amsgrad: v_hat = self._velocity_hats[self._index_dict[var_key]] v_hat.assign(tf.maximum(v_hat, v)) v = v_hat variable.assign_sub((m * alpha) / (tf.sqrt(v) + self.epsilon))
def create_masks(inputs, target): """Function that creates all masks for training/validation""" enc_padding_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) enc_padding_mask = enc_padding_mask[:, tf.newaxis, tf.newaxis, :] dec_padding_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) dec_padding_mask = dec_padding_mask[:, tf.newaxis, tf.newaxis, :] size = target.shape[1] look_ahead_mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0) dec_target_p_mask = tf.cast(tf.math.equal(target, 0), tf.float32) dec_target_p_mask = dec_target_p_mask[:, tf.newaxis, tf.newaxis, :] combined_mask = tf.maximum(dec_target_p_mask, look_ahead_mask) return (enc_padding_mask, combined_mask, dec_padding_mask)
def transition_fn_no_batch_shape(_, particles): # Autobatched JD defaults to treating num_particles as event shape, but # we need it to be batch shape to get per-particle logprobs. return tfd.JointDistributionNamedAutoBatched( { 'sales': tfd.Poisson(10. * tf.ones_like(particles['inventory'])), 'inventory': lambda sales: tfd.Deterministic( # pylint: disable=g-long-lambda tf.maximum(0., particles['inventory'] - sales)) }, validate_args=True)
def transition_fn_partial_batch_shape(_, particles): return tfd.JointDistributionNamed( # Using `Sample` ensures iid proposals for each particle, but not # per-particle log probs. { 'sales': tfd.Sample(tfd.Poisson(10.), ps.shape(particles['sales'])), 'inventory': lambda sales: tfd.Deterministic( # pylint: disable=g-long-lambda tf.maximum(0., particles['inventory'] - sales)) }, validate_args=True)
def random_resize(features, scale=(0.5, 2.0), ensure_small=None, keys=("image", ), methods=("bilinear", )): """Randomly resize the image and label by a uniformly sampled scale. Args: features: Input dictionary containing "image", "label", and other keys. scale: Output image and label will be scaled by a scale sampled uniformly at random in this range. ensure_small: Ignored if None. Else, if input image size * min(scale) is less than ensure_small, it will adjust the scale so that the output image is always at least as big as ensure_small. This is useful so that subsequent crop operations do not go out of range. keys: Keys to apply resize op to. Note that keys starting with prefix "label" will be resized using nearest neighbour. methods: Resize methods per key. Returns: features with randomly scaled "images" defined by keys. """ if ensure_small is None: scale_min, scale_max = scale else: width, height = _get_image_size(features["image"], dynamic=True) scale_min = tf.maximum( ensure_small / tf.cast(width, dtype=tf.float32), ensure_small / tf.cast(height, dtype=tf.float32)) scale_max = tf.maximum(scale[1], scale_min) scale_chosen = tf.random.uniform(shape=(), minval=scale_min, maxval=scale_max, dtype=tf.float32) width, height = _get_image_size(features["image"], dynamic=True) new_width = tf.cast(tf.cast(width, tf.float32) * scale_chosen, tf.int32) new_height = tf.cast(tf.cast(height, tf.float32) * scale_chosen, tf.int32) return resize(features, (new_height, new_width), keys, methods)
def pad_reflecting(x, padding_below, padding_above, axis): """Pads `x` with reflecting conditions above and/or below it along some axis. Pads `x` with reflecting conditions for `padding_below` entries below the tensor and `padding_above` entries above the tensor in the direction along `axis`. This is like using tf.pad(x, --, 'REFLECT'), except that this code allows for an unbounded number of reflections while tf.pad() only supports one reflection. Multiple reflections are necessary for for wavelet decompositions to guard against cases where the wavelet filters are larger than the input tensor along `axis`, which happens often at coarse scales. Note that "reflecting" boundary conditions are different from "symmetric" boundary conditions, in that it doesn't repeat the last element: reflect([A, B, C, D], 2) = [C, B, A, B, C, D, C, B] symmet.([A, B, C, D], 2) = [B, A, A, B, C, D, D, C] Args: x: The tensor to be padded with reflecting boundary conditions. padding_below: The number of elements being padded below the tensor. padding_above: The number of elements being padded above the tensor. axis: The axis in x in which padding will be performed. Returns: `x` padded according to `padding_below` and `padding_above` along `axis` with reflecting boundary conditions. """ if not isinstance(padding_below, int): raise ValueError( 'Expected `padding_below` of type int, but is of type {}'.format( type(padding_below))) if not isinstance(padding_above, int): raise ValueError( 'Expected `padding_above` of type int, but is of type {}'.format( type(padding_above))) if not isinstance(axis, int): raise ValueError( 'Expected `axis` of type int, but is of type {}'.format( type(axis))) if not (axis >= 0 and axis < len(x.shape)): raise ValueError('Expected `axis` in [0, {}], but is = {}'.format( len(x.shape) - 1, axis)) if padding_below == 0 and padding_above == 0: return tf.convert_to_tensor(x) n = tf.shape(x)[axis] # `i' contains the indices of the output padded tensor in the frame of # reference of the input tensor. i = tf.range(-padding_below, n + padding_above, dtype=tf.int32) # `j` contains the indices of the input tensor corresponding to the output # padded tensor. i_mod = tf.math.mod(i, tf.maximum(1, 2 * (n - 1))) j = tf.minimum(2 * (n - 1) - i_mod, i_mod) return tf.gather(x, j, axis=axis)
def call(self, inputs): inputs = tf.convert_to_tensor(inputs) if inputs.shape.rank == 1: inputs = tf.compat.v1.expand_dims(inputs, 1) # If the inputs are not floats, cast them to floats. This avoids issues # with int-float multiplication and division below. if inputs.dtype != K.floatx(): inputs = tf.cast(inputs, K.floatx()) # We need to reshape the mean and variance data to ensure that Tensorflow # broadcasts the data correctly. mean = tf.reshape(self.mean, self._broadcast_shape) variance = tf.reshape(self.variance, self._broadcast_shape) return ((inputs - mean) / tf.maximum(tf.sqrt(variance), K.epsilon()))
def pairwise_l2_distance(embs1, embs2): """Computes pairwise distances between all rows of embs1 and embs2.""" norm1 = tf.reduce_sum(tf.square(embs1), 1) norm1 = tf.reshape(norm1, [-1, 1]) norm2 = tf.reduce_sum(tf.square(embs2), 1) norm2 = tf.reshape(norm2, [1, -1]) # Max to ensure matmul doesn't produce anything negative due to floating # point approximations. dist = tf.maximum( norm1 + norm2 - 2.0 * tf.matmul(embs1, embs2, False, True), 0.0) return dist
def logmap0(y, c): """Hyperbolic logarithmic map at zero in the Poincare ball model. Args: y: Tensor of size B x dimension representing hyperbolic points c: Tensor of size 1 representing the absolute hyperbolic curvature. Returns: Tensor of shape B x dimension. """ sqrt_c = tf.sqrt(c) y_norm = tf.maximum(tf.norm(y, axis=-1, keepdims=True), MIN_NORM) return y / y_norm / sqrt_c * artanh(sqrt_c * y_norm)
def _calc_loudness(audio, n_fft=2048, top_db=200.0, pmin=1e-20): """Perceptual loudness in tf, following librosa implementation.""" librosa = tfds.core.lazy_imports.librosa log10 = lambda x: tf.math.log(x) / tf.math.log(10.0) spectra = tf.signal.stft(signals=audio, frame_length=n_fft, frame_step=int(_AUDIO_RATE // _F0_AND_LOUDNESS_RATE), fft_length=n_fft, pad_end=True) power = tf.abs(spectra)**2.0 power_db = 10.0 * log10(tf.maximum(pmin, power)) power_db = tf.maximum(power_db, tf.reduce_max(power_db) - top_db) fft_frequencies = librosa.fft_frequencies(n_fft=n_fft) a_weighting = librosa.A_weighting(fft_frequencies) loudness = power_db + a_weighting[tf.newaxis, :] loudness = tf.reduce_mean(loudness, axis=-1) return loudness
def clip_to_window(boxlist, window, filter_nonoverlapping=True, scope=None): """Clip bounding boxes to a window. This op clips any input bounding boxes (represented by bounding box corners) to a window, optionally filtering out boxes that do not overlap at all with the window. Args: boxlist: BoxList holding M_in boxes window: a tensor of shape [4] representing the [y_min, x_min, y_max, x_max] window to which the op should clip boxes. filter_nonoverlapping: whether to filter out boxes that do not overlap at all with the window. scope: name scope. Returns: a BoxList holding M_out boxes where M_out <= M_in """ with tf.name_scope(scope, 'ClipToWindow'): y_min, x_min, y_max, x_max = tf.split(value=boxlist.get(), num_or_size_splits=4, axis=1) win_y_min, win_x_min, win_y_max, win_x_max = tf.unstack(window) y_min_clipped = tf.maximum(tf.minimum(y_min, win_y_max), win_y_min) y_max_clipped = tf.maximum(tf.minimum(y_max, win_y_max), win_y_min) x_min_clipped = tf.maximum(tf.minimum(x_min, win_x_max), win_x_min) x_max_clipped = tf.maximum(tf.minimum(x_max, win_x_max), win_x_min) clipped = box_list.BoxList( tf.concat( [y_min_clipped, x_min_clipped, y_max_clipped, x_max_clipped], 1)) clipped = _copy_extra_fields(clipped, boxlist) if filter_nonoverlapping: areas = area(clipped) nonzero_area_indices = tf.cast( tf.reshape(tf.where(tf.greater(areas, 0.0)), [-1]), tf.int32) clipped = gather(clipped, nonzero_area_indices) return clipped
def _values_transform_fn(t, grid, value_grid): zero = tf.zeros_like(value_grid) is_ex_time = _is_exercise_time(t) def _at_least_one_swaption_pays(): payoff_swap = tf.nn.relu(_get_swap_payoff(t)) return tf.where( tf.reshape(is_ex_time, tf.concat([batch_shape, [1] * dim], axis=0)), payoff_swap, zero) v_star = tf.cond(tf.math.reduce_any(is_ex_time), _at_least_one_swaption_pays, lambda: zero) return grid, tf.maximum(value_grid, v_star)
def get_linear_warmup_rsqrt_decay_lr(init_lr, hidden_size, num_warmup_steps): """Calculate learning rate with linear warmup and rsqrt decay.""" num_warmup_steps = tf.cast(num_warmup_steps, tf.float32) global_step = tf.compat.v1.train.get_or_create_global_step() global_step = tf.cast(global_step, tf.float32) learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32) learning_rate *= tf.math.rsqrt(tf.cast(hidden_size, tf.float32)) # Apply linear warmup learning_rate *= tf.minimum(1.0, global_step / num_warmup_steps) # Apply rsqrt decay learning_rate *= tf.math.rsqrt(tf.maximum(global_step, num_warmup_steps)) return learning_rate
def compute_power(audio, sample_rate=16000, frame_rate=250, frame_size=1024, range_db=LD_RANGE, ref_db=20.7): """Compute power of audio in dB.""" # TODO(hanoih@): enable `use_tf` to be True or False like `compute_loudness` rms_energy = compute_rms_energy(audio, sample_rate, frame_rate, frame_size) power_db = amplitude_to_db(rms_energy**2, use_tf=True) # Set dynamic range. power_db -= ref_db power_db = tf.maximum(power_db, -range_db) return power_db
def create_masks(inputs, target): """function""" encoder_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) encoder_mask = encoder_mask[:, tf.newaxis, tf.newaxis, :] decoder_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) decoder_mask = decoder_mask[:, tf.newaxis, tf.newaxis, :] batch_size, seq_len_out = target.shape look_ahead_mask = tf.linalg.band_part(tf.ones((seq_len_out, seq_len_out)), -1, 0) look_ahead_mask = 1 - look_ahead_mask padding_mask = tf.cast(tf.math.equal(target, 0), tf.float32) padding_mask = padding_mask[:, tf.newaxis, tf.newaxis, :] combined_mask = tf.maximum(look_ahead_mask, padding_mask) return encoder_mask, combined_mask, decoder_mask
def _crop(x, max_length, sample): """Select (optionally random) crop from sequence.""" if sample: # Optionally sample random starting position. start = tf.random.uniform( (), dtype=tf.int32, maxval=tf.maximum(1, tf.shape(x)[0] - max_length + 1)) else: start = 0 x = x[start:(start + max_length)] return x
def create_masks(inputs, target): """ Returns: encoder_mask, look_ahead_mask, decoder_mask """ encoder_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) encoder_mask = encoder_mask[:, tf.newaxis, tf.newaxis, :] decoder_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) decoder_mask = decoder_mask[:, tf.newaxis, tf.newaxis, :] size = target.shape[1] look_ahead_mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0) dec_target_mask = tf.cast(tf.math.equal(target, 0), tf.float32) dec_target_mask = dec_target_mask[:, tf.newaxis, tf.newaxis, :] combined_mask = tf.maximum(dec_target_mask, look_ahead_mask) return (encoder_mask, combined_mask, decoder_mask)
def create_masks(inputs, target): """creates all masks for training/validation""" size = tf.shape(target)[1] encoder_mask = tf.cast(tf.math.equal(inputs, 0), tf.float32) encoder_mask = encoder_mask[:, tf.newaxis, tf.newaxis, :] look_ahead_mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0) dec_target_padd_mask = tf.cast(tf.math.equal(target, 0), tf.float32) dec_target_padd_mask = dec_target_padd_mask[:, tf.newaxis, tf.newaxis, :] look_ahead_mask = tf.maximum(dec_target_padd_mask, look_ahead_mask) return encoder_mask, look_ahead_mask, encoder_mask
def berp(global_step, start_step, end_step, start_val, end_val, alpha=5): """Beta interpolation.""" beta_dist = tfd.Beta(alpha, alpha) mode = beta_dist.mode() interp = (tf.cast(global_step - start_step, tf.float32) / tf.cast(end_step - start_step, tf.float32)) interp = tf.maximum(0.0, tf.minimum(1.0, interp)) interp = tf.where(tf.math.is_nan(interp), tf.zeros_like(interp), interp) interp *= mode val = beta_dist.prob(interp) val /= beta_dist.prob(mode) val *= (end_val - start_val) val += start_val return val
def _pad_multichannel(image, ensure_small, pad_value, mode): """Pad to ensure `ensure_small`.""" pad_h = tf.maximum(ensure_small[0] - tf.shape(image)[0], 0) pad_h_l = pad_h // 2 pad_h_r = pad_h - pad_h_l pad_w = tf.maximum(ensure_small[1] - tf.shape(image)[1], 0) pad_w_l = pad_w // 2 pad_w_r = pad_w - pad_w_l def pad_2d(x, v): """Pad 2D input `x` with constant value `v`.""" return tf.pad( x, [[pad_h_l, pad_h_r], [pad_w_l, pad_w_r]] + [[0, 0]] * (len(x.shape) - 2), mode=mode, constant_values=v) if isinstance(pad_value, (list, tuple)): image_new = tf.stack( [pad_2d(image[:, :, i], v) for i, v in enumerate(pad_value)], axis=2) else: image_new = pad_2d(image, pad_value) return image_new
def get_max_num_levels(sz): """Returns the maximum number of levels that construct() can support. Args: sz: A tuple of ints representing some input size (batch, width, height). Returns: The maximum value for num_levels, when calling construct(im, num_levels), assuming `sz` is the shape of `im`. """ min_sz = tf.minimum(sz[1], sz[2]) log2 = lambda x: tf.math.log(tf.cast(x, tf.float32)) / tf.math.log(2.) max_num_levels = tf.cast(tf.math.ceil(log2(tf.maximum(1, min_sz))), tf.int32) return max_num_levels