gpbench 1.0.0__py3-none-any.whl

gpbench/method_class/GEFormer/GEFormer_he_class.py

@@ -0,0 +1,137 @@
+import os
+import time
+import psutil
+import random
+import torch
+import numpy as np
+import optuna
+from sklearn.model_selection import StratifiedKFold, train_test_split
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+from torch.utils.data import DataLoader, TensorDataset
+from optuna.exceptions import TrialPruned
+from .gMLP_class import GEFormer
+
+def run_nested_cv_with_early_stopping(
+    data, label, nsnp,
+    learning_rate, patience, batch_size,
+    epoch=1000
+):
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    print("Starting 10-fold cross-validation...")
+
+    kf = StratifiedKFold(n_splits=10, shuffle=True, random_state=42)
+
+    all_acc, all_prec, all_rec, all_f1 = [], [], [], []
+
+    num_classes = len(np.unique(label))
+
+    for fold, (train_index, test_index) in enumerate(kf.split(data, label)):
+        print(f"Running fold {fold}...")
+        process = psutil.Process(os.getpid())
+        fold_start_time = time.time()
+
+        X_train, X_test = data[train_index], data[test_index]
+        y_train, y_test = label[train_index], label[test_index]
+
+        X_train_sub, X_valid, y_train_sub, y_valid = train_test_split(
+            X_train, y_train, test_size=0.1, stratify=y_train, random_state=42
+        )
+
+        x_train = torch.from_numpy(X_train_sub).float().to(device)
+        y_train = torch.from_numpy(y_train_sub).long().to(device)
+        x_valid = torch.from_numpy(X_valid).float().to(device)
+        y_valid = torch.from_numpy(y_valid).long().to(device)
+        x_test  = torch.from_numpy(X_test).float().to(device)
+        y_test_tensor = torch.from_numpy(y_test).long().to(device)
+
+        train_loader = DataLoader(
+            TensorDataset(x_train, y_train), batch_size, shuffle=True
+        )
+        valid_loader = DataLoader(
+            TensorDataset(x_valid, y_valid), batch_size, shuffle=False
+        )
+        test_loader = DataLoader(
+            TensorDataset(x_test, y_test_tensor), batch_size, shuffle=False
+        )
+
+        model = GEFormer(nsnp=nsnp, num_classes=num_classes)
+        model.train_model(
+            train_loader, valid_loader,
+            epoch, learning_rate, patience, device
+        )
+
+        logits = model.predict(test_loader)
+        y_pred = np.argmax(logits, axis=1)
+
+        acc = accuracy_score(y_test, y_pred)
+        prec, rec, f1, _ = precision_recall_fscore_support(
+            y_test, y_pred, average="macro", zero_division=0
+        )
+
+        if np.isnan(f1) or f1 <= 0:
+            print(f"Fold {fold} resulted in NaN or zero F1, pruning trial")
+            raise TrialPruned()
+
+        all_acc.append(acc)
+        all_prec.append(prec)
+        all_rec.append(rec)
+        all_f1.append(f1)
+
+        fold_time = time.time() - fold_start_time
+        fold_cpu_mem = process.memory_info().rss / 1024**2
+
+        print(
+            f"Fold {fold}: "
+            f"ACC={acc:.4f}, PREC={prec:.4f}, REC={rec:.4f}, F1={f1:.4f}, "
+            f"Time={fold_time:.2f}s, CPU={fold_cpu_mem:.2f}MB"
+        )
+
+    print("\n===== CV Summary =====")
+    print(f"ACC : {np.mean(all_acc):.4f} ± {np.std(all_acc):.4f}")
+    print(f"PREC: {np.mean(all_prec):.4f} ± {np.std(all_prec):.4f}")
+    print(f"REC : {np.mean(all_rec):.4f} ± {np.std(all_rec):.4f}")
+    print(f"F1  : {np.mean(all_f1):.4f} ± {np.std(all_f1):.4f}")
+
+    return np.mean(all_f1) if all_f1 else 0.0
+
+def set_seed(seed=42):
+    random.seed(seed)
+    np.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
+
+def Hyperparameter(data, label, nsnp):
+    set_seed(42)
+
+    def objective(trial):
+        learning_rate = trial.suggest_float(
+            "learning_rate", 1e-4, 0.1, log=True
+        )
+        batch_size = trial.suggest_categorical(
+            "batch_size", [32, 64, 128]
+        )
+        patience = trial.suggest_int("patience", 1, 10)
+
+        try:
+            f1_score = run_nested_cv_with_early_stopping(
+                data=data,
+                label=label,
+                nsnp=nsnp,
+                learning_rate=learning_rate,
+                patience=patience,
+                batch_size=batch_size
+            )
+        except TrialPruned:
+            return float("-inf")
+
+        return f1_score
+
+    study = optuna.create_study(direction="maximize")
+    study.optimize(objective, n_trials=20)
+
+    print("best params:", study.best_params)
+    print("successfully")
+    return study.best_params

gpbench/method_class/GEFormer/__init__.py

@@ -0,0 +1,5 @@
+from .GEFormer_class import GEFormer_class
+
+GEFormer = GEFormer_class
+
+__all__ = ["GEFormer","GEFormer_class"]

gpbench/method_class/GEFormer/gMLP_class.py

@@ -0,0 +1,357 @@
+from random import randrange
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange, Reduce
+
+# functions
+
+def exists(val):
+    return val is not None
+
+def pair(val):
+    return (val, val) if not isinstance(val, tuple) else val
+
+def dropout_layers(layers, prob_survival):
+    if prob_survival == 1:
+        return layers
+
+    num_layers = len(layers)
+    to_drop = torch.zeros(num_layers).uniform_(0., 1.) > prob_survival
+
+    # make sure at least one layer makes it
+    if all(to_drop):
+        rand_index = randrange(num_layers)
+        to_drop[rand_index] = False
+
+    layers = [layer for (layer, drop) in zip(layers, to_drop) if not drop]
+    return layers
+
+
+# helper classes
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(x) + x
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+class Attention(nn.Module):
+    def __init__(self, dim_in, dim_out, dim_inner, causal = False):
+        super().__init__()
+        self.scale = dim_inner ** -0.5
+        self.causal = causal
+
+        self.to_qkv = nn.Linear(dim_in, dim_inner * 3, bias = False)
+        self.to_out = nn.Linear(dim_inner, dim_out)
+
+    def forward(self, x):
+        device = x.device
+        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
+        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+        if self.causal:
+            mask = torch.ones(sim.shape[-2:], device = device).triu(1).bool()
+            sim.masked_fill_(mask[None, ...], -torch.finfo(q.dtype).max)
+
+        attn = sim.softmax(dim = -1)
+        out = einsum('b i j, b j d -> b i d', attn, v)
+        return self.to_out(out)
+
+class SpatialGatingUnit(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_seq,
+        causal = False,
+        act = nn.Identity(),
+        heads = 1,
+        init_eps = 1e-3,
+        circulant_matrix = False
+    ):
+        super().__init__()
+        dim_out = dim // 2
+        self.heads = heads
+        self.causal = causal
+        self.norm = nn.LayerNorm(dim_out)
+
+        self.act = act
+
+        # parameters
+
+        if circulant_matrix:
+            self.circulant_pos_x = nn.Parameter(torch.ones(heads, dim_seq))
+            self.circulant_pos_y = nn.Parameter(torch.ones(heads, dim_seq))
+
+        self.circulant_matrix = circulant_matrix
+        shape = (heads, dim_seq,) if circulant_matrix else (heads, dim_seq, dim_seq)
+        weight = torch.zeros(shape)
+
+        self.weight = nn.Parameter(weight)
+        init_eps /= dim_seq
+        nn.init.uniform_(self.weight, -init_eps, init_eps)
+
+        self.bias = nn.Parameter(torch.ones(heads, dim_seq))
+
+    def forward(self, x, gate_res = None):
+        device, n, h = x.device, x.shape[1], self.heads
+
+        res, gate = x.chunk(2, dim = -1)
+        gate = self.norm(gate)
+
+        weight, bias = self.weight, self.bias
+
+        if self.circulant_matrix:
+            # build the circulant matrix
+
+            dim_seq = weight.shape[-1]
+            weight = F.pad(weight, (0, dim_seq), value = 0)
+            weight = repeat(weight, '... n -> ... (r n)', r = dim_seq)
+            weight = weight[:, :-dim_seq].reshape(h, dim_seq, 2 * dim_seq - 1)
+            weight = weight[:, :, (dim_seq - 1):]
+
+            # give circulant matrix absolute position awareness
+
+            pos_x, pos_y = self.circulant_pos_x, self.circulant_pos_y
+            weight = weight * rearrange(pos_x, 'h i -> h i ()') * rearrange(pos_y, 'h j -> h () j')
+
+        if self.causal:
+            weight, bias = weight[:, :n, :n], bias[:, :n]
+            mask = torch.ones(weight.shape[-2:], device = device).triu_(1).bool()
+            mask = rearrange(mask, 'i j -> () i j')
+            weight = weight.masked_fill(mask, 0.)
+
+        gate = rearrange(gate, 'b n (h d) -> b h n d', h = h)
+
+        gate = einsum('b h n d, h m n -> b h m d', gate, weight)
+        gate = gate + rearrange(bias, 'h n -> () h n ()')
+
+        gate = rearrange(gate, 'b h n d -> b n (h d)')
+
+        if exists(gate_res):
+            gate = gate + gate_res
+
+        return self.act(gate) * res
+
+class gMLPBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_ff,
+        seq_len,
+        heads = 1,
+        attn_dim = None,
+        causal = False,
+        act = nn.Identity(),
+        circulant_matrix = False
+    ):
+        super().__init__()
+        self.proj_in = nn.Sequential(
+            nn.Linear(dim, dim_ff),
+            nn.GELU()
+        )
+
+        self.attn = Attention(dim, dim_ff // 2, attn_dim, causal) if exists(attn_dim) else None
+
+        self.sgu = SpatialGatingUnit(dim_ff, seq_len, causal, act, heads, circulant_matrix = circulant_matrix)
+        self.proj_out = nn.Linear(dim_ff // 2, dim)
+
+    def forward(self, x):
+        gate_res = self.attn(x) if exists(self.attn) else None
+        x = self.proj_in(x)
+        x = self.sgu(x, gate_res=gate_res)
+        x = self.proj_out(x)
+        return x
+
+
+# main classes
+class gMLPVision(nn.Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        snp_len,
+        heads = 1,
+        ff_mult = 4,
+        channels = 1,
+        attn_dim = None,
+        prob_survival = 1.
+    ):
+        super().__init__()
+        assert (dim % heads) == 0, 'dimension must be divisible by number of heads'
+
+        image_height, image_width = pair(image_size)       
+        patch_height, patch_width = pair(patch_size)      
+        #assert (image_height % patch_height) == 0 and (image_width % patch_width) == 0, 'image height and width must be divisible by patch size'
+        #num_patches = (image_height[0] // patch_height[0]) * (image_width[1] // patch_width[1])
+        num_patches = 200 
+        dim_ff = dim * ff_mult
+
+        self.to_patch_embed = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (c p1 p2)', p1 = patch_height, p2 = patch_width),
+            nn.Linear(1*snp_len*1, dim) 
+        )
+        
+       
+        self.prob_survival = prob_survival
+
+        self.layers = nn.ModuleList([Residual(PreNorm(dim, gMLPBlock(dim = dim, heads = heads, dim_ff = dim_ff, seq_len = num_patches, attn_dim = attn_dim))) for i in range(depth)])
+
+        self.to_logits = nn.Sequential(
+            nn.LayerNorm(dim),
+            Reduce('b n d -> b d', 'mean'),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, x):
+        x = self.to_patch_embed(x)
+        layers = self.layers if not self.training else dropout_layers(self.layers, self.prob_survival)
+        x = nn.Sequential(*layers)(x)
+        return self.to_logits(x)
+
+class EarlyStopping:
+    def __init__(self, patience=10, delta=0):
+        self.patience = patience
+        self.delta = delta
+        self.best_score = None
+        self.counter = 0
+        self.early_stop = False
+
+    def __call__(self, score):
+        if self.best_score is None:
+            self.best_score = score
+        elif score < self.best_score + self.delta:
+            self.counter += 1
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.counter = 0
+
+def exists(val):
+    return val is not None
+
+class GEFormer(nn.Module):
+    def __init__(self, nsnp, num_classes):
+        super(GEFormer, self).__init__()
+
+        self.gmlp = gMLPVision(
+            image_size=(nsnp, 1),
+            patch_size=(nsnp, 1),
+            num_classes=126,
+            dim=126,
+            depth=1,
+            snp_len=nsnp
+        )
+
+        self.MLP = nn.Sequential(
+            nn.Linear(126, 128),
+            nn.LeakyReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(128, 64),
+            nn.LeakyReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(64, num_classes)
+        )
+
+        self.numsnp = nsnp
+
+    def forward(self, x):
+        x = x.view(x.size(0), 1, self.numsnp, 1)
+        x = self.gmlp(x)
+        logits = self.MLP(x)
+        return logits
+
+
+    def train_model(
+        self, train_loader, valid_loader,
+        num_epochs, learning_rate, patience, device
+    ):
+        optimizer = torch.optim.Adam(
+            self.parameters(), lr=learning_rate, weight_decay=1e-4
+        )
+
+        criterion = nn.CrossEntropyLoss() 
+        self.to(device)
+
+        best_loss = float('inf')
+        best_state = None
+        trigger_times = 0
+
+        for epoch in range(num_epochs):
+            self.train()
+            train_loss = 0.0
+
+            for inputs, labels in train_loader:
+                inputs = inputs.to(device)
+                labels = labels.to(device).long()
+
+                optimizer.zero_grad()
+                logits = self(inputs)
+                loss = criterion(logits, labels)
+                loss.backward()
+                optimizer.step()
+
+                train_loss += loss.item() * inputs.size(0)
+
+            self.eval()
+            valid_loss = 0.0
+            with torch.no_grad():
+                for inputs, labels in valid_loader:
+                    inputs = inputs.to(device)
+                    labels = labels.to(device).long()
+                    logits = self(inputs)
+                    loss = criterion(logits, labels)
+                    valid_loss += loss.item() * inputs.size(0)
+
+            train_loss /= len(train_loader.dataset)
+            valid_loss /= len(valid_loader.dataset)
+
+            if valid_loss < best_loss:
+                best_loss = valid_loss
+                best_state = self.state_dict()
+                trigger_times = 0
+            else:
+                trigger_times += 1
+                if trigger_times >= patience:
+                    print(f"Early stopping at epoch {epoch+1}")
+                    break
+
+        if best_state is not None:
+            self.load_state_dict(best_state)
+
+        return best_loss
+
+    def predict(self, test_loader):
+        self.eval()
+        logits_all = []
+
+        with torch.no_grad():
+            for inputs, _ in test_loader:
+                inputs = inputs.to(next(self.parameters()).device)
+                logits = self(inputs)
+                logits_all.append(logits.cpu().numpy())
+
+        return np.concatenate(logits_all, axis=0)  # logits
+
+