Training Code for BrO Vertical Profile Retrieval Using a CNN-LSTM Model with Attention Mechanism

This code implements a CNN-LSTM + Multi-Head Self-Attention model to retrieve BrO vertical profiles from MAX-DOAS differential slant column densities (DSCDs).Main features:- Data loading: uses a memory-safe MemorySafeDataGenerator to read large-scale training/validation data in Parquet format batch by batch.- Model definition: includes a two-stage residual 1D convolutional encoder, multi-head self-attention bottleneck, residual decoder with skip connections, and a BiLSTM temporal module.- Training pipeline: supports mixed precision training (mixed_float16), multi-GPU distributed strategy (MirroredStrategy), custom R² metric, learning rate warmup + cosine decay scheduling.- Callbacks and saving: automatically saves the best model checkpoint, records training logs (JSON format), and generates training curve plots.- Crash recovery: automatically saves a crash recovery model if training fails.Input features (length 32):- SZA (Solar Zenith Angle)- RAA (Relative Azimuth Angle)- 9 DSCD values (DSCD_1 … DSCD_9)- 21 aerosol extinction coefficients (Aerosol_1 … Aerosol_21)Output target (length 191):- BrO vertical profile (191 altitude levels), column names BrO_1 … BrO_191Requirements:- Python 3.8+- TensorFlow 2.12+- Pandas, NumPy, PyArrow, MatplotlibUsage steps:1. Install dependencies: pip install -r requirements.txt2. Prepare data: organize Parquet files into directories: train/input, train/output, verify/input, verify/output.3. Modify the current_dir path and training parameters (e.g., batch_size, epochs) in train.py.4. Run training: python train.py