195 lines
6.1 KiB
Python
Executable File
195 lines
6.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""Train hybrid diffusion on HAI (configurable runnable example)."""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import random
|
||
from pathlib import Path
|
||
from typing import Dict
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
from data_utils import load_split, windowed_batches
|
||
from hybrid_diffusion import (
|
||
HybridDiffusionModel,
|
||
cosine_beta_schedule,
|
||
q_sample_continuous,
|
||
q_sample_discrete,
|
||
)
|
||
from platform_utils import resolve_device, safe_path, ensure_dir
|
||
|
||
|
||
BASE_DIR = Path(__file__).resolve().parent
|
||
REPO_DIR = BASE_DIR.parent.parent
|
||
|
||
DEFAULTS = {
|
||
"data_path": REPO_DIR / "dataset" / "hai" / "hai-21.03" / "train1.csv.gz",
|
||
"split_path": BASE_DIR / "feature_split.json",
|
||
"stats_path": BASE_DIR / "results" / "cont_stats.json",
|
||
"vocab_path": BASE_DIR / "results" / "disc_vocab.json",
|
||
"out_dir": BASE_DIR / "results",
|
||
"device": "auto",
|
||
"timesteps": 1000,
|
||
"batch_size": 8,
|
||
"seq_len": 64,
|
||
"epochs": 1,
|
||
"max_batches": 50,
|
||
"lambda": 0.5,
|
||
"lr": 1e-3,
|
||
"seed": 1337,
|
||
"log_every": 10,
|
||
"ckpt_every": 50,
|
||
}
|
||
|
||
|
||
def load_json(path: str) -> Dict:
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
|
||
|
||
def set_seed(seed: int):
|
||
random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed_all(seed)
|
||
torch.backends.cudnn.deterministic = True
|
||
torch.backends.cudnn.benchmark = False
|
||
|
||
|
||
# 使用 platform_utils 中的 resolve_device 函数
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description="Train hybrid diffusion on HAI.")
|
||
parser.add_argument("--config", default=None, help="Path to JSON config.")
|
||
parser.add_argument("--device", default="auto", help="cpu, cuda, or auto")
|
||
return parser.parse_args()
|
||
|
||
|
||
def resolve_config_paths(config, base_dir: Path):
|
||
keys = ["data_path", "split_path", "stats_path", "vocab_path", "out_dir"]
|
||
for key in keys:
|
||
if key in config:
|
||
# 如果值是字符串，转换为Path对象
|
||
if isinstance(config[key], str):
|
||
path = Path(config[key])
|
||
else:
|
||
path = config[key]
|
||
|
||
if not path.is_absolute():
|
||
config[key] = str((base_dir / path).resolve())
|
||
else:
|
||
config[key] = str(path)
|
||
return config
|
||
|
||
|
||
def main():
|
||
args = parse_args()
|
||
config = dict(DEFAULTS)
|
||
if args.config:
|
||
cfg_path = Path(args.config).resolve()
|
||
config.update(load_json(str(cfg_path)))
|
||
config = resolve_config_paths(config, cfg_path.parent)
|
||
else:
|
||
config = resolve_config_paths(config, BASE_DIR)
|
||
|
||
# 优先使用命令行传入的device参数
|
||
if args.device != "auto":
|
||
config["device"] = args.device
|
||
|
||
set_seed(int(config["seed"]))
|
||
|
||
split = load_split(config["split_path"])
|
||
time_col = split.get("time_column", "time")
|
||
cont_cols = [c for c in split["continuous"] if c != time_col]
|
||
disc_cols = [c for c in split["discrete"] if not c.startswith("attack") and c != time_col]
|
||
|
||
stats = load_json(config["stats_path"])
|
||
mean = stats["mean"]
|
||
std = stats["std"]
|
||
|
||
vocab = load_json(config["vocab_path"])["vocab"]
|
||
vocab_sizes = [len(vocab[c]) for c in disc_cols]
|
||
|
||
device = resolve_device(str(config["device"]))
|
||
print("device", device)
|
||
model = HybridDiffusionModel(cont_dim=len(cont_cols), disc_vocab_sizes=vocab_sizes).to(device)
|
||
opt = torch.optim.Adam(model.parameters(), lr=float(config["lr"]))
|
||
|
||
betas = cosine_beta_schedule(int(config["timesteps"])).to(device)
|
||
alphas = 1.0 - betas
|
||
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
||
|
||
os.makedirs(config["out_dir"], exist_ok=True)
|
||
log_path = os.path.join(config["out_dir"], "train_log.csv")
|
||
with open(log_path, "w", encoding="utf-8") as f:
|
||
f.write("epoch,step,loss,loss_cont,loss_disc\n")
|
||
|
||
total_step = 0
|
||
for epoch in range(int(config["epochs"])):
|
||
for step, (x_cont, x_disc) in enumerate(
|
||
windowed_batches(
|
||
config["data_path"],
|
||
cont_cols,
|
||
disc_cols,
|
||
vocab,
|
||
mean,
|
||
std,
|
||
batch_size=int(config["batch_size"]),
|
||
seq_len=int(config["seq_len"]),
|
||
max_batches=int(config["max_batches"]),
|
||
)
|
||
):
|
||
x_cont = x_cont.to(device)
|
||
x_disc = x_disc.to(device)
|
||
|
||
bsz = x_cont.size(0)
|
||
t = torch.randint(0, int(config["timesteps"]), (bsz,), device=device)
|
||
|
||
x_cont_t, noise = q_sample_continuous(x_cont, t, alphas_cumprod)
|
||
|
||
mask_tokens = torch.tensor(vocab_sizes, device=device)
|
||
x_disc_t, mask = q_sample_discrete(x_disc, t, mask_tokens, int(config["timesteps"]))
|
||
|
||
eps_pred, logits = model(x_cont_t, x_disc_t, t)
|
||
|
||
loss_cont = F.mse_loss(eps_pred, noise)
|
||
loss_disc = 0.0
|
||
for i, logit in enumerate(logits):
|
||
if mask[:, :, i].any():
|
||
loss_disc = loss_disc + F.cross_entropy(
|
||
logit[mask[:, :, i]], x_disc[:, :, i][mask[:, :, i]]
|
||
)
|
||
|
||
lam = float(config["lambda"])
|
||
loss = lam * loss_cont + (1 - lam) * loss_disc
|
||
opt.zero_grad()
|
||
loss.backward()
|
||
opt.step()
|
||
|
||
if step % int(config["log_every"]) == 0:
|
||
print("epoch", epoch, "step", step, "loss", float(loss))
|
||
with open(log_path, "a", encoding="utf-8") as f:
|
||
f.write(
|
||
"%d,%d,%.6f,%.6f,%.6f\n"
|
||
% (epoch, step, float(loss), float(loss_cont), float(loss_disc))
|
||
)
|
||
|
||
total_step += 1
|
||
if total_step % int(config["ckpt_every"]) == 0:
|
||
ckpt = {
|
||
"model": model.state_dict(),
|
||
"optim": opt.state_dict(),
|
||
"config": config,
|
||
"step": total_step,
|
||
}
|
||
torch.save(ckpt, os.path.join(config["out_dir"], "model_ckpt.pt"))
|
||
|
||
torch.save(model.state_dict(), os.path.join(config["out_dir"], "model.pt"))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|