def test_diffusive_convolutional_model_gives_bounded_output(): """ The model with diffusive enabled should give outputs in the range of the inputs. """ fv3fit.set_random_seed(1) # don't set n_feature too high for this, because of curse of dimensionality n_sample, n_tile, nx, ny, n_feature = 10, 6, 12, 12, 2 low, high = 0.0, 1.0 train_sample_func = get_uniform_sample_func(size=(n_sample, n_tile, nx, ny, n_feature), low=low - 1.0, high=high + 1.0) test_sample_func = get_uniform_sample_func(size=(n_sample, n_tile, nx, ny, n_feature), low=low, high=high) input_variables, output_variables, test_dataset = get_dataset( model_type="convolutional", sample_func=test_sample_func) hyperparameters = ConvolutionalHyperparameters( input_variables=input_variables, output_variables=output_variables, convolutional_network=ConvolutionalNetworkConfig(diffusive=True), ) result = train_identity_model("convolutional", sample_func=train_sample_func, hyperparameters=hyperparameters) output = result.model.predict(test_dataset) for name, data_array in output.data_vars.items(): assert np.all(data_array.values >= low), name assert np.all(data_array.values <= high), name
def test_convolutional_network_build_initial_loss_near_one(): fv3fit.set_random_seed(0) nt, nx, ny, nz = 5, 12, 12, 15 ds = xr.Dataset( data_vars={ "var_in": xr.DataArray( np.random.randn(nt, 6, nx, ny, nz), dims=["sample", "tile", "x", "y", "z"], ), "var_out": xr.DataArray( np.random.randn(nt, 6, nx, ny, nz), dims=["sample", "tile", "x", "y", "z"], ), }) config = ConvolutionalHyperparameters(input_variables=["var_in"], output_variables=["var_out"]) X, y = XyMultiArraySequence( X_names=["var_in"], y_names=["var_out"], dataset_sequence=[ds], unstacked_dims=["x", "y", "z"], n_halo=config.convolutional_network.halos_required, )[0] _, predict_model = build_model(config=config, X=X, y=y) out = predict_model.predict(X) np.testing.assert_allclose(np.std(y - out), 1.0, atol=0.3) loss = ConvolutionalHyperparameters.loss.loss( std=np.std(ds["var_out"], axis=(0, 1, 2, 3))) np.testing.assert_allclose(loss(y, out), 1.0, atol=0.3)
def test_default_convolutional_model_is_transpose_invariant(): """ The model with default configuration options can learn the identity function, using gaussian-sampled data around 0 with unit variance. """ fv3fit.set_random_seed(1) # don't set n_feature too high for this, because of curse of dimensionality n_sample, n_tile, nx, ny, n_feature = 10, 6, 12, 12, 2 sample_func = get_uniform_sample_func(size=(n_sample, n_tile, nx, ny, n_feature)) result = train_identity_model("convolutional", sample_func=sample_func) transpose_input = result.test_dataset.copy(deep=True) transpose_input["var_in_3d"].values[:] = np.transpose( transpose_input["var_in_3d"].values, axes=(0, 1, 3, 2, 4)) transpose_output = result.model.predict(result.test_dataset) transpose_output["var_out"].values[:] = np.transpose( transpose_output["var_out"].values, axes=(0, 1, 3, 2, 4)) output_from_transpose = result.model.predict(transpose_input) n_halo = result.hyperparameters.convolutional_network.halos_required # transposing tile data messes up neighbors, so we have to assess only on # data that has no halo dependence assert n_halo * 2 < nx assert n_halo * 2 < ny xr.testing.assert_allclose( output_from_transpose.isel(x=slice(n_halo, -n_halo), y=slice(n_halo, -n_halo)), transpose_output.isel(x=slice(n_halo, -n_halo), y=slice(n_halo, -n_halo)), atol=1e-5, )
def test_default_output(regtest): """default output should not change""" fv3fit.set_random_seed(0) config = fv3fit.DenseNetworkConfig() array = np.random.randn(5, 10) out = config.build(array, n_features_out=3) print_result(out, decimals=5, file=regtest)
def test_convolutional_network_build_standard_input_gives_standard_output(): fv3fit.set_random_seed(0) nt, nx, ny, nz = 5, 12, 12, 15 ds = xr.Dataset( data_vars={ "var_in": xr.DataArray( np.random.randn(nt, 6, nx, ny, nz), dims=["sample", "tile", "x", "y", "z"], ), "var_out": xr.DataArray( np.random.randn(nt, 6, nx, ny, nz), dims=["sample", "tile", "x", "y", "z"], ), }) config = ConvolutionalHyperparameters(input_variables=["var_in"], output_variables=["var_out"]) X, y = XyMultiArraySequence( X_names=["var_in"], y_names=["var_out"], dataset_sequence=[ds], unstacked_dims=["x", "y", "z"], n_halo=config.convolutional_network.halos_required, )[0] _, predict_model = build_model(config=config, X=X, y=y) out = predict_model.predict(X) np.testing.assert_almost_equal(np.mean(out), 0.0, decimal=1) out_std = np.std(out) # std isn't going to be 1 because of relu activation function, should be less assert out_std < 1.0 assert out_std > 0.1
def test_convolutional_network_concat_2d_and_3d_inputs(): fv3fit.set_random_seed(0) nt, nx, ny, nz = 5, 12, 12, 15 ds = xr.Dataset( data_vars={ "var_in_2d": xr.DataArray( np.random.randn(nt, 6, nx, ny), dims=["sample", "tile", "x", "y"], ), "var_in_3d": xr.DataArray( np.random.randn(nt, 6, nx, ny, nz), dims=["sample", "tile", "x", "y", "z"], ), "var_out": xr.DataArray( np.random.randn(nt, 6, nx, ny, nz), dims=["sample", "tile", "x", "y", "z"], ), }) config = ConvolutionalHyperparameters( input_variables=["var_in_2d", "var_in_3d"], output_variables=["var_out"]) X, y = XyMultiArraySequence( X_names=["var_in_2d", "var_in_3d"], y_names=["var_out"], dataset_sequence=[ds], unstacked_dims=["x", "y", "z"], n_halo=config.convolutional_network.halos_required, )[0] _, predict_model = build_model(config=config, X=X, y=y) predict_model.predict(X)
def test_network_has_correct_number_of_hidden_layers(depth): fv3fit.set_random_seed(0) config = fv3fit.DenseNetworkConfig(depth=depth) n_features_in, n_features_out = 5, 5 input = tf.keras.layers.Input(shape=(n_features_in, )) dense_network = config.build(input, n_features_out=n_features_out) # one layer of depth is the output assert len(dense_network.hidden_outputs) == depth - 1
def test_network_has_gaussian_noise_layer(): fv3fit.set_random_seed(0) config = fv3fit.DenseNetworkConfig(gaussian_noise=0.1) n_features_in, n_features_out = 5, 5 input = tf.keras.layers.Input(shape=(n_features_in, )) dense_network = config.build(input, n_features_out=n_features_out) model = tf.keras.Model(inputs=input, outputs=dense_network.output) assert any( isinstance(layer, tf.keras.layers.GaussianNoise) for layer in model.layers)
def main(args, unknown_args=None): with open(args.training_config, "r") as f: config_dict = yaml.safe_load(f) if unknown_args is not None: config_dict = get_arg_updated_config_dict(args=unknown_args, config_dict=config_dict) training_config = fv3fit.TrainingConfig.from_dict(config_dict) with open(args.training_data_config, "r") as f: training_data_config = loaders.BatchesLoader.from_dict( yaml.safe_load(f)) fv3fit.set_random_seed(training_config.random_seed) dump_dataclass(training_config, os.path.join(args.output_path, "train.yaml")) dump_dataclass(training_data_config, os.path.join(args.output_path, "training_data.yaml")) train_batches: loaders.typing.Batches = training_data_config.load_batches( variables=training_config.variables) if args.validation_data_config is not None: with open(args.validation_data_config, "r") as f: validation_data_config = loaders.BatchesLoader.from_dict( yaml.safe_load(f)) dump_dataclass( validation_data_config, os.path.join(args.output_path, "validation_data.yaml"), ) val_batches = validation_data_config.load_batches( variables=training_config.variables) else: val_batches: Sequence[xr.Dataset] = [] if args.local_download_path: train_batches = loaders.to_local( train_batches, os.path.join(args.local_download_path, "train")) val_batches = loaders.to_local( val_batches, os.path.join(args.local_download_path, "validation")) train = fv3fit.get_training_function(training_config.model_type) model = train( hyperparameters=training_config.hyperparameters, train_batches=train_batches, validation_batches=val_batches, ) if len(training_config.derived_output_variables) > 0: model = fv3fit.DerivedModel(model, training_config.derived_output_variables) fv3fit.dump(model, args.output_path)
def test_standard_input_gives_standard_output(): fv3fit.set_random_seed(0) config = fv3fit.ConvolutionalNetworkConfig() array = np.random.randn(2, 10, 10, 20) np.testing.assert_almost_equal(np.mean(array), 0.0, decimal=1) np.testing.assert_almost_equal(np.std(array), 1.0, decimal=1) convolutional_network = config.build(array, n_features_out=3) np.testing.assert_almost_equal(np.mean(convolutional_network.output), 0.0, decimal=1) out_std = np.std(convolutional_network.output) # std isn't going to be 1 because of relu activation function, should be less assert out_std < 1.0 assert out_std > 0.1
def main(config: TrainConfig, seed: int = 0, model_url: str = None): logging.basicConfig(level=getattr(logging, config.log_level)) set_random_seed(seed) if config.use_wandb: d = asdict(config) d["model_url_override"] = model_url config.wandb.init(config=d) if model_url is None: model = load_final_model_or_checkpoint(config.out_url) else: logger.info(f"Loading user specified model from {model_url}") model = tf.keras.models.load_model(model_url) train_ds = nc_dir_to_tf_dataset(config.train_url, config.transform, nfiles=config.nfiles) test_ds = nc_dir_to_tf_dataset(config.test_url, config.transform, nfiles=config.nfiles_valid) train_set = next(iter(train_ds.shuffle(100_000).batch(50_000))) test_set = next(iter(test_ds.shuffle(160_000).batch(80_000))) train_scores, train_profiles = score_model(model, train_set) test_scores, test_profiles = score_model(model, test_set) logger.debug("Scoring Complete") if config.use_wandb: pred_sample = model.predict(test_set) log_profile_plots(test_set, pred_sample) # add level for dataframe index, assumes equivalent feature dims sample_profile = next(iter(train_profiles.values())) train_profiles["level"] = np.arange(len(sample_profile)) test_profiles["level"] = np.arange(len(sample_profile)) log_to_table("score/train", train_scores, index=[config.wandb.job.name]) log_to_table("score/test", test_scores, index=[config.wandb.job.name]) log_to_table("profiles/train", train_profiles) log_to_table("profiles/test", test_profiles) with put_dir(config.out_url) as tmpdir: with open(os.path.join(tmpdir, "scores.json"), "w") as f: json.dump({"train": train_scores, "test": test_scores}, f)
def test_train_with_same_seed_gives_same_result(model_type): n_sample, n_tile, nx, ny, n_feature = 1, 6, 12, 12, 2 fv3fit.set_random_seed(0) sample_func = get_uniform_sample_func(size=(n_sample, n_tile, nx, ny, n_feature)) first_result = train_identity_model(model_type, sample_func) fv3fit.set_random_seed(0) sample_func = get_uniform_sample_func(size=(n_sample, n_tile, nx, ny, n_feature)) second_result = train_identity_model(model_type, sample_func) xr.testing.assert_equal(first_result.test_dataset, second_result.test_dataset) first_output = first_result.model.predict(first_result.test_dataset) second_output = second_result.model.predict(first_result.test_dataset) xr.testing.assert_equal(first_output, second_output)
def test_train_default_model_on_identity(model_type, regtest): """ The model with default configuration options can learn the identity function, using gaussian-sampled data around 0 with unit variance. """ fv3fit.set_random_seed(1) # don't set n_feature too high for this, because of curse of dimensionality n_sample, n_tile, nx, ny, n_feature = 5, 6, 12, 12, 2 sample_func = get_uniform_sample_func(size=(n_sample, n_tile, nx, ny, n_feature)) assert_can_learn_identity( model_type, sample_func=sample_func, max_rmse=0.2, regtest=regtest, )
def test_fit_loop_calls_in_reproducible_order(): n_batches = 5 n_epochs = 2 config = fv3fit.TrainingLoopConfig(epochs=n_epochs) first_mock_model = mock.MagicMock(spec=tf.keras.Model) second_mock_model = mock.MagicMock(spec=tf.keras.Model) mock_Xy = [] for _ in range(n_batches): mock_Xy.append( (mock.MagicMock(spec=np.ndarray), mock.MagicMock(spec=np.ndarray)) ) validation_data = (mock.MagicMock(spec=np.ndarray), mock.MagicMock(spec=np.ndarray)) fv3fit.set_random_seed(0) config.fit_loop(first_mock_model, mock_Xy, validation_data) fv3fit.set_random_seed(0) config.fit_loop(second_mock_model, mock_Xy, validation_data) assert first_mock_model.fit.call_args_list == second_mock_model.fit.call_args_list
def main(config: TrainConfig, seed: int = 0): logging.basicConfig(level=getattr(logging, config.log_level)) set_random_seed(seed) callbacks = [] if config.use_wandb: config.wandb.init(config=asdict(config)) callbacks.append(config.wandb.get_callback()) train_ds = nc_dir_to_tf_dataset(config.train_url, config.get_dataset_convertor(), nfiles=config.nfiles) test_ds = nc_dir_to_tf_dataset(config.test_url, config.get_dataset_convertor(), nfiles=config.nfiles_valid) train_set = next(iter(train_ds.shuffle(100_000).batch(50_000))) model = config.build_model(train_set) if config.shuffle_buffer_size is not None: train_ds = train_ds.shuffle(config.shuffle_buffer_size) if config.checkpoint_model: callbacks.append( ModelCheckpointCallback(filepath=os.path.join( config.out_url, "checkpoints", "epoch.{epoch:03d}.tf"))) with tempfile.TemporaryDirectory() as train_temp: with tempfile.TemporaryDirectory() as test_temp: train_ds_batched = train_ds.batch(config.batch_size) test_ds_batched = test_ds.batch(config.batch_size) if config.cache: train_ds_batched = train_ds_batched.cache(train_temp) test_ds_batched = test_ds_batched.cache(test_temp) history = train( model, train_ds_batched, config.build_loss(train_set), optimizer=config.loss.optimizer.instance, epochs=config.epochs, validation_data=test_ds_batched, validation_freq=config.valid_freq, verbose=config.verbose, callbacks=callbacks, ) logger.debug("Training complete") with put_dir(config.out_url) as tmpdir: with open(os.path.join(tmpdir, "history.json"), "w") as f: json.dump(history.params, f) with open(os.path.join(tmpdir, "config.yaml"), "w") as f: f.write(yaml.safe_dump(asdict(config))) local_model_path = save_model(model, tmpdir) if config.use_wandb: store_model_artifact(local_model_path, name=config._model.name) # Jacobians after model storing in case of "out of memory" errors std_jacobians = compute_standardized_jacobians( model, config.get_transform().forward(train_set), config.input_variables) save_jacobians(std_jacobians, config.out_url, "jacobians.npz") if config.use_wandb: plot_all_output_sensitivities(std_jacobians)