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