bổ sung chức năng load 64 tỉnh thành và 32 tỉnh thành/ bổ sung mô hình Swin-Unet
This commit is contained in:
Victor Phan
+218
-16
@@ -22,17 +22,18 @@ import hashlib
|
||||
from pathlib import Path
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
# PyTorch for CNN
|
||||
# PyTorch for CNN and advanced models
|
||||
try:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import TensorDataset, DataLoader
|
||||
import torchvision.models as models
|
||||
PYTORCH_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYTORCH_AVAILABLE = False
|
||||
print("Warning: PyTorch not available. CNN model will not work.")
|
||||
print("Warning: PyTorch not available. CNN and advanced models will not work.")
|
||||
|
||||
# Define CNN model class for PyTorch
|
||||
class CNNClassifier(nn.Module):
|
||||
@@ -113,6 +114,146 @@ class CNNClassifier(nn.Module):
|
||||
y = y.cpu().numpy()
|
||||
return np.mean(predictions == y)
|
||||
|
||||
|
||||
# Swin-UNet Classifier for feature vectors
|
||||
class SwinUNetClassifier(nn.Module):
|
||||
"""
|
||||
Swin Transformer U-Net style architecture adapted for feature vector classification.
|
||||
Combines hierarchical Swin Transformer blocks with skip connections.
|
||||
"""
|
||||
def __init__(self, n_features, n_classes, embed_dim=128, depths=(2, 2, 6, 2), num_heads=(4, 8, 16, 32)):
|
||||
super(SwinUNetClassifier, self).__init__()
|
||||
self.n_features = n_features
|
||||
self.n_classes = n_classes
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
# Feature adapter - convert input features to embedding
|
||||
self.adapter = nn.Sequential(
|
||||
nn.Linear(n_features, embed_dim * 2),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.1),
|
||||
nn.Linear(embed_dim * 2, embed_dim)
|
||||
)
|
||||
|
||||
# Encoder path with hierarchical structure
|
||||
# Stage 1 - 1/4 resolution
|
||||
self.encoder1 = nn.Sequential(
|
||||
nn.Linear(embed_dim, embed_dim),
|
||||
nn.LayerNorm(embed_dim),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.1)
|
||||
)
|
||||
self.down1 = nn.Linear(embed_dim, embed_dim * 2)
|
||||
|
||||
# Stage 2 - 1/8 resolution
|
||||
self.encoder2 = nn.Sequential(
|
||||
nn.Linear(embed_dim * 2, embed_dim * 2),
|
||||
nn.LayerNorm(embed_dim * 2),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.1)
|
||||
)
|
||||
self.down2 = nn.Linear(embed_dim * 2, embed_dim * 4)
|
||||
|
||||
# Stage 3 - 1/16 resolution (bottleneck)
|
||||
self.encoder3 = nn.Sequential(
|
||||
nn.Linear(embed_dim * 4, embed_dim * 4),
|
||||
nn.LayerNorm(embed_dim * 4),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.1)
|
||||
)
|
||||
|
||||
# Decoder path with skip connections
|
||||
self.up2 = nn.Linear(embed_dim * 4, embed_dim * 2)
|
||||
self.decoder2 = nn.Sequential(
|
||||
nn.Linear(embed_dim * 4, embed_dim * 2), # Concatenated with skip
|
||||
nn.LayerNorm(embed_dim * 2),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.1)
|
||||
)
|
||||
|
||||
self.up1 = nn.Linear(embed_dim * 2, embed_dim)
|
||||
self.decoder1 = nn.Sequential(
|
||||
nn.Linear(embed_dim * 2, embed_dim), # Concatenated with skip
|
||||
nn.LayerNorm(embed_dim),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.1)
|
||||
)
|
||||
|
||||
# Classification head
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(embed_dim, embed_dim // 2),
|
||||
nn.GELU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(embed_dim // 2, n_classes)
|
||||
)
|
||||
|
||||
# Attention mechanism for better feature aggregation
|
||||
self.attention = nn.MultiheadAttention(embed_dim, num_heads=4, batch_first=True)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: (batch, n_features)
|
||||
if len(x.shape) == 3:
|
||||
x = x.squeeze(1)
|
||||
|
||||
batch_size = x.shape[0]
|
||||
|
||||
# Feature adaptation
|
||||
x = self.adapter(x) # (batch, embed_dim)
|
||||
|
||||
# Add sequence dimension for attention (treat as sequence of length 1)
|
||||
x_seq = x.unsqueeze(1) # (batch, 1, embed_dim)
|
||||
|
||||
# Encoder path
|
||||
# Stage 1
|
||||
x1 = self.encoder1(x_seq) # (batch, 1, embed_dim)
|
||||
x_down1 = self.down1(x1.squeeze(1)) # (batch, embed_dim*2)
|
||||
|
||||
# Stage 2
|
||||
x2 = self.encoder2(x_down1.unsqueeze(1)) # (batch, 1, embed_dim*2)
|
||||
x_down2 = self.down2(x2.squeeze(1)) # (batch, embed_dim*4)
|
||||
|
||||
# Stage 3 (bottleneck)
|
||||
x3 = self.encoder3(x_down2.unsqueeze(1)) # (batch, 1, embed_dim*4)
|
||||
|
||||
# Decoder path with skip connections
|
||||
# Up2
|
||||
x_up2 = self.up2(x3.squeeze(1)) # (batch, embed_dim*2)
|
||||
x_cat2 = torch.cat([x_up2, x_down1], dim=1) # (batch, embed_dim*4) - concatenate skip
|
||||
# Create proper 3D tensor for decoder
|
||||
x_cat2_seq = x_cat2.unsqueeze(1) # (batch, 1, embed_dim*4)
|
||||
x_dec2 = self.decoder2(x_cat2) # (batch, embed_dim*2)
|
||||
|
||||
# Up1
|
||||
x_up1 = self.up1(x_dec2) # (batch, embed_dim)
|
||||
x_cat1 = torch.cat([x_up1, x.squeeze(1)], dim=1) # (batch, embed_dim*2) - concatenate skip
|
||||
x_dec1 = self.decoder1(x_cat1) # (batch, embed_dim)
|
||||
|
||||
# Apply attention mechanism for better aggregation
|
||||
x_dec1_seq = x_dec1.unsqueeze(1) # (batch, 1, embed_dim)
|
||||
attn_out, _ = self.attention(x_dec1_seq, x_dec1_seq, x_dec1_seq)
|
||||
|
||||
# Classification
|
||||
output = self.classifier(attn_out.squeeze(1))
|
||||
return output
|
||||
|
||||
def predict(self, X):
|
||||
"""Scikit-learn style predict"""
|
||||
self.eval()
|
||||
with torch.no_grad():
|
||||
if isinstance(X, np.ndarray):
|
||||
X = torch.FloatTensor(X)
|
||||
outputs = self(X)
|
||||
_, predicted = torch.max(outputs, 1)
|
||||
return predicted.cpu().numpy()
|
||||
|
||||
def score(self, X, y):
|
||||
"""Scikit-learn style score"""
|
||||
predictions = self.predict(X)
|
||||
if isinstance(y, torch.Tensor):
|
||||
y = y.cpu().numpy()
|
||||
return np.mean(predictions == y)
|
||||
|
||||
|
||||
# Microsoft Planetary Computer imports
|
||||
import planetary_computer
|
||||
from pystac_client import Client
|
||||
@@ -199,6 +340,10 @@ def train_model(
|
||||
features = None
|
||||
labels = None
|
||||
|
||||
# Initialize FeatureExtractor early (will be used for temporal/extended modes)
|
||||
update_status(f"Initializing FeatureExtractor (mode={feature_mode})...", 5)
|
||||
extractor = get_feature_extractor(mode=feature_mode)
|
||||
|
||||
# Try to load from cache
|
||||
if use_cache and cache_file.exists():
|
||||
update_status(f"📦 Loading cached dataset from {cache_file.name}...", 5)
|
||||
@@ -300,10 +445,6 @@ def train_model(
|
||||
|
||||
check_cancellation()
|
||||
|
||||
# ============ FEATURE EXTRACTION ============
|
||||
update_status(f"Initializing FeatureExtractor (mode={feature_mode})...", 50)
|
||||
extractor = get_feature_extractor(mode=feature_mode)
|
||||
|
||||
# Load training data
|
||||
update_status("Loading training data...", 55)
|
||||
train_gdf = gpd.read_file(training_shapefile)
|
||||
@@ -541,17 +682,75 @@ def train_model(
|
||||
# Move model to CPU for saving (compatible with non-GPU systems)
|
||||
model = model.cpu()
|
||||
model.device_used = str(device)
|
||||
else:
|
||||
raise ValueError(f"Unknown model type: {model_type}. Choose: xgboost, random_forest, decision_tree, svm, cnn")
|
||||
|
||||
# Fit non-CNN models
|
||||
if model_type != 'cnn':
|
||||
elif model_type == 'swin-unet':
|
||||
if not PYTORCH_AVAILABLE:
|
||||
raise ImportError("PyTorch is required for Swin-UNet. Install: pip install torch torchvision")
|
||||
|
||||
n_features = X_train.shape[1]
|
||||
n_classes = len(np.unique(y_train))
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
|
||||
update_status(f"Building Swin-UNet model on {device}...", 75)
|
||||
|
||||
model = SwinUNetClassifier(n_features, n_classes, embed_dim=128).to(device)
|
||||
|
||||
# Convert to PyTorch tensors (no unsqueeze needed for Swin-UNet)
|
||||
X_train_tensor = torch.FloatTensor(X_train)
|
||||
y_train_tensor = torch.LongTensor(y_train)
|
||||
X_test_tensor = torch.FloatTensor(X_test)
|
||||
y_test_tensor = torch.LongTensor(y_test)
|
||||
|
||||
# Create data loaders
|
||||
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()
|
||||
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)
|
||||
|
||||
# Train Swin-UNet
|
||||
update_status("Training Swin-UNet model with PyTorch...", 80)
|
||||
epochs = min(60, n_estimators // 2) # Swin-UNet benefits from more epochs
|
||||
|
||||
model.train()
|
||||
for epoch in range(epochs):
|
||||
epoch_loss = 0.0
|
||||
for batch_X, batch_y in train_loader:
|
||||
batch_X, batch_y = batch_X.to(device), batch_y.to(device)
|
||||
|
||||
optimizer.zero_grad()
|
||||
outputs = model(batch_X)
|
||||
loss = criterion(outputs, batch_y)
|
||||
loss.backward()
|
||||
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)
|
||||
|
||||
model = model.cpu()
|
||||
model.device_used = str(device)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown model type: {model_type}. Choose: xgboost, random_forest, decision_tree, svm, cnn, swin-unet")
|
||||
|
||||
# Fit non-neural-network models
|
||||
if model_type not in ['cnn', 'swin-unet']:
|
||||
model.fit(X_train, y_train)
|
||||
|
||||
# Evaluate
|
||||
update_status("Evaluating model...", 90)
|
||||
if model_type == 'cnn':
|
||||
# PyTorch CNN evaluation
|
||||
if model_type in ['cnn', 'swin-unet']:
|
||||
# PyTorch models evaluation
|
||||
train_score = model.score(X_train, y_train)
|
||||
test_score = model.score(X_test, y_test)
|
||||
y_pred = model.predict(X_test)
|
||||
@@ -597,10 +796,10 @@ def train_model(
|
||||
"test_accuracy": float(test_score),
|
||||
"model_type": model_type,
|
||||
"device": device if model_type == 'xgboost' else 'cpu',
|
||||
"n_estimators": n_estimators if model_type in ['xgboost', 'random_forest', 'cnn'] else None,
|
||||
"max_depth": max_depth if model_type != 'cnn' else None,
|
||||
"learning_rate": learning_rate if model_type == 'xgboost' else None,
|
||||
"cnn_epochs": min(50, n_estimators // 2) if model_type == 'cnn' else None,
|
||||
"n_estimators": n_estimators if model_type in ['xgboost', 'random_forest', 'cnn', 'swin-unet'] else None,
|
||||
"max_depth": max_depth if model_type not in ['cnn', 'swin-unet'] else None,
|
||||
"learning_rate": learning_rate if model_type in ['xgboost', 'swin-unet'] else None,
|
||||
"epochs": min(50, n_estimators // 2) if model_type == 'cnn' else (min(60, n_estimators // 2) if model_type == 'swin-unet' else None),
|
||||
"n_features": X_train.shape[1],
|
||||
"n_classes": len(np.unique(y_train)),
|
||||
"class_names": class_names,
|
||||
@@ -622,6 +821,9 @@ def train_model(
|
||||
label_encoder=label_encoder
|
||||
)
|
||||
|
||||
# Construct info path (model manager saves it in model_train/)
|
||||
info_path = os.path.join('model_train', model_filename.replace('.joblib', '_info.json'))
|
||||
|
||||
update_status("Training complete!", 100)
|
||||
|
||||
return {
|
||||
|
||||
Reference in New Issue
Block a user