hoàn thành model swing-unet
This commit is contained in:
Victor Phan
+150
-12
@@ -277,7 +277,7 @@ def train_model(
|
||||
use_gpu=True,
|
||||
use_cache=True,
|
||||
test_size=0.2,
|
||||
feature_mode='simple', # Changed from 'odc' - simple mode works with B04, B08, SCL only
|
||||
feature_mode='odc', # ODC mode: 8 features (NDVI stats + NDWI/NDBI/EVI) for better accuracy
|
||||
output_model_path=None,
|
||||
status_callback=None,
|
||||
cancel_check=None
|
||||
@@ -406,9 +406,11 @@ def train_model(
|
||||
# Load different bands based on feature mode
|
||||
if feature_mode == 'simple':
|
||||
bands_to_load = ["B04", "B08", "SCL"]
|
||||
else: # temporal or extended
|
||||
else: # odc, temporal, or extended - all need full spectral bands
|
||||
bands_to_load = ["B02", "B03", "B04", "B08", "B11", "SCL"]
|
||||
|
||||
update_status(f"Loading bands: {bands_to_load} for mode={feature_mode}", 26)
|
||||
|
||||
ds_s2 = stac_load(
|
||||
items_s2,
|
||||
bands=bands_to_load,
|
||||
@@ -428,9 +430,11 @@ def train_model(
|
||||
if 'time' in ds_s2.dims:
|
||||
print(f"[DEBUG S2] Time range: {ds_s2.time.min().values} to {ds_s2.time.max().values}")
|
||||
|
||||
# Rename for compatibility (simple mode)
|
||||
if "B04" in ds_s2 and "red" not in ds_s2:
|
||||
# Rename bands ONLY for simple mode (simple mode uses 'red', 'nir', 'scl' names)
|
||||
# Other modes (odc, extended, temporal) use original band names (B02, B03, B04, B08, B11, SCL)
|
||||
if feature_mode == 'simple' and "B04" in ds_s2 and "red" not in ds_s2:
|
||||
ds_s2 = ds_s2.rename({"B04": "red", "B08": "nir", "SCL": "scl"})
|
||||
print(f"[DEBUG S2] Renamed bands for simple mode: B04→red, B08→nir, SCL→scl")
|
||||
|
||||
check_cancellation()
|
||||
|
||||
@@ -575,6 +579,7 @@ def train_model(
|
||||
update_status("Extracting features from satellite data...", 60)
|
||||
|
||||
print(f"[DEBUG] Starting feature extraction...")
|
||||
print(f"[DEBUG] Feature mode: {feature_mode}")
|
||||
print(f"[DEBUG] Training GDF has {len(train_gdf)} points")
|
||||
print(f"[DEBUG] Training GDF CRS: {train_gdf.crs}")
|
||||
print(f"[DEBUG] Training GDF bounds (UTM): {train_gdf.total_bounds}")
|
||||
@@ -637,7 +642,85 @@ def train_model(
|
||||
features = np.array(features)
|
||||
labels = np.array(labels)
|
||||
|
||||
else: # temporal or extended mode
|
||||
elif feature_mode in ['odc', 'extended']:
|
||||
# For odc/extended: Extract features for full raster first, then sample at points
|
||||
update_status(f"Extracting {feature_mode} features from full raster...", 62)
|
||||
|
||||
# Apply cloud mask first
|
||||
if 'SCL' in ds_s2:
|
||||
scl_band = ds_s2['SCL']
|
||||
cloud_mask = scl_band.isin([1, 3, 8, 9, 10])
|
||||
for band in ds_s2.data_vars:
|
||||
if band != 'SCL':
|
||||
ds_s2[band] = ds_s2[band].where(~cloud_mask)
|
||||
|
||||
# Extract features using FeatureExtractor for entire raster
|
||||
raster_features = extractor.extract(
|
||||
s2_data=ds_s2,
|
||||
vh_data=None, # ODC/extended don't use radar in aggregate
|
||||
vv_data=None
|
||||
)
|
||||
|
||||
print(f"[DEBUG] Extracted raster features: shape={raster_features.shape}")
|
||||
print(f"[DEBUG] Feature range: [{raster_features.min()}, {raster_features.max()}]")
|
||||
|
||||
# Now sample at each training point
|
||||
features = []
|
||||
labels = []
|
||||
failed_extractions = 0
|
||||
|
||||
# Get spatial dimensions
|
||||
y_coords = ds_s2.y.values
|
||||
x_coords = ds_s2.x.values
|
||||
|
||||
print(f"[DEBUG] S2 spatial grid: x=[{x_coords.min()}, {x_coords.max()}], y=[{y_coords.min()}, {y_coords.max()}]")
|
||||
|
||||
for idx, row in train_gdf.iterrows():
|
||||
point = row.geometry
|
||||
x_coord = point.x
|
||||
y_coord = point.y
|
||||
label = row[label_column]
|
||||
|
||||
try:
|
||||
# Find nearest pixel indices
|
||||
x_idx = np.argmin(np.abs(x_coords - x_coord))
|
||||
y_idx = np.argmin(np.abs(y_coords - y_coord))
|
||||
|
||||
# Get features at this pixel
|
||||
# raster_features shape: (n_pixels, n_features)
|
||||
# Need to convert 2D (y, x) index to 1D pixel index
|
||||
pixel_idx = y_idx * len(x_coords) + x_idx
|
||||
|
||||
if pixel_idx < len(raster_features):
|
||||
feature_vec = raster_features[pixel_idx]
|
||||
|
||||
if idx < 3:
|
||||
print(f"[DEBUG] Point {idx}: coords=({x_coord:.2f}, {y_coord:.2f}) -> pixel[{y_idx},{x_idx}] -> idx={pixel_idx}, features={feature_vec[:3]}...")
|
||||
|
||||
if not np.isnan(feature_vec).any():
|
||||
features.append(feature_vec)
|
||||
labels.append(label)
|
||||
else:
|
||||
failed_extractions += 1
|
||||
if idx < 3:
|
||||
print(f"[DEBUG] Point {idx} has NaN features")
|
||||
else:
|
||||
failed_extractions += 1
|
||||
if idx < 3:
|
||||
print(f"[DEBUG] Point {idx} pixel_idx {pixel_idx} out of range (max={len(raster_features)})")
|
||||
except Exception as e:
|
||||
failed_extractions += 1
|
||||
if idx < 3:
|
||||
print(f"[DEBUG] Point {idx} extraction failed: {e}")
|
||||
continue
|
||||
|
||||
if failed_extractions > 0:
|
||||
update_status(f"⚠️ {failed_extractions}/{len(train_gdf)} points had NaN/missing data", 65)
|
||||
|
||||
features = np.array(features)
|
||||
labels = np.array(labels)
|
||||
|
||||
else: # temporal mode
|
||||
# Apply cloud mask for temporal/extended modes
|
||||
if 'scl' in ds_s2 or 'SCL' in ds_s2:
|
||||
scl_band = ds_s2['scl'] if 'scl' in ds_s2 else ds_s2['SCL']
|
||||
@@ -882,19 +965,39 @@ def train_model(
|
||||
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
|
||||
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
|
||||
|
||||
# Loss and optimizer with weight decay
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
# Calculate class weights for imbalanced data
|
||||
class_counts = np.bincount(y_train)
|
||||
class_weights = 1.0 / (class_counts + 1e-6) # Avoid division by zero
|
||||
class_weights = class_weights / class_weights.sum() * len(class_counts) # Normalize
|
||||
class_weights_tensor = torch.FloatTensor(class_weights).to(device)
|
||||
|
||||
print(f"[SWIN-UNET] Class distribution: {class_counts}")
|
||||
print(f"[SWIN-UNET] Class weights: {class_weights}")
|
||||
|
||||
# Loss with class weights and optimizer with weight decay
|
||||
criterion = nn.CrossEntropyLoss(weight=class_weights_tensor)
|
||||
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
|
||||
|
||||
# LR scheduler for better convergence
|
||||
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
|
||||
|
||||
# Early stopping to prevent overfitting
|
||||
best_val_loss = float('inf')
|
||||
patience = 10
|
||||
patience_counter = 0
|
||||
|
||||
# Train Swin-UNet
|
||||
update_status("Training Swin-UNet model with PyTorch...", 80)
|
||||
update_status("Training Swin-UNet model with PyTorch (with class weights)...", 80)
|
||||
epochs = min(60, n_estimators // 2) # Swin-UNet benefits from more epochs
|
||||
|
||||
# Validation dataset
|
||||
val_dataset = TensorDataset(X_test_tensor, y_test_tensor)
|
||||
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
|
||||
|
||||
model.train()
|
||||
for epoch in range(epochs):
|
||||
# Training phase
|
||||
model.train()
|
||||
epoch_loss = 0.0
|
||||
for batch_X, batch_y in train_loader:
|
||||
batch_X, batch_y = batch_X.to(device), batch_y.to(device)
|
||||
@@ -903,16 +1006,51 @@ def train_model(
|
||||
outputs = model(batch_X)
|
||||
loss = criterion(outputs, batch_y)
|
||||
loss.backward()
|
||||
|
||||
# Gradient clipping to prevent exploding gradients
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||||
|
||||
optimizer.step()
|
||||
|
||||
epoch_loss += loss.item()
|
||||
|
||||
scheduler.step()
|
||||
|
||||
if (epoch + 1) % 10 == 0:
|
||||
avg_loss = epoch_loss / len(train_loader)
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
update_status(f"Swin-UNet Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}, LR: {lr:.6f}", 80 + (epoch / epochs) * 10)
|
||||
# Validation phase
|
||||
model.eval()
|
||||
val_loss = 0.0
|
||||
correct = 0
|
||||
total = 0
|
||||
with torch.no_grad():
|
||||
for batch_X, batch_y in val_loader:
|
||||
batch_X, batch_y = batch_X.to(device), batch_y.to(device)
|
||||
outputs = model(batch_X)
|
||||
loss = criterion(outputs, batch_y)
|
||||
val_loss += loss.item()
|
||||
|
||||
_, predicted = torch.max(outputs, 1)
|
||||
total += batch_y.size(0)
|
||||
correct += (predicted == batch_y).sum().item()
|
||||
|
||||
avg_train_loss = epoch_loss / len(train_loader)
|
||||
avg_val_loss = val_loss / len(val_loader)
|
||||
val_acc = 100 * correct / total
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
|
||||
if (epoch + 1) % 5 == 0:
|
||||
update_status(f"Swin-UNet Epoch {epoch+1}/{epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%, LR: {lr:.6f}", 80 + (epoch / epochs) * 10)
|
||||
print(f"[SWIN-UNET] Epoch {epoch+1}/{epochs} - Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%")
|
||||
|
||||
# Early stopping check
|
||||
if avg_val_loss < best_val_loss:
|
||||
best_val_loss = avg_val_loss
|
||||
patience_counter = 0
|
||||
else:
|
||||
patience_counter += 1
|
||||
if patience_counter >= patience:
|
||||
print(f"[SWIN-UNET] Early stopping at epoch {epoch+1} (best val loss: {best_val_loss:.4f})")
|
||||
update_status(f"Swin-UNet early stopped at epoch {epoch+1}", 90)
|
||||
break
|
||||
|
||||
model = model.cpu()
|
||||
model.device_used = str(device)
|
||||
|
||||
Reference in New Issue
Block a user