File size: 7,997 Bytes
f884fa5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
import torch
from torch.utils.data import DataLoader, random_split, Subset
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm
import numpy as np
import os
import datetime
import pandas as pd
import matplotlib.pyplot as plt
import math
import joblib
from dataloader import MultiHouseDataset
from hierarchical_diffusion_model import HierarchicalDiffusionModel
if torch.cuda.is_available():
DEVICE = "cuda"
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
print("Using NVIDIA CUDA backend.")
elif torch.backends.mps.is_available():
DEVICE = "mps"
print("Using Apple MPS backend.")
else:
DEVICE = "cpu"
print("Using CPU.")
EPOCHS = 200
LEARNING_RATE = 1e-4
BATCH_SIZE = 512
USE_AMP = True
GRADIENT_CLIP_VAL = 0.1
WINDOW_DURATION = '14_days'
DATA_DIRECTORY = './data/per_house'
NUM_WORKERS = os.cpu_count() // 2
PIN_MEMORY = True
USE_ATTENTION = True
DROPOUT = 0.1
HIDDEN_SIZE = 512
EMBEDDING_DIM = 64
DIFFUSION_TIMESTEPS = 500
DOWNSCALE_FACTOR = 4
def calculate_window_size(duration: str) -> int:
SAMPLES_PER_DAY = 48
mapping = {
'2_days': 2 * SAMPLES_PER_DAY,
'7_days': 7 * SAMPLES_PER_DAY,
'14_days': 14 * SAMPLES_PER_DAY,
'15_days': 15 * SAMPLES_PER_DAY,
'30_days': 30 * SAMPLES_PER_DAY
}
if duration not in mapping:
raise ValueError(f"Invalid WINDOW_DURATION: {duration}")
return mapping[duration]
def denormalize_data(normalized_data, scaler_path='global_scaler.gz'):
scaler = joblib.load(scaler_path)
original_shape = normalized_data.shape
if len(original_shape) == 3:
batch_size, seq_len, features = original_shape
normalized_flat = normalized_data.reshape(-1, features)
denormalized_flat = scaler.inverse_transform(normalized_flat)
return denormalized_flat.reshape(original_shape)
else:
return scaler.inverse_transform(normalized_data)
def moving_average(data, window_size):
return np.convolve(data, np.ones(window_size), 'valid') / window_size
def save_and_plot_loss(loss_dict, title, filepath, window_size=10):
plt.figure(figsize=(12, 6))
for label, losses in loss_dict.items():
pd.DataFrame({label: losses}).to_csv(f"{filepath}_{label.lower().replace(' ', '_')}.csv", index=False)
plt.plot(losses, label=f'Raw {label}', alpha=0.3)
if len(losses) > window_size:
smoothed_losses = moving_average(losses, window_size)
plt.plot(np.arange(window_size - 1, len(losses)), smoothed_losses, label=f'Smoothed {label}')
plt.title(title)
plt.xlabel('Epoch'); plt.ylabel('Loss')
plt.legend(); plt.grid(True)
plt.savefig(f"{filepath}.png"); plt.close()
print(f" Loss plot saved to {filepath}.png")
def train_diffusion(log_dir, model_save_path):
print("--- Starting Hierarchical Diffusion Training ---")
window_size = calculate_window_size(WINDOW_DURATION)
print(f"Using window duration: {WINDOW_DURATION} ({window_size} samples)")
dataset = MultiHouseDataset(
data_dir=DATA_DIRECTORY,
window_size=window_size,
step_size=window_size//2,
limit_to_one_year=False
)
print(f"Dataset loaded: {len(dataset)} samples, {dataset.num_houses} houses, {dataset[0][0].shape[1]} features.")
val_split = 0.1
val_size = int(len(dataset) * val_split)
train_size = len(dataset) - val_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
print(f"Train size: {train_size}, Validation size: {val_size}")
train_dataloader = DataLoader(
train_dataset, batch_size=BATCH_SIZE, shuffle=True,
num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, drop_last=True
)
val_dataloader = DataLoader(
val_dataset, batch_size=BATCH_SIZE*2, shuffle=False,
num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY
)
channel_weights = torch.tensor([1.0, 8.0, 1.0, 1.0], device=DEVICE)
print(f"Using channel weights: {channel_weights}")
model = HierarchicalDiffusionModel(
in_channels=dataset[0][0].shape[1],
num_houses=dataset.num_houses,
downscale_factor=DOWNSCALE_FACTOR,
channel_weights=channel_weights,
embedding_dim=EMBEDDING_DIM,
hidden_dims=[HIDDEN_SIZE // 4, HIDDEN_SIZE // 2, HIDDEN_SIZE],
dropout=DROPOUT,
use_attention=USE_ATTENTION,
num_timesteps=DIFFUSION_TIMESTEPS,
blocks_per_level=3
).to(DEVICE)
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
scaler = GradScaler(enabled=(USE_AMP and DEVICE == "cuda"))
train_losses, val_losses = [], []
best_val_loss = float('inf')
print(f"Starting training for {EPOCHS} epochs...")
for epoch in range(EPOCHS):
model.train()
total_train_loss = 0.0
pbar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{EPOCHS} (Train)")
for clean_data, conditions in pbar:
clean_data = clean_data.to(DEVICE, non_blocking=PIN_MEMORY)
conditions = {k: v.to(DEVICE, non_blocking=PIN_MEMORY) for k, v in conditions.items()}
optimizer.zero_grad(set_to_none=True)
with autocast(enabled=(USE_AMP and DEVICE == "cuda")):
loss = model(clean_data, conditions)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), GRADIENT_CLIP_VAL)
scaler.step(optimizer)
scaler.update()
total_train_loss += loss.item()
pbar.set_postfix({'loss': f'{loss.item():.6f}', 'lr': f'{scheduler.get_last_lr()[0]:.2e}'})
avg_train_loss = total_train_loss / len(train_dataloader)
train_losses.append(avg_train_loss)
model.eval()
total_val_loss = 0.0
with torch.no_grad():
for clean_data, conditions in tqdm(val_dataloader, desc="Validating"):
clean_data = clean_data.to(DEVICE, non_blocking=PIN_MEMORY)
conditions = {k: v.to(DEVICE, non_blocking=PIN_MEMORY) for k, v in conditions.items()}
with autocast(enabled=(USE_AMP and DEVICE == "cuda")):
loss = model(clean_data, conditions)
total_val_loss += loss.item()
avg_val_loss = total_val_loss / len(val_dataloader)
val_losses.append(avg_val_loss)
print(f"Epoch {epoch+1}/{EPOCHS} | Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f}")
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
torch.save(model.state_dict(), model_save_path)
print(f"New best model saved to {model_save_path} (Val Loss: {best_val_loss:.6f})")
scheduler.step()
print("--- Training complete ---")
save_and_plot_loss(
{'Train Loss': train_losses, 'Validation Loss': val_losses},
'Hierarchical Diffusion Model Training & Validation Loss',
os.path.join(log_dir, 'diffusion_loss_curves')
)
return dataset
if __name__ == "__main__":
timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
run_name = f"hierarchical_diffusion_{WINDOW_DURATION}_{timestamp}"
log_dir = os.path.join("./training_logs", run_name)
os.makedirs(log_dir, exist_ok=True)
model_path = os.path.join(log_dir, 'best_hierarchical_model.pth')
print(f"Starting new run: {run_name}")
print(f"Logs and models will be saved to: {log_dir}")
full_dataset = train_diffusion(log_dir=log_dir, model_save_path=model_path)
print("\nTraining and best model saving complete.")
print(f"Model saved to: {model_path}")
print(f"Loss curves saved to: {os.path.join(log_dir, 'diffusion_loss_curves.png')}") |