ftt_train.py

import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
import toml
from torch.utils.data import DataLoader

from utils.dataset import CensusDataset
from utils.models import FTTransformer

import torch.nn as nn
import torch.optim as optim
import torch
import math

def cus_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
    scale = 1.0 / math.sqrt(query.size(-1))
    attn_weights = torch.matmul(query, key.transpose(-2, -1)) * scale

    if attn_mask is not None:
        attn_weights += attn_mask

    attn_weights = torch.softmax(attn_weights, dim=-1)
    

    if dropout_p > 0.0:
        attn_weights = torch.nn.functional.dropout(
            attn_weights, p=dropout_p)

    output = torch.matmul(attn_weights, value)
    return output

from contextlib import contextmanager

@contextmanager
def use_cus_scaled_dot_product_attention():
    original_sdp_attention = torch.nn.functional.scaled_dot_product_attention

    try:
        torch.nn.functional.scaled_dot_product_attention = cus_scaled_dot_product_attention
        yield
    finally:
        torch.nn.functional.scaled_dot_product_attention = original_sdp_attention

def preprocess_bank_marketing(df):

    df = df.copy()
    
    df['deposit'] = (df['deposit'] == 'yes').astype(int)
    

    df['age_10'] = df['age'] // 10

    education_mapping = {
        'primary': 0,
        'secondary': 1,
        'tertiary': 2,
        'unknown': -1
    }

    df['education'] = df['education'].map(education_mapping)
    

    categorical_features = ['job', 'marital', 'default', 'housing', 'loan', 
                          'contact', 'month', 'poutcome']
    

    le_dict = {}
    for feature in categorical_features:
        le = LabelEncoder()
        df[feature] = le.fit_transform(df[feature])
        le_dict[feature] = le
    

    numerical_features = ['balance', 'day', 'duration', 'campaign', 
                         'pdays', 'previous']
    

    scaler = StandardScaler()
    df[numerical_features] = scaler.fit_transform(df[numerical_features])
    

    main_features = categorical_features + numerical_features

    feature_names = main_features + ['age_10', 'education']
    

    X = df[feature_names].values
    y = df['deposit'].values
    

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    return X_train, X_test, y_train, y_test


df = pd.read_csv('./dataset/bank/bank.csv')
config = toml.load('./config/bank.toml')

X_train, X_test, y_train, y_test = preprocess_bank_marketing(df)

train_dataset = CensusDataset(X_train, y_train)
test_dataset = CensusDataset(X_test, y_test)

train_loader = DataLoader(train_dataset, batch_size=config["train"]["batch_size"], shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=config["train"]["batch_size"], shuffle=False)

device = "cuda"

model = FTTransformer(**config["model"]).to(device)

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=config["train"]["lr"], weight_decay=config["train"]["weight_decay"])

for epoch in range(config["train"]["num_epochs"]):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch} loss: {loss.item()}")
    
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    
    with use_cus_scaled_dot_product_attention():
        for batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion(output, target)
            test_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
    
    test_loss /= len(test_loader.dataset)
    test_accuracy = 100. * correct / total
    print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.2f}%")
    
torch.save(model.state_dict(), config["model_path"])