LordZeee/ndvi-convlstm
0
1# model_convlstm_train_focused.py2import os3import numpy as np4import torch5from torch.utils.data import Dataset, DataLoader6import torch.nn as nn7import torch.nn.functional as F # For attention8import torch.optim as optim9from sklearn.model_selection import train_test_split10from sklearn.preprocessing import StandardScaler11import glob12import yaml13import argparse14import pickle15import pandas as pd16 17# Global variable to hold the loaded configuration18CONFIG = None19 20class NpzSequenceDataset(Dataset):21 def __init__(self, file_paths, config, scalers=None, is_train=False, fit_scalers_on_subset_size=0):22 self.file_paths = file_paths23 self.config = config24 self.train_opts = config.get('training_options', {})25 self.input_len = config['time_parameters']['input_len']26 self.pred_len = config['time_parameters']['pred_len']27 self.model_input_band_order = self.train_opts.get('input_band_names', [])28 if not self.model_input_band_order:29 raise ValueError("'input_band_names' must be defined in config's training_options.")30 self.num_model_input_channels = len(self.model_input_band_order)31 self.target_band_name = self.train_opts.get('target_band_name', 'NDVI')32 self.scalers = scalers33 self.is_train = is_train34 self.nodata_value = config.get('nodata_value_in_export', -9999.0) 35 36 if self.is_train and self.scalers is None and fit_scalers_on_subset_size > 0:37 print(f"Fitting scalers on a subset of up to {fit_scalers_on_subset_size} training files...")38 self._fit_scalers_on_subset(min(fit_scalers_on_subset_size, len(self.file_paths)))39 elif self.scalers is None:40 print("Warning: No scalers provided for NpzSequenceDataset. Data will be unscaled.")41 42 def _fit_scalers_on_subset(self, num_files_to_sample):43 self.scalers = {44 'input_scalers': [StandardScaler() for _ in range(self.num_model_input_channels)],45 'target_scaler': StandardScaler()46 }47 if num_files_to_sample == 0:48 print("Scaler fitting skipped (0 files to sample). Scalers remain un-fitted.")49 return50 51 print(f"Performing scaler fitting on {num_files_to_sample} random training files.")52 sample_file_paths = np.random.choice(self.file_paths, size=num_files_to_sample, replace=False)53 54 for f_path in sample_file_paths:55 try:56 data = np.load(f_path)57 input_d_raw = data['input_data']58 target_d_raw = data['target_data']59 bands_in_file = list(data['input_bands'])60 61 ordered_input_d_for_fitting = np.full(62 (input_d_raw.shape[0], input_d_raw.shape[1], input_d_raw.shape[2], self.num_model_input_channels),63 self.nodata_value, 64 dtype=np.float3265 )66 for config_band_idx, config_band_name in enumerate(self.model_input_band_order):67 try:68 file_band_idx = bands_in_file.index(config_band_name)69 ordered_input_d_for_fitting[..., config_band_idx] = input_d_raw[..., file_band_idx]70 except ValueError:71 pass 72 73 for band_idx in range(self.num_model_input_channels):74 band_data = ordered_input_d_for_fitting[..., band_idx]75 valid_pixels = band_data[band_data != self.nodata_value].reshape(-1, 1)76 if len(valid_pixels) > 1: # partial_fit needs at least 2 samples for variance calculation77 self.scalers['input_scalers'][band_idx].partial_fit(valid_pixels)78 79 valid_target_pixels = target_d_raw[target_d_raw != self.nodata_value].reshape(-1, 1)80 if len(valid_target_pixels) > 1:81 self.scalers['target_scaler'].partial_fit(valid_target_pixels)82 except Exception as e:83 print(f"Error fitting scaler with file {f_path}: {e}")84 print("Scaler fitting on subset complete.")85 for i, scaler in enumerate(self.scalers['input_scalers']):86 if not hasattr(scaler, 'mean_'): # Check if scaler was actually fitted87 print(f"Warning: Input scaler for channel {i} ({self.model_input_band_order[i] if i < len(self.model_input_band_order) else 'Unknown'}) was not fitted (e.g., all NoData in sample).")88 if not hasattr(self.scalers['target_scaler'], 'mean_'):89 print(f"Warning: Target scaler was not fitted.")90 91 92 def __len__(self):93 return len(self.file_paths)94 95 def __getitem__(self, idx):96 file_path = self.file_paths[idx]97 try:98 npz_content = np.load(file_path)99 input_data_raw = npz_content['input_data'].astype(np.float32)100 target_data_raw = npz_content['target_data'].astype(np.float32)101 bands_in_file = list(npz_content['input_bands'])102 avg_input_clear_pixel_count = npz_content.get('avg_input_clear_pixel_count', 1.0).astype(np.float32) 103 except Exception as e:104 print(f"Error loading .npz file {file_path}: {e}. Returning dummy data.")105 dummy_input = np.zeros((self.input_len, self.config['minicube_parameters']['minicube_size_pixels'], self.config['minicube_parameters']['minicube_size_pixels'], self.num_model_input_channels), dtype=np.float32)106 dummy_target = np.zeros((self.pred_len, self.config['minicube_parameters']['minicube_size_pixels'], self.config['minicube_parameters']['minicube_size_pixels'], 1), dtype=np.float32)107 dummy_weight = np.array(1.0, dtype=np.float32)108 return torch.from_numpy(np.transpose(dummy_input, (0, 3, 1, 2))).float(), \109 torch.from_numpy(np.transpose(dummy_target, (0, 3, 1, 2))).float(), \110 torch.from_numpy(dummy_weight).float()111 112 ordered_input_data = np.full((self.input_len, input_data_raw.shape[1], input_data_raw.shape[2], self.num_model_input_channels), self.nodata_value, dtype=np.float32)113 for config_band_idx, config_band_name in enumerate(self.model_input_band_order):114 try:115 file_band_idx = bands_in_file.index(config_band_name)116 ordered_input_data[..., config_band_idx] = input_data_raw[..., file_band_idx]117 except ValueError: pass 118 119 input_data_scaled = ordered_input_data.copy()120 target_data_scaled = target_data_raw.copy()121 122 if self.scalers:123 for band_idx in range(self.num_model_input_channels): 124 if band_idx < len(self.scalers['input_scalers']) and hasattr(self.scalers['input_scalers'][band_idx], 'mean_'):125 band_data = ordered_input_data[..., band_idx]126 valid_mask_ch = (band_data != self.nodata_value)127 if np.any(valid_mask_ch):128 data_to_scale_ch = band_data[valid_mask_ch].reshape(-1,1)129 scaled_data_ch = self.scalers['input_scalers'][band_idx].transform(data_to_scale_ch)130 temp_scaled_band = np.full_like(band_data, self.nodata_value) 131 temp_scaled_band[valid_mask_ch] = scaled_data_ch.flatten()132 input_data_scaled[..., band_idx] = temp_scaled_band133 134 if hasattr(self.scalers.get('target_scaler'), 'mean_'):135 valid_mask_target = (target_data_raw != self.nodata_value)136 if np.any(valid_mask_target):137 data_to_scale_target = target_data_raw[valid_mask_target].reshape(-1,1)138 scaled_data_target = self.scalers['target_scaler'].transform(data_to_scale_target)139 temp_scaled_target = np.full_like(target_data_raw, self.nodata_value)140 temp_scaled_target[valid_mask_target] = scaled_data_target.flatten()141 target_data_scaled = temp_scaled_target.reshape(target_data_raw.shape)142 143 input_data_final = np.transpose(input_data_scaled, (0, 3, 1, 2))144 target_data_final = np.transpose(target_data_scaled, (0, 3, 1, 2))145 146 max_expected_clear_count = self.train_opts.get('max_avg_clear_pixel_for_weighting', 3.0) # e.g. S2 might have up to 3-4 clear views in a week147 sample_weight = np.clip(avg_input_clear_pixel_count / max_expected_clear_count, 0.1, 1.0) # Normalize and clip148 149 return torch.from_numpy(input_data_final).float(), \150 torch.from_numpy(target_data_final).float(), \151 torch.tensor(sample_weight, dtype=torch.float32)152 153 154# --- ConvLSTM Model Definitions with Temporal Attention ---155class ConvLSTMCell(nn.Module):156 def __init__(self, input_dim, hidden_dim, kernel_size, bias):157 super(ConvLSTMCell, self).__init__()158 self.input_dim = input_dim159 self.hidden_dim = hidden_dim160 self.kernel_size = kernel_size 161 self.padding = kernel_size[0] // 2, kernel_size[1] // 2162 self.bias = bias163 self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,164 out_channels=4 * self.hidden_dim,165 kernel_size=self.kernel_size,166 padding=self.padding,167 bias=self.bias)168 def forward(self, input_tensor, cur_state):169 h_cur, c_cur = cur_state170 combined_conv = self.conv(torch.cat([input_tensor, h_cur], dim=1))171 cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)172 i, f, o, g = torch.sigmoid(cc_i), torch.sigmoid(cc_f), torch.sigmoid(cc_o), torch.tanh(cc_g)173 c_next = f * c_cur + i * g174 h_next = o * torch.tanh(c_next)175 return h_next, c_next176 def init_hidden(self, batch_size, image_size):177 height, width = image_size178 return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),179 torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))180 181class ConvLSTMEncoder(nn.Module):182 def __init__(self, input_dim, hidden_dims_list, kernel_sizes_list, num_layers, bias=True):183 super(ConvLSTMEncoder, self).__init__()184 self.input_dim, self.hidden_dims_list, self.kernel_sizes_list, self.num_layers = input_dim, hidden_dims_list, kernel_sizes_list, num_layers185 cell_list = []186 for i in range(self.num_layers):187 cur_input_dim = self.input_dim if i == 0 else self.hidden_dims_list[i-1]188 cell_list.append(ConvLSTMCell(cur_input_dim, self.hidden_dims_list[i], self.kernel_sizes_list[i], bias))189 self.cell_list = nn.ModuleList(cell_list)190 191 def forward(self, input_tensor): 192 b, seq_len, _, h, w = input_tensor.size()193 current_layer_input_sequence = input_tensor 194 all_layer_final_states = []195 last_layer_output_sequence = None196 197 for layer_idx in range(self.num_layers):198 h_state, c_state = self.cell_list[layer_idx].init_hidden(b, (h, w))199 output_hidden_states_this_layer_for_next_input = [] 200 for t in range(seq_len):201 h_state, c_state = self.cell_list[layer_idx](current_layer_input_sequence[:, t, :, :, :], [h_state, c_state])202 output_hidden_states_this_layer_for_next_input.append(h_state)203 204 all_layer_final_states.append([h_state, c_state]) 205 current_layer_input_sequence = torch.stack(output_hidden_states_this_layer_for_next_input, dim=1)206 if layer_idx == self.num_layers - 1:207 last_layer_output_sequence = current_layer_input_sequence 208 209 return all_layer_final_states, last_layer_output_sequence210 211 212class TemporalAttention(nn.Module):213 def __init__(self, hidden_dim_encoder_last_layer):214 super(TemporalAttention, self).__init__()215 # This attention mechanism will operate on the sequence of outputs from the last encoder layer.216 # These outputs are spatial feature maps (Batch, SeqLen, Channels_last_hidden, H, W).217 # We need to reduce the spatial dimensions to get a vector per time step for attention scoring.218 self.hidden_dim = hidden_dim_encoder_last_layer219 220 # A common way is to flatten or average pool the spatial dimensions.221 # Then apply a linear layer to get an energy score for each time step.222 # For ConvLSTM outputs, a simple approach is to use a 1x1 convolution to reduce channels to 1 (energy),223 # then softmax over time. Or, adaptive average pool then linear.224 225 # Let's use adaptive average pooling to (1,1) then a linear layer for scoring.226 self.attention_scorer = nn.Linear(self.hidden_dim, 1)227 228 def forward(self, encoder_output_sequence):229 # encoder_output_sequence shape: (Batch, SeqLen, Channels_last_hidden, H, W)230 b, seq_len, c, h, w = encoder_output_sequence.shape231 232 # Reshape for pooling: (Batch * SeqLen, Channels, H, W)233 reshaped_for_pooling = encoder_output_sequence.view(b * seq_len, c, h, w)234 # Adaptive average pool to reduce H, W to 1, 1235 pooled_spatial = F.adaptive_avg_pool2d(reshaped_for_pooling, (1,1)) # (B*S, C, 1, 1)236 # Reshape back: (Batch, SeqLen, Channels)237 pooled_temporal_sequence = pooled_spatial.view(b, seq_len, c)238 239 # Calculate attention scores (energies)240 # energies shape: (Batch, SeqLen, 1)241 energies = self.attention_scorer(pooled_temporal_sequence) 242 243 # Apply softmax over the sequence length dimension to get weights244 # attn_weights shape: (Batch, SeqLen, 1)245 attn_weights = F.softmax(energies, dim=1)246 247 # Calculate context vector: weighted sum of encoder outputs (pooled_temporal_sequence)248 # attn_weights needs to be (B, S, 1), pooled_temporal_sequence is (B, S, C)249 # We want context_vector to be (B, C)250 context_vector = torch.sum(attn_weights * pooled_temporal_sequence, dim=1) # (B, C)251 252 return context_vector # This is the context vector for the decoder253 254 255class ConvLSTMDecoderWithAttention(nn.Module):256 def __init__(self, output_dim, hidden_dims_list, kernel_sizes_list, num_layers, attention_context_dim, bias=True):257 super(ConvLSTMDecoderWithAttention, self).__init__()258 self.output_dim = output_dim 259 self.hidden_dims_list = hidden_dims_list 260 self.kernel_sizes_list = kernel_sizes_list 261 self.num_layers = num_layers262 self.attention_context_dim = attention_context_dim # Dim of context vector from TemporalAttention263 264 cell_list = []265 for i in range(self.num_layers):266 # The input to the first ConvLSTM cell of the decoder at each step needs to incorporate267 # the previous prediction AND potentially the attention context.268 # A common way is to concatenate the context with the input at each step,269 # or use the context to initialize/modify the decoder's initial hidden state.270 271 # Let's try modifying the initial hidden state of the first decoder layer using the context.272 # The input_dim for the first cell (i=0) will be self.output_dim (from the previous prediction).273 cur_input_dim_cell = self.output_dim if i == 0 else self.hidden_dims_list[i-1]274 cell_list.append(ConvLSTMCell(cur_input_dim_cell, self.hidden_dims_list[i], self.kernel_sizes_list[i], bias))275 self.cell_list = nn.ModuleList(cell_list)276 277 # Linear layers to transform the attention context vector to initialize/modify278 # the first decoder layer's hidden (h) and cell (c) states.279 self.W_h_context = nn.Linear(self.attention_context_dim, self.hidden_dims_list[0])280 self.W_c_context = nn.Linear(self.attention_context_dim, self.hidden_dims_list[0])281 282 self.output_conv = nn.Conv2d(self.hidden_dims_list[-1], self.output_dim, 1, padding=0)283 284 def forward(self, encoder_final_states_list, attention_context_vector, pred_len):285 # encoder_final_states_list: list of (h,c) from each encoder layer286 # attention_context_vector: (batch, attention_context_dim)287 b, _, h_enc, w_enc = encoder_final_states_list[0][0].size() 288 289 current_decoder_hidden_states = []290 for i in range(self.num_layers):291 h_enc_layer, c_enc_layer = encoder_final_states_list[i]292 if i == 0: # Modify initial state of the first decoder layer with attention context293 # Project context vector and expand to spatial dimensions294 # W_h_context(attention_context_vector) -> (B, hidden_dims_list[0])295 # .unsqueeze(-1).unsqueeze(-1) -> (B, hidden_dims_list[0], 1, 1)296 # .expand_as(h_enc_layer) -> (B, hidden_dims_list[0], H, W)297 h_modifier = self.W_h_context(attention_context_vector).unsqueeze(-1).unsqueeze(-1).expand_as(h_enc_layer)298 c_modifier = self.W_c_context(attention_context_vector).unsqueeze(-1).unsqueeze(-1).expand_as(c_enc_layer)299 300 # Combine, e.g., by adding and applying tanh301 h_init_dec0 = torch.tanh(h_enc_layer + h_modifier)302 c_init_dec0 = torch.tanh(c_enc_layer + c_modifier)303 current_decoder_hidden_states.append([h_init_dec0, c_init_dec0])304 else: # For other layers, initialize with encoder's final state for that layer305 current_decoder_hidden_states.append([h_enc_layer, c_enc_layer])306 307 predictions = []308 # Initial input to the decoder (e.g., a zero tensor with shape of one output frame)309 decoder_input_frame = torch.zeros(b, self.output_dim, h_enc, w_enc, device=attention_context_vector.device)310 311 for t in range(pred_len):312 current_layer_input_for_cell_propagation = decoder_input_frame 313 next_time_step_hidden_states = [] # Store states for this time step from all layers314 for layer_idx in range(self.num_layers):315 h_prev_layer_t, c_prev_layer_t = current_decoder_hidden_states[layer_idx]316 317 # Input to current cell:318 # For layer 0: it's the decoder_input_frame (previous prediction or zeros)319 # For layer > 0: it's the h_state from the previous layer *at this same prediction time step*320 input_to_cell = current_layer_input_for_cell_propagation if layer_idx == 0 else next_time_step_hidden_states[-1][0]321 322 h_next_layer_t, c_next_layer_t = self.cell_list[layer_idx](input_to_cell, [h_prev_layer_t, c_prev_layer_t])323 next_time_step_hidden_states.append([h_next_layer_t, c_next_layer_t])324 current_layer_input_for_cell_propagation = h_next_layer_t # Output h becomes input for next cell in stack325 326 current_decoder_hidden_states = next_time_step_hidden_states # Update states for the next time step327 output_frame = self.output_conv(current_decoder_hidden_states[-1][0]) # Predict from last layer's h_state328 predictions.append(output_frame)329 decoder_input_frame = output_frame # Use current prediction as input for the next time step (teacher forcing can be added here)330 331 return torch.stack(predictions, dim=1)332 333 334class Seq2SeqConvLSTMAttention(nn.Module):335 def __init__(self, num_input_channels, num_output_channels, hidden_dims_list, kernel_sizes_list, num_layers, bias=True):336 super(Seq2SeqConvLSTMAttention, self).__init__()337 self.encoder = ConvLSTMEncoder(num_input_channels, hidden_dims_list, kernel_sizes_list, num_layers, bias)338 self.attention = TemporalAttention(hidden_dims_list[-1]) # Attention on last encoder layer's hidden dim339 self.decoder = ConvLSTMDecoderWithAttention(num_output_channels, hidden_dims_list, kernel_sizes_list, num_layers, 340 attention_context_dim=hidden_dims_list[-1], bias=bias)341 342 def forward(self, input_tensor, pred_len):343 encoder_final_states_list, encoder_output_sequence = self.encoder(input_tensor)344 context_vector = self.attention(encoder_output_sequence)345 predictions = self.decoder(encoder_final_states_list, context_vector, pred_len)346 return predictions347 348 349class ConvLSTMDecoder(nn.Module):350 def __init__(self, output_dim, hidden_dims_list, kernel_sizes_list, num_layers, bias=True):351 super(ConvLSTMDecoder, self).__init__()352 self.output_dim, self.hidden_dims_list, self.kernel_sizes_list, self.num_layers = output_dim, hidden_dims_list, kernel_sizes_list, num_layers353 cell_list = []354 for i in range(self.num_layers):355 cur_input_dim = self.output_dim if i == 0 else self.hidden_dims_list[i-1]356 cell_list.append(ConvLSTMCell(cur_input_dim, self.hidden_dims_list[i], self.kernel_sizes_list[i], bias))357 self.cell_list = nn.ModuleList(cell_list)358 self.output_conv = nn.Conv2d(self.hidden_dims_list[-1], self.output_dim, 1, padding=0)359 360 def forward(self, encoder_final_states, pred_len):361 b, _, h, w = encoder_final_states[0][0].size()362 current_hidden_states = encoder_final_states 363 predictions = []364 decoder_input_frame = torch.zeros(b, self.output_dim, h, w, device=encoder_final_states[0][0].device)365 for t in range(pred_len):366 current_layer_input_for_cell = decoder_input_frame 367 next_time_step_hidden_states = []368 for layer_idx in range(self.num_layers):369 h_state, c_state = current_hidden_states[layer_idx]370 input_to_cell = current_layer_input_for_cell if layer_idx == 0 else next_time_step_hidden_states[-1][0]371 h_state, c_state = self.cell_list[layer_idx](input_to_cell, [h_state, c_state])372 next_time_step_hidden_states.append([h_state, c_state])373 current_layer_input_for_cell = h_state 374 current_hidden_states = next_time_step_hidden_states 375 output_frame = self.output_conv(current_hidden_states[-1][0]) 376 predictions.append(output_frame)377 decoder_input_frame = output_frame 378 return torch.stack(predictions, dim=1)379 380 381# Standard Seq2SeqConvLSTM (without attention, for comparison or if attention is disabled)382class Seq2SeqConvLSTM(nn.Module):383 def __init__(self, num_input_channels, num_output_channels, hidden_dims_list, kernel_sizes_list, num_layers, bias=True):384 super(Seq2SeqConvLSTM, self).__init__()385 self.encoder = ConvLSTMEncoder(num_input_channels, hidden_dims_list, kernel_sizes_list, num_layers, bias)386 # Standard decoder does not use attention_context_dim387 # Its input_dim for the first cell in the feedback loop is num_output_channels388 self.decoder = ConvLSTMDecoder(num_output_channels, hidden_dims_list, kernel_sizes_list, num_layers, bias=bias) # Reusing the simpler decoder389 390 def forward(self, input_tensor, pred_len):391 encoder_final_states_list, _ = self.encoder(input_tensor) # Ignore encoder_output_sequence392 predictions = self.decoder.forward(encoder_final_states_list, pred_len) # Call original decoder's forward393 return predictions394 395 396# --- Training and Evaluation Functions ---397def calculate_rmse(predictions, targets, nodata_value=-9999.0):398 targets = targets.to(predictions.device)399 valid_mask = (targets != nodata_value) & (~torch.isnan(targets)) & (~torch.isnan(predictions))400 if not torch.any(valid_mask): return torch.tensor(float('nan'), device=predictions.device)401 squared_diff = torch.where(valid_mask, (predictions - targets)**2, torch.zeros_like(predictions))402 mean_squared_error = torch.sum(squared_diff) / torch.sum(valid_mask.float()).clamp(min=1e-6) # clamp to avoid div by zero403 return torch.sqrt(mean_squared_error)404 405def train_epoch(model, dataloader, optimizer, criterion, device, grad_clip_value=None, use_sample_weights=False):406 model.train()407 epoch_loss, epoch_rmse, num_batches = 0.0, 0.0, 0408 for batch_idx, batch_data in enumerate(dataloader):409 inputs, targets, sample_weights = batch_data[0].to(device), batch_data[1].to(device), batch_data[2].to(device)410 411 optimizer.zero_grad()412 outputs = model(inputs, targets.size(1))413 414 nodata_val = CONFIG.get('training_options', {}).get('nodata_value_for_plotting', -9999.0) 415 valid_target_mask = (targets != nodata_val) & (~torch.isnan(targets)) 416 valid_mask_for_loss = valid_target_mask & (~torch.isnan(outputs))417 418 if not torch.any(valid_mask_for_loss):419 # print(f"Warning: All targets/outputs are NoData/NaN in train batch {batch_idx}. Skipping.")420 continue421 422 per_pixel_loss = criterion(outputs, targets) 423 masked_loss_values = per_pixel_loss[valid_mask_for_loss] # Get only valid loss values424 425 if masked_loss_values.numel() == 0: # If no valid pixels after masking426 # print(f"Warning: Loss tensor is empty after masking in train batch {batch_idx}. Skipping.")427 continue428 429 # For sample weighting, we need to calculate loss per sample first, then weight, then mean.430 # This requires a bit more care if pixels within a sample can be masked.431 # Simpler approach for now: if using sample weights, apply it to the mean loss of valid pixels per sample.432 # Better: weight each valid pixel's loss, then average.433 434 # Calculate mean loss over valid pixels for each sample in the batch435 # masked_loss (B, T, C, H, W), valid_mask_for_loss (B, T, C, H, W)436 temp_masked_loss = torch.where(valid_mask_for_loss, per_pixel_loss, torch.zeros_like(per_pixel_loss))437 sum_loss_per_sample = torch.sum(temp_masked_loss, dim=[1,2,3,4])438 num_valid_pixels_per_sample = torch.sum(valid_mask_for_loss.float(), dim=[1,2,3,4]).clamp(min=1e-6)439 mean_loss_per_sample = sum_loss_per_sample / num_valid_pixels_per_sample440 441 442 if use_sample_weights:443 # sample_weights shape: (B)444 weighted_mean_loss_per_sample = mean_loss_per_sample * sample_weights445 batch_loss = weighted_mean_loss_per_sample.mean() # Mean of weighted sample losses446 else:447 batch_loss = mean_loss_per_sample.mean() # Mean of unweighted sample losses448 449 rmse = calculate_rmse(outputs, targets, nodata_val)450 451 if torch.isnan(batch_loss).item(): 452 print(f"Warning: NaN loss in train batch {batch_idx}. Skipping."); 453 continue454 batch_loss.backward()455 if grad_clip_value: torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip_value)456 optimizer.step()457 epoch_loss += batch_loss.item()458 if not torch.isnan(rmse).item(): 459 epoch_rmse += rmse.item() 460 num_batches += 1461 if batch_idx % CONFIG.get('training_options', {}).get('log_interval', 10) == 0:462 print(f" Train Batch {batch_idx+1}/{len(dataloader)}: Loss={batch_loss.item():.4f}, RMSE={rmse.item() if not torch.isnan(rmse).item() else float('nan'):.4f}")463 if num_batches == 0: return float('nan'), float('nan')464 return epoch_loss / num_batches, epoch_rmse / num_batches465 466def evaluate_epoch(model, dataloader, criterion, device): 467 model.eval()468 epoch_loss, epoch_rmse, num_batches = 0.0, 0.0, 0469 with torch.no_grad():470 for batch_idx, batch_data in enumerate(dataloader):471 inputs, targets, _ = batch_data[0].to(device), batch_data[1].to(device), batch_data[2].to(device) # Ignore sample_weights for eval loss472 outputs = model(inputs, targets.size(1))473 474 nodata_val = CONFIG.get('training_options', {}).get('nodata_value_for_plotting', -9999.0)475 valid_target_mask = (targets != nodata_val) & (~torch.isnan(targets))476 valid_mask_for_loss = valid_target_mask & (~torch.isnan(outputs))477 478 if not torch.any(valid_mask_for_loss):479 continue480 481 per_pixel_loss = criterion(outputs, targets)482 masked_loss_values = per_pixel_loss[valid_mask_for_loss]483 if masked_loss_values.numel() == 0: continue484 485 batch_loss = masked_loss_values.mean() # Mean of valid pixel losses across the batch486 rmse = calculate_rmse(outputs, targets, nodata_val)487 488 if torch.isnan(batch_loss).item(): 489 print(f"Warning: NaN loss in val batch {batch_idx}. Skipping."); 490 continue491 epoch_loss += batch_loss.item()492 if not torch.isnan(rmse).item(): 493 epoch_rmse += rmse.item() 494 num_batches += 1495 if num_batches == 0: return float('nan'), float('nan')496 return epoch_loss / num_batches, epoch_rmse / num_batches497 498# --- Main Script ---499if __name__ == '__main__':500 parser = argparse.ArgumentParser(description="Train ConvLSTM model on .npz minicube sequences.")501 parser.add_argument("config_file", type=str, help="Path to the YAML configuration file.")502 args = parser.parse_args()503 504 print(f"Loading configuration from: {args.config_file}")505 try:506 with open(args.config_file, 'r') as f:507 CONFIG = yaml.safe_load(f)508 except FileNotFoundError:509 print(f"Error: Configuration file '{args.config_file}' not found.")510 exit()511 except yaml.YAMLError as e:512 print(f"Error parsing YAML configuration file: {e}")513 exit()514 515 516 # # --- START DEBUG BLOCK 1 ---517 # print("\n--- DEBUG: Immediately after loading CONFIG ---")518 # if CONFIG is None:519 # print("CONFIG is None after loading!")520 # else:521 # print(f"Type of CONFIG: {type(CONFIG)}")522 # # print(f"Full CONFIG content: {CONFIG}") # Can be very verbose523 524 # training_options_from_config = CONFIG.get('training_options')525 # if training_options_from_config is None:526 # print("'training_options' key NOT FOUND in CONFIG.")527 # else:528 # print(f"Type of training_options_from_config: {type(training_options_from_config)}")529 # model_parameters_from_config = training_options_from_config.get('model_parameters')530 # if model_parameters_from_config is None:531 # print("'model_parameters' key NOT FOUND under 'training_options'.")532 # else:533 # print(f"Type of model_parameters_from_config: {type(model_parameters_from_config)}")534 # hidden_dims_val_from_config = model_parameters_from_config.get('hidden_dims')535 # if hidden_dims_val_from_config is None:536 # print("'hidden_dims' key NOT FOUND under 'model_parameters'. Default will be used.")537 # else:538 # print(f"SUCCESS: 'hidden_dims' found in config: {hidden_dims_val_from_config}")539 # print("--- END DEBUG BLOCK 1 ---\n")540 541 train_opts = CONFIG.get('training_options', {})542 data_dir = CONFIG['directories'].get('processed_training_data_dir')543 model_output_dir = CONFIG['directories'].get('model_output_dir', 'model_outputs') 544 os.makedirs(model_output_dir, exist_ok=True)545 546 if not data_dir: print("Error: 'processed_training_data_dir' not found in config."); exit()547 device = torch.device(train_opts.get('device', 'cuda' if torch.cuda.is_available() else 'cpu'))548 print(f"Using device: {device}")549 550 all_npz_files = []551 for ecoregion_info in CONFIG['ecoregions']:552 ecoregion_name = ecoregion_info['name']553 ecoregion_processed_dir = os.path.join(data_dir, ecoregion_name)554 if os.path.exists(ecoregion_processed_dir):555 all_npz_files.extend(glob.glob(os.path.join(ecoregion_processed_dir, "*.npz")))556 else: print(f"Warning: Processed data directory not found: {ecoregion_processed_dir}")557 558 if not all_npz_files: print(f"Error: No .npz files found under {data_dir}. Run prepare_training_data.py first."); exit()559 print(f"Found {len(all_npz_files)} total .npz sequence files.")560 561 train_files, val_files = train_test_split(all_npz_files, test_size=train_opts.get('validation_split', 0.25), random_state=train_opts.get('random_seed', 42))562 val_files, test_files = train_test_split(val_files, test_size=0.2, random_state=train_opts.get('random_seed', 42))563 print(f"Training files: {len(train_files)}, Validation files: {len(val_files)}, Test files: {len(test_files)}")564 test_files_save_path = os.path.join(model_output_dir, f"test_files_seed{train_opts.get('random_seed', 42)}.txt")565 val_files_save_path = os.path.join(model_output_dir, f"validation_files_seed{train_opts.get('random_seed', 42)}.txt")566 567 for file_list, name_suffix in [(train_files, "train"), (val_files, "validation"), (test_files, "test")]:568 if file_list: # Only save if the list is not empty569 list_save_path = os.path.join(model_output_dir, f"{name_suffix}_files_seed{train_opts.get('random_seed', 42)}.txt")570 with open(list_save_path, 'w') as f:571 for item in file_list:572 f.write(f"{item}\n")573 print(f"{name_suffix.capitalize()} file list saved to {list_save_path}")574 575 print("Initializing training dataset and fitting scalers (if configured)...")576 fit_subset_size = train_opts.get('fit_scalers_on_subset_size', 100 if len(train_files) > 100 else max(1, len(train_files)))577 if not train_files: 578 print("Error: No files available for training dataset after train/val split. Cannot fit scalers.")579 fit_subset_size = 0 580 581 train_dataset = NpzSequenceDataset(file_paths=train_files, config=CONFIG, scalers=None, is_train=True, fit_scalers_on_subset_size=fit_subset_size)582 583 fitted_scalers = None584 if hasattr(train_dataset, 'scalers') and train_dataset.scalers is not None:585 if train_dataset.scalers.get('input_scalers') and \586 len(train_dataset.scalers['input_scalers']) > 0 and \587 hasattr(train_dataset.scalers['input_scalers'][0], 'mean_') and \588 hasattr(train_dataset.scalers.get('target_scaler'), 'mean_'):589 fitted_scalers = train_dataset.scalers590 scaler_save_path = os.path.join(model_output_dir, 'fitted_scalers.pkl')591 try:592 with open(scaler_save_path, 'wb') as f: pickle.dump(fitted_scalers, f)593 print(f"Fitted scalers saved to {scaler_save_path}")594 except Exception as e: print(f"Error saving scalers: {e}")595 else:596 print("Warning: Scalers were initialized but not fitted (likely due to no training files or small subset). Data will be unscaled.")597 else: print("Warning: Training dataset did not produce scalers. Data will be unscaled.")598 599 val_dataset = NpzSequenceDataset(file_paths=val_files, config=CONFIG, scalers=fitted_scalers, is_train=False)600 train_loader = DataLoader(train_dataset, batch_size=train_opts.get('batch_size', 4), shuffle=True, num_workers=train_opts.get('num_workers', 0), pin_memory=True if device.type == 'cuda' else False)601 val_loader = DataLoader(val_dataset, batch_size=train_opts.get('batch_size', 4), shuffle=False, num_workers=train_opts.get('num_workers', 0), pin_memory=True if device.type == 'cuda' else False)602 603 input_band_names_from_config = train_opts.get('input_band_names', [])604 if not input_band_names_from_config: print("Error: 'input_band_names' must be defined in config."); exit()605 num_input_channels = len(input_band_names_from_config)606 num_output_channels = 1 # Predicting NDVI607 608 model_params = train_opts.get('model_parameters', {})609 hidden_dims = model_params.get('hidden_dims') 610 kernel_sizes_config = model_params.get('kernel_sizes', [[3,3], [3,3], [3,3]]) 611 kernel_sizes_tuples = [tuple(ks) for ks in kernel_sizes_config]612 num_layers = model_params.get('num_layers', 3)613 use_attention = model_params.get('use_attention', True) 614 # # --- START DEBUG BLOCK 2 ---615 # print("\n--- DEBUG: After fetching model_params ---")616 # print(f"train_opts dictionary: {train_opts}")617 # print(f"model_params dictionary: {model_params}")618 # # --- END DEBUG BLOCK 2 ---\n")619 620 print(f"\n hidden_dims value being used for model: {hidden_dims}\n") #621 if len(hidden_dims) != num_layers or len(kernel_sizes_tuples) != num_layers: print("Error: Length of hidden_dims/kernel_sizes must match num_layers."); exit()622 623 if use_attention:624 print("Using Seq2SeqConvLSTM with Temporal Attention.")625 model = Seq2SeqConvLSTMAttention(num_input_channels, num_output_channels, hidden_dims, kernel_sizes_tuples, num_layers).to(device)626 else:627 print("Using standard Seq2SeqConvLSTM (no attention).")628 model = Seq2SeqConvLSTM(num_input_channels, num_output_channels, hidden_dims, kernel_sizes_tuples, num_layers).to(device)629 630 print(f"Model initialized: {num_input_channels} input channels, {num_output_channels} output channels, {num_layers} layers. Attention: {use_attention}")631 632 criterion = nn.MSELoss(reduction='none') 633 optimizer = optim.Adam(model.parameters(), lr=train_opts.get('learning_rate', 0.001))634 scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=train_opts.get('lr_patience', 5))635 636 num_epochs = train_opts.get('num_epochs', 20)637 best_val_rmse = float('inf')638 model_save_full_path = os.path.join(model_output_dir, train_opts.get('model_save_path', 'convlstm_model_best.pth'))639 grad_clip = train_opts.get('gradient_clipping', None)640 use_sample_weighting_from_config = train_opts.get('use_sample_weighting', True) 641 training_log = []642 643 print(f"Sample weighting during training: {use_sample_weighting_from_config}")644 645 for epoch in range(1, num_epochs + 1):646 print(f"\nEpoch {epoch}/{num_epochs}")647 current_lr = optimizer.param_groups[0]['lr']648 print(f"Current learning rate: {current_lr}")649 train_loss, train_rmse = train_epoch(model, train_loader, optimizer, criterion, device, 650 grad_clip_value=grad_clip, 651 use_sample_weights=use_sample_weighting_from_config)652 val_loss, val_rmse = evaluate_epoch(model, val_loader, criterion, device) 653 654 print(f"Epoch Summary: Train Loss={train_loss:.4f}, Train RMSE={train_rmse:.4f} | Val Loss={val_loss:.4f}, Val RMSE={val_rmse:.4f}")655 training_log.append({'epoch': epoch, 'train_loss': train_loss, 'train_rmse': train_rmse, 'val_loss': val_loss, 'val_rmse': val_rmse})656 657 scheduler.step(val_rmse if not np.isnan(val_rmse) else float('inf')) 658 659 if not np.isnan(val_rmse) and val_rmse < best_val_rmse:660 best_val_rmse = val_rmse661 torch.save(model.state_dict(), model_save_full_path)662 print(f"Model improved and saved to {model_save_full_path} (Val RMSE: {best_val_rmse:.4f})")663 664 log_df = pd.DataFrame(training_log)665 log_save_path = os.path.join(model_output_dir, 'training_log.csv')666 log_df.to_csv(log_save_path, index=False)667 print(f"Training log saved to {log_save_path}")668 669 print("\nTraining finished.")670 print(f"Best validation RMSE: {best_val_rmse:.4f}") 671 672 print("\nTo visualize predictions, run 'plot_model_predictions.py' with the appropriate config and saved model path.")673