# train_ad_qwen_vl.py import argparse import json import numpy as np import pandas as pd import torch from torch.utils.data import DataLoader from accelerate import Accelerator from accelerate.logging import get_logger from Test_TSB import PASS_LIST,TSB_test from evaluation.metrics import fast_get_metrics from model.Vision_encoder.V_encoder import V_model from loss.loss import load_balance_loss from model.TS_encoder.ts_model import TS_Model from model.TS_encoder.config import default_config_t from dataset.dataloader import AnomalyDataset,collate_fn import logging from tqdm.auto import tqdm import os from datetime import datetime from model.VETime import VETIME from Test_TSB import EarlyStopping from functools import partial logging.basicConfig(level=logging.INFO) logger = get_logger(__name__) os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ["TRANSFORMERS_VERBOSITY"] = "error" import torch torch.cuda.empty_cache() def main(args): accelerator = Accelerator( mixed_precision="bf16", gradient_accumulation_steps=4, log_with="tensorboard", project_dir="./output/logs" ) logger.info(f"Using {accelerator.num_processes} {'GPUs' if accelerator.num_processes > 1 else 'CPU'}") vision_model = V_model(args.vision_name,unpatch=True) config_v = vision_model.config if 'mae' in args.vision_name: patch_size=config_v['patch_size'] else: patch_size=config_v.patch_size ts_model = TS_Model(default_config_t) if args.ts_path!=None: state_ts_dict = torch.load(args.ts_path, map_location='cpu')['model_state_dict'] ts_model.load_state_dict(state_ts_dict) model = VETIME(config_v,vision_model,default_config_t,ts_model,args.model_name) if args.vetime_path!=None: state_dict = torch.load(args.vetime_path, map_location='cpu') model.load_state_dict(state_dict) del vision_model,ts_model collatefn = partial(collate_fn, patch_size=patch_size) train_dataset = AnomalyDataset(args.dataset_path, patch_size=patch_size, split="train") train_loader = DataLoader(train_dataset, batch_size=args.batch_size, collate_fn=collatefn, shuffle=False, num_workers=args.num_workers, pin_memory=True, drop_last=True) trainable_params = (param for param in model.parameters() if param.requires_grad) optimizer = torch.optim.Adam(trainable_params,lr=1e-4, weight_decay=1e-2) model, optimizer,train_loader = accelerator.prepare( model, optimizer, train_loader ) model.train() global_step = 0 epochs = args.num_epochs output = [] device = accelerator.device data_setting=args.data_setting img_size=data_setting['img_size'] name_save=f'./output/{args.model_name}__{img_size}_best.pth' early_stopping = EarlyStopping(patience=4, verbose=True, path=name_save) output_path0=f'./output/score/uni/{args.model_name}_train' os.makedirs(output_path0, exist_ok=True) for epoch in range(epochs): model.train() total_loss = 0 all_probs,all_preds, all_labels = [], [], [] progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}[Train]", disable=not accelerator.is_local_main_process) j=0 for batch in progress_bar: labels = batch["labels"] images = batch["image"] # (B, C, H, W) time_series, att_mask = batch['time_series'],batch['attention_mask'] mask = batch['mask'] period = batch['period'] p_value = batch['padding_value'] if labels.shape[1]>model.MAX_L: data_splits = model.split_data(images, time_series, att_mask, labels) loss1=0 loss2=0 logits_list=[] for data_part in data_splits: img_part, ts_part, att_mask, label_part = data_part images, init_img_size = model.vit_encoder.fold_image(img_part, period,p_value ,**data_setting) local_embeddings1, m_w, loss_cl,local_embeddings2 = model(images, ts_part, att_mask, init_img_size,label_part) loss01, logit = model.anomaly_detection_loss(local_embeddings1, label_part) loss02, rec= model.weighted_reconstruction_loss(local_embeddings2, ts_part, att_mask, label_part) loss2=loss2+loss02 loss2 = loss2 + 0.1*loss_cl + 0.2*load_balance_loss(m_w) loss1 = loss1+loss01 logits_list.append(logit) logits = torch.cat(logits_list, dim=1) else: images, init_img_size = model.vit_encoder.fold_image(images, period,p_value ,**data_setting) local_embeddings1, m_w, loss_cl,local_embeddings2 = model(images, time_series, att_mask, init_img_size,labels) loss1, logits = model.anomaly_detection_loss(local_embeddings1, labels) loss2, rec= model.weighted_reconstruction_loss(local_embeddings2, time_series, att_mask, labels) loss2 = loss2 + 0.2*load_balance_loss(m_w)+0.1*loss_cl accelerator.backward(loss1+loss2) global_step+=1 if global_step % accelerator.gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() total_loss += loss1.item()+loss2.item() progress_bar.set_postfix({"loss": total_loss}) probs = torch.softmax(logits, dim=-1)[:, :,1] preds = (probs > 0.5).float() probs,preds, labels = accelerator.gather_for_metrics((probs,preds, labels)) j+=1 for i in range(probs.shape[0]): all_probs.append(probs[i].detach().cpu().numpy().reshape(-1)) all_preds.append(preds[i].detach().cpu().numpy().reshape(-1)) all_labels.append(labels[i].detach().cpu().numpy().reshape(-1).astype(int)) del images,logits, loss1, probs, preds, labels, loss2 torch.cuda.empty_cache() all_probs = np.concatenate(all_probs) # (total_points,) all_preds= np.concatenate(all_preds) all_labels = np.concatenate(all_labels) if np.any(np.isnan(all_probs)): print("⚠️ Warning: all_probs contains NaN values!") train_metrics = fast_get_metrics(all_probs, all_labels) avg_train_loss = total_loss / len(train_loader) accelerator.log({"epoch_train_loss": avg_train_loss}, step=epoch) print(f"\n[Epoch {epoch + 1}/{epochs}] 🟩 Training Summary:") print(f" Avg Loss: {avg_train_loss:.4f}") for k, v in train_metrics.items(): print(f" Train {k}: {v:.4f}") if (epoch+1) % 2 == 0 or epoch == epochs - 1: model.eval() avg_val_loss=TSB_test(model,args.model_name,args.data_setting,device,dataset_setting=PASS_LIST) accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) timestamp = datetime.now().strftime("%m%d-%H") name_save=f'./output/{args.model_name}__{img_size}_{avg_val_loss:.4f}_{timestamp}.pth' torch.save(unwrapped_model.state_dict(), name_save) logger.info(f"Best model saved at epoch {epoch+1} with val_loss={avg_val_loss:.4f}") epoch_log = { "epoch": epoch + 1, "train_loss": round(avg_train_loss, 6), "train_metrics": {k: round(v, 6) for k, v in train_metrics.items()}, "val_loss": round(avg_val_loss, 6) if avg_val_loss is not None else None, } output.append(epoch_log) early_stopping(avg_val_loss, model) if early_stopping.early_stop: print("Early stopping triggered.") loss_all=TSB_test(model,args.model_name,args.data_setting,device,dataset_setting=PASS_LIST) print(loss_all) accelerator.end_training() logger.info("Training completed!") return output if __name__ == "__main__": DATA_INIT_SETTING = { "img_size": 224, "T_sqrt": False, } parser = argparse.ArgumentParser(description='Training') parser.add_argument('--dataset_path', default='./dataset' , type=str, help='Path to the data file') parser.add_argument('--dataset_test_dir', type=str, default='./dataset/TSB-AD/Datasets/TSB-AD-U') parser.add_argument('--file_list', type=str, default='./dataset/TSB-AD/Datasets/File_List/TSB-AD-U.csv') parser.add_argument('--model_name', default= 'VETime', type=str, help='Name of the model') parser.add_argument('--seed', type=int, default=64, help='Random seed') parser.add_argument('--batch_size', type=int, default=32, help='Batch size') parser.add_argument('--num_workers', type=int, default=5, help='Number of data loader workers') parser.add_argument('--num_epochs', type=int, default=4, help='epochs number') parser.add_argument('--output_file_path', default='./output/result.json',type=str, help='Path to the output file') parser.add_argument('--keep_idx_path', type=str, required=False, help='Path to the keep idx file') parser.add_argument('--device', type=str, default='auto', help='Device to use for evaluation') parser.add_argument('--data_setting', type=str, default=DATA_INIT_SETTING, help='Device to use for evaluation') parser.add_argument('--vision_path', type=str, default='./checkpoints/weight_v' , help='vision_weight') parser.add_argument('--ts_path', type=str, default=None , help='TS_weight') parser.add_argument('--vetime_path', type=str, default=None , help='VETime_weight') parser.add_argument('--vision_name', type=str, default='mae_visualize_base.pth' , help='vision_weight_name') args = parser.parse_args() output_file_path = args.output_file_path.replace('result.json', f'{args.model_name.replace("/", "-")}_result.json') results = main(args) with open(output_file_path, 'w') as f: json.dump(results, f, indent=4)