if acc_loss:
# Generate the data to fit the scaler. The generator will by default apply scaling because it is necessary
# to automate its use in the Keras fit_generator method, so disable it when dealing with data to fit the scaler
print('Finding climatology for ACC loss...')
p_fit, t_fit = generator.generate([], scale_and_impute=False)
climo = t_fit.mean(axis=0, keepdims=True)
p_fit, t_fit = (None, None)
loss_function = anomaly_correlation_loss(climo,
regularize_mean='mse',
reverse=True)
else:
loss_function = mean_squared_error
if weight_loss:
loss_function = latitude_weighted_loss(loss_function,
generator.ds.lat.values,
generator.convolution_shape,
axis=-2,
weighting='midlatitude')
# Build the model
try:
dlwp.build_model(layers,
loss=loss_function,
optimizer='adam',
metrics=['mae'],
gpus=n_gpu)
except (ValueError, IndexError):
for layer in dlwp.base_model.layers:
print(layer.name, layer.output_shape)
raise
print(dlwp.base_model.summary())
forecast_steps=num_forecast_steps, dt=dt)
verification = verification.sel(lat=((verification.lat >= lat_min) & (verification.lat <= lat_max)))
#%% Iterate through the models and calculate their stats
for m, model in enumerate(models):
print('Loading model %s...' % model)
# Some tolerance for using a weighted loss function. Unreliable but doesn't hurt.
if 'weight' in model.lower():
lats = validation_data.lat.values
output_shape = (validation_data.dims['lat'], validation_data.dims['lon'])
if crop_north_pole:
lats = lats[1:]
customs = {'loss': latitude_weighted_loss(mean_squared_error, lats, output_shape, axis=-2,
weighting='midlatitude')}
else:
customs = None
# Load the model
dlwp, history = load_model('%s/%s' % (root_directory, model), True, custom_objects=customs)
# Assign forecast hour coordinate
if isinstance(dlwp, DLWPFunctional):
if not hasattr(dlwp, '_n_steps'):
dlwp._n_steps = 6
if not hasattr(dlwp, 'time_dim'):
dlwp.time_dim = 2
sequence = dlwp._n_steps
else:
sequence = None
model_function = skip_model if skip_connections else basic_model
outputs = [model_function(input_0)]
for o in range(1, integration_steps):
outputs.append(model_function(outputs[o - 1]))
if loss_by_step is None:
loss_by_step = [1. / integration_steps] * integration_steps
model = Model(inputs=input_0, outputs=outputs)
# Example custom loss function: pass to loss= in build_model()
if weight_loss:
loss_function = latitude_weighted_loss(mean_squared_error,
generator.ds.lat.values,
generator.output_convolution_shape,
axis=-2,
weighting='midlatitude')
else:
loss_function = 'mse'
# Build the DLWP model
dlwp.build_model(model,
loss=loss_function,
loss_weights=loss_by_step,
optimizer='adam',
metrics=['mae'],
gpus=n_gpu)
print(dlwp.base_model.summary())
#%% Train, evaluate, and save the model