def determine_loss_scales( snapshots: np.ndarray, hparams: tf.contrib.training.HParams) -> Tuple[np.ndarray, np.ndarray]: """Determine scale factors for the loss. When passed into model.compute_loss, predictions of all zero should result in a loss of 1.0 when averaged over the full dataset. Args: snapshots: np.ndarray with shape [examples, x] with high-resolution training data. hparams: hyperparameters to use for training. Returns: Tuple of two numpy arrays: error_scale: array with dimensions [2, derivative] indicating the scaling in the loss to use on squared error and relative squared error for each derivative target. error_floor: numpy array with scale for weighting of relative errors. """ with tf.Graph().as_default(): dataset = model.make_dataset(snapshots, hparams, repeat=False) data = load_dataset(dataset) baseline_error = (data['labels'] - data['baseline'])**2 percentile = 100 * hparams.error_floor_quantile error_floor = np.maximum( np.percentile(baseline_error, percentile, axis=(0, 1)), 1e-12) # predict zero for all derivatives, and a constant value for the integrated # solution over time. equation_type = equations.equation_type_from_hparams(hparams) num_zero_predictions = len(equation_type.DERIVATIVE_ORDERS) + 1 labels_shape = data['labels'].shape predictions = np.concatenate([ np.zeros(labels_shape[:-1] + (num_zero_predictions, )), np.repeat(data['inputs'][..., np.newaxis], labels_shape[-1] - num_zero_predictions, axis=-1) ], axis=-1) components = np.stack( model.abs_and_rel_error(predictions=predictions, labels=data['labels'], baseline=data['baseline'], error_floor=error_floor)) baseline_error = np.mean(components, axis=(1, 2)) logging.info('baseline_error: %s', baseline_error) error_scale = np.where(baseline_error > 0, 1.0 / baseline_error, 0) return error_floor, error_scale
def weighted_loss(normalized_loss_per_head: tf.Tensor, hparams: tf.contrib.training.HParams) -> tf.Tensor: """Calculate overall training loss. Weights are normalized to sum to 1.0 (`relative_error+absolute_error` and `space_derivatives_weight+time_derivatives_weight+integrated_solution_weight`) before being used. Args: normalized_loss_per_head: tensor with dimensions [abs/rel error, channel]. hparams: model hyperparameters. Returns: Scalar float32 Tensor indicating the loss. """ # dimensions [abs/rel error] abs_rel_weights = tf.convert_to_tensor( [hparams.absolute_error_weight, hparams.relative_error_weight]) abs_rel_weights /= tf.reduce_sum(abs_rel_weights) equation_type = equations.equation_type_from_hparams(hparams) num_space = len(equation_type.DERIVATIVE_ORDERS) num_integrated = normalized_loss_per_head.shape[-1].value - num_space - 1 # dimensions [channel] weights_list = ( [hparams.space_derivatives_weight / num_space] * num_space + [hparams.time_derivative_weight]) if num_integrated: weights_list.extend( [hparams.integrated_solution_weight / num_integrated] * num_integrated) channel_weights = tf.convert_to_tensor(weights_list) channel_weights /= tf.reduce_sum(channel_weights) # dimensions [abs/rel error, channel] weights = abs_rel_weights[:, tf.newaxis] * channel_weights[tf.newaxis, :] return tf.reduce_sum(weights * normalized_loss_per_head)
def training_loop(snapshots: np.ndarray, checkpoint_dir: str, hparams: tf.contrib.training.HParams, master: str = '') -> pd.DataFrame: """Run training. Args: snapshots: np.ndarray with shape [examples, x] with high-resolution training data. checkpoint_dir: directory to which to save model checkpoints. hparams: hyperparameters for training, as created by create_hparams(). master: string master to use for MonitoredTrainingSession. Returns: pd.DataFrame with metrics for the full training run. """ hparams = copy.deepcopy(hparams) set_data_dependent_hparams(hparams, snapshots) logging.info('Training with hyperparameters:\n%r', hparams) hparams_path = os.path.join(checkpoint_dir, 'hparams.pbtxt') with tf.gfile.GFile(hparams_path, 'w') as f: f.write(str(hparams.to_proto())) logging.info('Setting up training') _, train_step = setup_training(snapshots, hparams) train_inferer = Inferer(snapshots, hparams, training=True) test_inferer = Inferer(snapshots, hparams, training=False) global_step = tf.train.get_or_create_global_step() logging.info('Variables: %s', '\n'.join(map(str, tf.trainable_variables()))) logged_metrics = [] equation_type = equations.equation_type_from_hparams(hparams) with tf.train.MonitoredTrainingSession( master=master, checkpoint_dir=checkpoint_dir, save_checkpoint_secs=300, config=_disable_rewrite_config(), hooks=[SaveAtEnd(checkpoint_dir_to_path(checkpoint_dir))]) as sess: test_writer = tf.summary.FileWriter(os.path.join( checkpoint_dir, 'test'), sess.graph, flush_secs=60) train_writer = tf.summary.FileWriter(os.path.join( checkpoint_dir, 'train'), sess.graph, flush_secs=60) initial_step = sess.run(global_step) with test_writer, train_writer: for step in range(initial_step, hparams.learning_stops[-1]): sess.run(train_step) if (step + 1) % hparams.eval_interval == 0: train_inference_data = train_inferer.run(sess) test_inference_data = test_inferer.run(sess) train_metrics = calculate_metrics(train_inference_data, equation_type) test_metrics = calculate_metrics(test_inference_data, equation_type) logged_metrics.append((step, test_metrics, train_metrics)) logging.info(metrics_one_linear(test_metrics)) save_summaries(test_metrics, test_writer, global_step=step) save_summaries(train_metrics, train_writer, global_step=step) return metrics_to_dataframe(logged_metrics)
def test_integrate_exact_baseline_and_model(self, warmup=0, conservative=False, resample_factor=4, exact_filter_interval=None, **hparam_values): hparams = training.create_hparams(learning_rates=[1e-3], learning_stops=[20], eval_interval=10, equation_kwargs=json.dumps( {'num_points': NUM_X_POINTS}), conservative=conservative, resample_factor=resample_factor, **hparam_values) self.train(hparams) results = integrate.integrate_exact_baseline_and_model( self.checkpoint_dir, random_seed=RANDOM_SEED, times=np.linspace(0, 1, num=11), warmup=warmup, exact_filter_interval=exact_filter_interval) self.assertIsInstance(results, xarray.Dataset) self.assertEqual( dict(results.dims), { 'time': 11, 'x_high': NUM_X_POINTS, 'x_low': NUM_X_POINTS // resample_factor }) self.assertEqual(results['y_exact'].dims, ('time', 'x_high')) self.assertEqual(results['y_baseline'].dims, ('time', 'x_low')) self.assertEqual(results['y_model'].dims, ('time', 'x_low')) with self.subTest('average should be zero'): y_exact_mean = results.y_exact.mean('x_high') xarray.testing.assert_allclose(y_exact_mean, xarray.zeros_like(y_exact_mean), atol=1e-3) with self.subTest('matching initial conditions'): if conservative: resample = duckarray.resample_mean else: resample = duckarray.subsample y_exact = resample( results.y_exact.isel(time=0).values, resample_factor) np.testing.assert_allclose(y_exact, results.y_baseline.isel(time=0).values) np.testing.assert_allclose(y_exact, results.y_model.isel(time=0).values) with self.subTest('matches integrate_baseline'): equation_type = equations.equation_type_from_hparams(hparams) assert equation_type.CONSERVATIVE == conservative equation = equation_type(NUM_X_POINTS // resample_factor, resample_factor=resample_factor, random_seed=RANDOM_SEED) results2 = integrate.integrate_baseline(equation, times=np.linspace(0, 1, num=11), warmup=warmup) np.testing.assert_allclose(results['y_baseline'].data, results2['y'].data, atol=1e-5)