def feature_model_fn(features, params): band_settings = band_settings_params.BandSettings.from_params(params) per_band_data = band_settings.per_band_sub_model_fn(per_band_model_fn, features, params=params) result = tf.stack(per_band_data, axis=1) graph_typecheck.assert_shape(result, [params["batch_size"], band_settings.nbands]) return {"is_max_soft": result}
def score(self): left = self._normalize(self.left) right = self._normalize(self.right) dotprod = tf.reduce_sum(left * right, axis=1) graph_typecheck.assert_shape(dotprod, [self.batch_size]) # Normalize output to between 0 and 1. result = 0.5 * (dotprod + 1.0) graph_typecheck.assert_shape(result, [self.batch_size]) return result
def loss(self, labels: tf.Tensor, score: tf.Tensor = None): """Computes the log loss of the scores. :param labels: Labels tensor :param score: Result from `self.score()`, if reusing such computation is desired. :return: Tensor with the loss. """ graph_typecheck.assert_shape(labels, [self.batch_size]) score = self.score() if score is None else score return tf.losses.log_loss(labels=labels, predictions=score, epsilon=1e-5)
def __init__(self, left: tf.Tensor, right: tf.Tensor, *, batch_size: int, hidden_size: int, norm: bool = True): graph_typecheck.assert_shape(left, [batch_size, hidden_size]) graph_typecheck.assert_shape(right, [batch_size, hidden_size]) self.left = left self.right = right self.batch_size = batch_size self.hidden_size = hidden_size self.norm = norm
def initial_layer_binned( initial_layer_features: tf.Tensor, cutoff_data: CutoffData, band: str, soft_onehot: Nonlinearity = Nonlinearity.SIGMOID ): batch_size, twice_window_size, channels = map(int, initial_layer_features.shape) nonlinearity = nonlinearity_fcn(soft_onehot) if channels == 3: scales = cutoff_data.dflux_dt_dflux_dtime_scales(band) cutoffs = cutoff_data.dflux_dt_dflux_dtime_cutoffs(band) cutoffs_batch_window = tf.expand_dims(tf.expand_dims(cutoffs, 0), 0) scales_batch_window = tf.expand_dims( tf.expand_dims(tf.expand_dims(scales, 0), 0), -1 ) init_layer_per_cutoff = tf.expand_dims(initial_layer_features, -1) graph_typecheck.assert_shape( cutoffs_batch_window, [1, 1, channels, cutoff_data.embedding_size] ) graph_typecheck.assert_shape(scales_batch_window, [1, 1, channels, 1]) graph_typecheck.assert_shape( init_layer_per_cutoff, [batch_size, twice_window_size, channels, 1] ) result = nonlinearity( (init_layer_per_cutoff - cutoffs_batch_window) / scales_batch_window ) return graph_typecheck.assert_shape( result, [batch_size, twice_window_size, channels, cutoff_data.embedding_size] ) else: raise NotImplementedError(f"{channels}-size data not implemented.")
def masked( self, expanded_tensor: tf.Tensor, value_if_masked: float, expected_extra_dims: typing.List[int] ): """Masks a tensor which was calculated from dflux_dt. :param expanded_tensor: <float>[batch_size, window_size, ...] Tensor with first dimensions being batch_size and window_size. :param value_if_masked: Value to fill for masked positions. :param expected_extra_dims: Expected extra dimensions. :returns: Tensor of same shape as expanded_tensor, but with `value_if_masked` filled in masked dimensions. """ mask_shape = list(map(int, self.mask.shape)) graph_typecheck.assert_shape(expanded_tensor, mask_shape + expected_extra_dims) value_if_masked = expanded_tensor.dtype.as_numpy_dtype(value_if_masked) if_masked_tensor = tf.fill(expanded_tensor.shape, value_if_masked) mask = self.mask for i in range(2, 2 + len(expected_extra_dims)): mask = tf.expand_dims(mask, axis=i) mask = tf.tile(mask, [1, 1] + expected_extra_dims) return tf.where(mask, expanded_tensor, if_masked_tensor)
def per_side_model(side_features, params): inputs = dense_extracted_features.feature_model_fn(side_features, params=params) # Channels is dflux/dt, etc. graph_typecheck.assert_shape( inputs, [batch_size, twice_window_size, channels, embedding_size, nbands] ) first_layer_type = params["first_layer"] if first_layer_type == "conv": inputs_shuffled = tf.transpose(inputs, perm=[0, 1, 4, 2, 3]) graph_typecheck.assert_shape( inputs_shuffled, [batch_size, twice_window_size, nbands, channels, embedding_size] ) conv_result = input_conv( tf.reshape(inputs_shuffled, [batch_size, twice_window_size, nbands, -1]) ) curr_layer = tf.reshape(conv_result, [batch_size, -1]) elif first_layer_type == "none": curr_layer = tf.reshape(inputs, [batch_size, -1]) else: raise ValueError(f"No known first_layer type {first_layer_type}") return per_side_model_keras(curr_layer)
def feature_model_fn(features, params): band_settings = band_settings_params.BandSettings.from_params(params) per_band_data = band_settings.per_band_sub_model_fn_with_band_name( per_band_model_fn, features, params=params, value_if_masked=params.get("value_if_masked", 0.0), soft_onehot=Nonlinearity[params.get("input_soft_onehot", "sigmoid").upper()] ) return graph_typecheck.assert_shape( tf.stack(per_band_data, axis=4), [ params["batch_size"], 2 * params["window_size"], 3, cutoff_data_for_window_size(params["window_size"]).embedding_size, band_settings.nbands, ] )
def batch_2win_shaped(t): return graph_typecheck.assert_shape(t, [batch_size, 2 * window_size])
def batch_shaped(t): return graph_typecheck.assert_shape(t, [batch_size])
def per_band_model_fn(band_features, params, debug_print=False): # NOTE(gatoatigrado): dense_extracted_features.WindowFeatures provides a convenient # API for many calculations here, but right now the unit test data for max_model_kernel_test # does not provide time tensors, making it inconvenient to modernize this code. batch_size = params["batch_size"] window_size = params["window_size"] inv_eps = 1.0 / params["flux_scale_epsilon"] graph_typecheck.assert_shape(band_features["before_padding"], [batch_size]) graph_typecheck.assert_shape(band_features["after_padding"], [batch_size]) closest_flux = graph_typecheck.assert_shape( band_features["closest_flux_in_band"], [batch_size]) before_flux = graph_typecheck.assert_shape(band_features["before_flux"], [batch_size, window_size]) after_flux = graph_typecheck.assert_shape(band_features["after_flux"], [batch_size, window_size]) # Return a soft-greater-than operator, product that all scores are greater. is_greater_than_before = masked_sigmoid( inv_eps * (tf.expand_dims(closest_flux, axis=1) - before_flux), _left_mask(band_features["before_padding"], window_size)) is_greater_than_after = masked_sigmoid( inv_eps * (tf.expand_dims(closest_flux, axis=1) - after_flux), _right_mask(band_features["after_padding"], window_size)) graph_typecheck.assert_shape(is_greater_than_before, [batch_size, window_size]) graph_typecheck.assert_shape(is_greater_than_after, [batch_size, window_size]) if debug_print: is_greater_than_before = graph_typecheck.print_single( is_greater_than_before, "is_greater_than_before:") is_greater_than_after = graph_typecheck.print_single( is_greater_than_after, "is_greater_than_after:") return tf.reduce_prod(is_greater_than_before, axis=1) * tf.reduce_prod( is_greater_than_after, axis=1)