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)
