run_PT.py

import numpy as np
import pandas as pd
import torch

from models.PatternTransformer import PT
import os
import time
import random
import argparse
from utils.evaluation import *
from torch.optim import lr_scheduler
from utils.data_process import *
from sklearn.metrics import precision_recall_curve
import logging
import warnings
import torch.nn as nn
from torch.utils.data import DataLoader

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
transformers_logger = logging.getLogger('transformer')
transformers_logger.setLevel(logging.ERROR)

warnings.filterwarnings("ignore", message="You are resizing the embedding layer without providing a pad_to_multiple_of parameter.*")
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)

def train(model, train_loader, epoch, optimizer, scheduler, device):
    loss_func = nn.MSELoss(reduction='mean')
    t = time.time()
    total_loss = []
    model.train()
    print1 = 0
    for x, y, labels in train_loader:
        x, y, labels = [item.float().to(device) for item in [x, y, labels]]
        optimizer.zero_grad()
        if print1 == 0 and epoch == 1:
            print("x.shape, y.shape, labels.shape", x.shape, y.shape, labels.shape)

        pred, vq_loss = model(x)
        loss = loss_func(pred, y)
        loss = loss + vq_loss
        total_loss.append(loss.item())
        loss.backward()
        optimizer.step()
        scheduler.step()
        print1 = 1

    # print("Epoch: {:04d}".format(epoch),
    #       "total_loss: {:.10f}".format(np.mean(total_loss)),
    #       "time: {:.4f}s".format(time.time() - t))


def test(model, test_loader, device):
    model.eval()
    error = []
    all_labels = []
    for x, y, labels in test_loader:
        x, y, labels = x.to(device), y.to(device), labels.to(device)

        with torch.no_grad():
            pred, _ = model(x)
            batch_error = model.compute_batch_error(pred, y)

            error += batch_error.detach().tolist()
            if len(all_labels) <= 0:
                all_labels = labels
            else:
                all_labels = torch.cat((all_labels, labels), dim=0)

    print("all_labels:", all_labels.shape)
    print("error:", len(error))
    all_labels = all_labels.cpu().numpy()
    all_labels = all_labels == 1
    precision, recall, thresholds = precision_recall_curve(all_labels, error)
    f1_scores = 2 * recall * precision / (recall + precision + 1e-10)
    round(f1_scores.max(), 4)

    argmax = f1_scores.argmax()
    best_threshold = thresholds[argmax]

    y_pred = error >= best_threshold

    accuracy = accuracy_score(all_labels, y_pred)

    pre = precision_score(all_labels, y_pred)
    rec = recall_score(all_labels, y_pred)
    f1 = f1_scores[argmax]

    print("accuracy:", accuracy)
    print("precision:", pre)
    print("recall:", rec)
    print("f1_score:", f1)

    return accuracy, pre, rec, f1

def solver(args):
    feature_list, _, train_dataset, test_dataset = get_dataset(args.train_path, args.test_path, args.feature_path, args.window_size)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    model = PT(args.node_num, args.window_size, args.head_num, args.embsize, args.inter_dim, args.pattern_space)
    # print("model:", model)
    model.to(args.device)

    optimizer = torch.optim.Adam(list(model.parameters()), lr=args.lr)
    scheduler = lr_scheduler.StepLR(optimizer, step_size=args.lr_decay, gamma=args.gamma)

    for epoch in range(1, args.epochs+1):
        train(model, train_loader, epoch, optimizer, scheduler, args.device)

    print("windows:{}, pattern_space:{}".format(args.window_size, args.pattern_space))
    accuracy, pre, rec, f1 = test(model, test_loader, device)
    return accuracy, pre, rec, f1



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="The parameter configuration of GDN")

    # Dataset parameters
    parser.add_argument("--train_path", default="./data/SMAP/train_D_1.csv", help="dataset path")
    parser.add_argument("--test_path", default="./data/SMAP/test_D_1.csv", help="dataset path")
    parser.add_argument("--feature_path", default="./data/SMAP/list.txt", help="dataset path")
    parser.add_argument("--window_size", default=8, type=int, help="Time Window length")

    # Model parameters
    parser.add_argument("--node_num", default=25, type=int, help='Data Dimension')
    parser.add_argument("--embsize", default=64, type=int, help='Embedding Size')
    parser.add_argument("--inter_dim", default=128, type=int, help='Intermedia Size')
    parser.add_argument("--pattern_space", default=1, type=int, help='Pattern Space Size')
    parser.add_argument("--head_num", default=4, type=int, help='Attention Head Number')
    parser.add_argument("--out_layer_num", default=1, type=int, help='Output Layer Number')
    parser.add_argument("--random_seed", default=2025, type=int, help='Random Seed')

    # train parameters
    parser.add_argument("--batch_size", default=128, type=int, help="Batch Size")
    parser.add_argument("--lr", default=4e-4, type=float, help="Learning Rate")
    parser.add_argument("--lr_decay", default=200, type=int, help="Learning Rate Decay")
    parser.add_argument("--gamma", default=1.0, type=float, help="Learning Rate Decay Factor")
    parser.add_argument("--epochs", default=100, type=int, help="Training Epochs")

    args = parser.parse_args()
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:0" if use_cuda else 'cpu')
    args.device = device

    random.seed(args.random_seed)
    np.random.seed(args.random_seed)
    torch.manual_seed(args.random_seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    os.environ['PYTHONHASHSEED'] = str(args.random_seed)

    accuracy, pre, rec, f1 = solver(args)