Update codes

codes/Fine-tuning/cell_type_classification/NuSPIRe_from_scratch.py

@@ -0,0 +1,292 @@
+import torch
+import random
+import numpy as np
+import os
+from torch.utils.tensorboard import SummaryWriter
+import pandas as pd
+from torchvision import transforms
+from torch.utils.data import DataLoader, SubsetRandomSampler
+from tqdm import tqdm
+from transformers import  ViTMAEConfig, ViTForImageClassification
+from torchvision.datasets import ImageFolder
+from sklearn.metrics import accuracy_score, f1_score, precision_score, roc_auc_score, recall_score, confusion_matrix
+from sklearn.preprocessing import label_binarize
+from torch.optim.lr_scheduler import LambdaLR
+import argparse
+
+def set_seeds(seed_value=42, cuda_deterministic=False):
+    """Set seeds for reproducibility."""
+    random.seed(seed_value)
+    os.environ['PYTHONHASHSEED'] = str(seed_value)
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed_value)
+        torch.cuda.manual_seed_all(seed_value)
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    if cuda_deterministic:  # slower, more reproducible
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:  # faster, less reproducible
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+        
+
+def warmup_lr_lambda(current_epoch: int, warmup_epochs: int):
+    if (current_epoch < warmup_epochs):
+        return float(current_epoch + 1) / float(max(1, warmup_epochs))
+    return 1.0
+
+# set up
+parser = argparse.ArgumentParser(description="Setup experiment parameters")
+parser.add_argument('--num', type=int, required=True, help='Number of samples per class')
+parser.add_argument('--device', type=int, default=0, help='CUDA device number (default: 0)')
+parser.add_argument('--rep', type=int, required=True, help='Number of replicate')
+args = parser.parse_args()
+num_samples_per_class = args.num
+device = args.device
+num_repeats = args.rep
+    
+SEED = 42
+DEVICE = torch.device(f"cuda:{device}")
+DATA_DIR  = '../lung5_rep1_cancer_nuclear_image_15micron/'
+BATCH_SIZE = 300
+NUM_EPOCHS = 30
+PORJECT_NAME = f'Nuspire_{num_samples_per_class}_r{num_repeats}_lung5_rep1_Classification'
+set_seeds(SEED)
+folder_name = f'./{PORJECT_NAME}_checkpoint'
+
+if not os.path.exists(folder_name):
+    os.makedirs(folder_name)
+    print(f"'{folder_name}'has been created.")
+else:
+    print(f"'{folder_name}' already exists.")
+
+# Dataset
+transform = transforms.Compose([
+    transforms.Resize((112, 112)),
+    transforms.Grayscale(),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+dataset = ImageFolder(DATA_DIR, transform=transform)
+labels = [dataset[i][1] for i in range(len(dataset))]
+
+# Define train and test sizes
+train_size = int(0.8 * len(dataset))
+valid_size = int(0.1 * len(dataset))
+test_size = len(dataset) - train_size - valid_size
+
+indices = np.arange(len(dataset))
+np.random.shuffle(indices)
+
+# Split
+train_indices = indices[:train_size]
+valid_indices = indices[train_size:train_size + valid_size]
+test_indices = indices[train_size + valid_size:]
+class_1_train_indices = [i for i in train_indices if labels[i] == 1]
+class_2_train_indices = [i for i in train_indices if labels[i] == 2]
+class_0_train_indices = [i for i in train_indices if labels[i] == 0]
+
+for repeat in range(num_repeats):
+    np.random.shuffle(class_1_train_indices)
+    np.random.shuffle(class_2_train_indices)
+    np.random.shuffle(class_0_train_indices)
+
+class_1_train_indices = class_1_train_indices[:num_samples_per_class]
+class_2_train_indices = class_2_train_indices[:num_samples_per_class]
+class_0_train_indices = class_0_train_indices[:num_samples_per_class]
+
+balanced_train_indices = (
+    class_1_train_indices +
+    class_2_train_indices +
+    class_0_train_indices
+)
+np.random.shuffle(balanced_train_indices)
+
+train_sampler = SubsetRandomSampler(balanced_train_indices)
+valid_sampler = SubsetRandomSampler(valid_indices)
+test_sampler = SubsetRandomSampler(test_indices)
+
+# print(balanced_train_indices)
+# print(valid_indices)
+# print(test_indices)
+
+train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers= 4)
+valid_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=valid_sampler, num_workers= 4)
+test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler, num_workers= 4)
+
+
+config_path = "/mnt/Storage/home/huayuwei/container_workspace/spCS/2.result/0.pretrain_model/V5/epoch69/config.json"
+config = ViTMAEConfig.from_json_file(config_path)
+config.architectures = ["ViTForImageClassification"]
+config.num_labels = 3
+config.image_size = 112
+config.num_channels = 1
+model = ViTForImageClassification(config)
+model.to(DEVICE)
+
+# Training
+optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)
+writer = SummaryWriter(f'./tensorboard/{PORJECT_NAME}')
+step1 = 0
+step2 = 0
+best_val_loss = float('inf')
+best_val_f1 = 0
+warmup_epochs = 5
+scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: warmup_lr_lambda(epoch, warmup_epochs))
+
+for epoch in range(NUM_EPOCHS):
+    print(f"Epoch: {epoch+1}/{NUM_EPOCHS}")
+    model.train()
+    train_preds, train_labels = [], []
+    loss_list = []
+    for i, (x, l) in tqdm(enumerate(train_loader), total=len(train_loader)):
+        x = x.to(DEVICE)
+        l = l.to(DEVICE)
+        
+        print(f"Input shape: {x.shape}")
+        print(f"Label shape: {l.shape}")
+        
+        optimizer.zero_grad()
+        
+        outputs = model(x, labels=l)
+        
+        loss = outputs.loss
+        
+        _, predicted = torch.max(outputs.logits, 1)
+        train_preds.extend(predicted.cpu().numpy())
+        train_labels.extend(l.cpu().numpy())
+        
+        writer.add_scalar("Step/Train Loss", loss.item(), step1)
+        loss_list.append(loss.item())
+        
+        step1 += 1
+        loss.backward()
+        optimizer.step()
+
+    train_loss = np.mean(loss_list)
+    train_accuracy = 100 * (np.array(train_preds) == np.array(train_labels)).mean()
+    train_f1 = f1_score(train_labels, train_preds, average='macro')
+    train_precision = precision_score(train_labels, train_preds, average='macro')
+    
+    model.eval()
+    val_preds, val_labels = [], []
+    loss_list = []
+    with torch.no_grad():
+        for i, (x, l) in tqdm(enumerate(valid_loader), total=len(valid_loader)):
+            x = x.to(DEVICE)
+            l = l.to(DEVICE)
+
+            outputs = model(x, labels=l)
+
+            loss = outputs.loss
+            
+            _, predicted = torch.max(outputs.logits, 1)
+            val_preds.extend(predicted.cpu().numpy())
+            val_labels.extend(l.cpu().numpy())
+            
+            writer.add_scalar("Step/Validation Loss", loss.item(), step2)
+
+            loss_list.append(loss.item())
+            step2 += 1
+    val_loss = np.mean(loss_list)
+    val_accuracy = 100 * (np.array(val_preds) == np.array(val_labels)).mean()
+    val_f1 = f1_score(val_labels, val_preds, average='macro')
+    val_precision = precision_score(val_labels, val_preds, average='macro')
+    
+    val_labels_bin = label_binarize(val_labels, classes=[0, 1, 2])
+    val_preds_bin = label_binarize(val_preds, classes=[0, 1, 2])
+    val_auc = roc_auc_score(val_labels_bin, val_preds_bin, average='macro', multi_class='ovr')
+
+    # Save the model if the validation loss is the best we've seen so far.
+    if val_loss < best_val_loss:
+        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_loss_model.pt')
+        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_loss_model')
+        best_val_loss = val_loss
+
+    # Save the model if the validation F1 score is the best we've seen so far.
+    if val_f1 > best_val_f1:
+        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_f1_model.pt')
+        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_f1_model')
+        best_val_f1 = val_f1
+
+    lr = optimizer.param_groups[0]['lr']
+    writer.add_scalar("Epoch/Lr", lr, epoch)
+    writer.add_scalar("Epoch/Validation ROC AUC", val_auc, epoch)
+    writer.add_scalars("Epoch/Loss", {'Train Loss': train_loss, 'Validation Loss': val_loss}, epoch)
+    writer.add_scalars("Epoch/ACC", {'Train ACC': train_accuracy, 'Validation ACC': val_accuracy}, epoch)
+    writer.add_scalars("Epoch/Precision", {'Train Precision': train_precision, 'Validation Precision': val_precision}, epoch)
+    writer.add_scalars("Epoch/F1_Score", {'Train F1 Score': train_f1, 'Validation F1 Score': val_f1}, epoch)
+
+    print(f"Epoch {epoch}, Train Loss: {train_loss:.4f}, Train ACC: {train_accuracy:.4f}%, Train F1: {train_f1:.4f}, Train Precision: {train_precision:.4f}")
+    print(f"Epoch {epoch}, Validation Loss: {val_loss:.4f}, Validation ACC: {val_accuracy:.4f}%, Validation F1: {val_f1:.4f}, Validation Precision: {val_precision:.4f}, Validation ROC AUC: {val_auc:.4f}")
+
+    scheduler.step()
+
+# Test with best f1 model
+transform = transforms.Compose([
+    transforms.Resize((112, 112)), 
+    transforms.Grayscale(),
+    # transforms.RandomHorizontalFlip(p=0.5), 
+    # transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+dataset = ImageFolder(DATA_DIR, transform=transform)
+test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler)
+model_path = f'{folder_name}/{PORJECT_NAME}_best_f1_model.pt'
+model.load_state_dict(torch.load(model_path))
+model.to(DEVICE)
+model.eval()
+test_preds, test_labels = [], []
+test_probs = []
+
+with torch.no_grad():
+    for x, l in tqdm(test_loader, total=len(test_loader)):
+        x = x.to(DEVICE)
+        l = l.to(DEVICE)
+
+        outputs = model(x)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        _, predicted = torch.max(probabilities, 1)
+
+        test_preds.extend(predicted.cpu().numpy())
+        test_labels.extend(l.cpu().numpy())
+        test_probs.extend(probabilities.cpu().numpy())
+
+test_probs = np.array(test_probs)
+
+df = pd.DataFrame({
+    'True Labels': test_labels,
+    'Predicted Labels': test_preds
+})
+
+for i in range(test_probs.shape[1]):
+    df[f'Prob_Class{i}'] = test_probs[:, i]
+
+df.to_csv(f'{PORJECT_NAME}.csv', index=False)
+print("Test labels, predictions, and probabilities have been saved")
+
+test_labels_binarized = label_binarize(test_labels, classes=[0, 1, 2])
+test_preds_binarized = label_binarize(test_preds, classes=[0, 1, 2])
+
+accuracy = accuracy_score(test_labels, test_preds)
+f1 = f1_score(test_labels, test_preds, average='macro')
+precision = precision_score(test_labels, test_preds, average='macro')
+recall = recall_score(test_labels, test_preds, average='macro')
+rocauc = roc_auc_score(test_labels_binarized, test_preds_binarized, average='macro')
+
+print(f'Accuracy: {accuracy:.4f}')
+print(f'F1 Score: {f1:.4f}')
+print(f'Precision: {precision:.4f}')
+print(f'Recall: {recall:.4f}')
+print(f'ROC AUC: {rocauc:.4f}')
+
+conf_matrix = confusion_matrix(test_labels, test_preds)
+print("Confusion Matrix:")
+print(conf_matrix)

codes/Fine-tuning/cell_type_classification/NuSPIRe_full_fine-tuning.py

@@ -0,0 +1,284 @@
+import torch
+import random
+import numpy as np
+import os
+from torch.utils.tensorboard import SummaryWriter
+import pandas as pd
+from torchvision import transforms
+from torch.utils.data import DataLoader, SubsetRandomSampler
+from tqdm import tqdm
+from transformers import ViTForImageClassification
+from torchvision.datasets import ImageFolder
+from sklearn.metrics import accuracy_score, f1_score, precision_score, roc_auc_score, recall_score, confusion_matrix
+from sklearn.preprocessing import label_binarize
+from torch.optim.lr_scheduler import LambdaLR
+import argparse
+
+
+def set_seeds(seed_value=42, cuda_deterministic=False):
+    """Set seeds for reproducibility."""
+    random.seed(seed_value)
+    os.environ['PYTHONHASHSEED'] = str(seed_value)
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed_value)
+        torch.cuda.manual_seed_all(seed_value)
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    if cuda_deterministic:  # slower, more reproducible
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:  # faster, less reproducible
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+
+def warmup_lr_lambda(current_epoch: int, warmup_epochs: int):
+    if (current_epoch < warmup_epochs):
+        return float(current_epoch + 1) / float(max(1, warmup_epochs))
+    return 1.0       
+
+# set up
+parser = argparse.ArgumentParser(description="Setup experiment parameters")
+parser.add_argument('--num', type=int, required=True, help='Number of samples per class')
+parser.add_argument('--device', type=int, default=0, help='CUDA device number (default: 0)')
+parser.add_argument('--rep', type=int, required=True, help='Number of replicate')
+args = parser.parse_args()
+num_samples_per_class = args.num
+device = args.device
+num_repeats = args.rep
+    
+SEED = 42
+DEVICE = torch.device(f"cuda:{device}")
+DATA_DIR  = '../lung5_rep1_cancer_nuclear_image_15micron/'
+BATCH_SIZE = 300
+NUM_EPOCHS = 30
+PORJECT_NAME = f'Nuspire_{num_samples_per_class}_lung5_rep1_Classification'
+set_seeds(SEED)
+folder_name = f'./{PORJECT_NAME}_checkpoint'
+if not os.path.exists(folder_name):
+    os.makedirs(folder_name)
+    print(f"'{folder_name}'has been created.")
+else:
+    print(f"'{folder_name}' already exists.")
+
+# Dataset
+transform = transforms.Compose([
+    transforms.Resize((112, 112)),
+    transforms.Grayscale(),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+dataset = ImageFolder(DATA_DIR, transform=transform)
+labels = [dataset[i][1] for i in range(len(dataset))]
+
+# Define train and test sizes
+train_size = int(0.8 * len(dataset))
+valid_size = int(0.1 * len(dataset))
+test_size = len(dataset) - train_size - valid_size
+
+indices = np.arange(len(dataset))
+np.random.shuffle(indices)
+
+# Split
+train_indices = indices[:train_size]
+valid_indices = indices[train_size:train_size + valid_size]
+test_indices = indices[train_size + valid_size:]
+class_1_train_indices = [i for i in train_indices if labels[i] == 1]
+class_2_train_indices = [i for i in train_indices if labels[i] == 2]
+class_0_train_indices = [i for i in train_indices if labels[i] == 0]
+
+
+for repeat in range(num_repeats):
+    np.random.shuffle(class_1_train_indices)
+    np.random.shuffle(class_2_train_indices)
+    np.random.shuffle(class_0_train_indices)
+
+class_1_train_indices = class_1_train_indices[:num_samples_per_class]
+class_2_train_indices = class_2_train_indices[:num_samples_per_class]
+class_0_train_indices = class_0_train_indices[:num_samples_per_class]
+
+balanced_train_indices = (
+    class_1_train_indices +
+    class_2_train_indices +
+    class_0_train_indices
+)
+np.random.shuffle(balanced_train_indices)
+
+train_sampler = SubsetRandomSampler(balanced_train_indices)
+valid_sampler = SubsetRandomSampler(valid_indices)
+test_sampler = SubsetRandomSampler(test_indices)
+
+# print(balanced_train_indices)
+# print(valid_indices)
+# print(test_indices)
+
+train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers= 4)
+valid_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=valid_sampler, num_workers= 4)
+test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler, num_workers= 4)
+
+
+# Model
+model = ViTForImageClassification.from_pretrained("/mnt/Storage/home/huayuwei/container_workspace/spCS/2.result/0.pretrain_model/V5/epoch69",num_labels=3)
+model.to(DEVICE)
+
+# Training
+optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)
+writer = SummaryWriter(f'./tensorboard/{PORJECT_NAME}')
+step1 = 0
+step2 = 0
+best_val_loss = float('inf')
+best_val_f1 = 0
+warmup_epochs = 5
+scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: warmup_lr_lambda(epoch, warmup_epochs))
+
+for epoch in range(NUM_EPOCHS):
+    print(f"Epoch: {epoch+1}/{NUM_EPOCHS}")
+    model.train()
+    train_preds, train_labels = [], []
+    loss_list = []
+    for i, (x, l) in tqdm(enumerate(train_loader), total=len(train_loader)):
+        x = x.to(DEVICE)
+        l = l.to(DEVICE)
+        
+        optimizer.zero_grad()
+        
+        outputs = model(x, labels=l)
+        
+        loss = outputs.loss
+        
+        _, predicted = torch.max(outputs.logits, 1)
+        train_preds.extend(predicted.cpu().numpy())
+        train_labels.extend(l.cpu().numpy())
+        
+        writer.add_scalar("Step/Train Loss", loss.item(), step1)
+        loss_list.append(loss.item())
+        
+        step1 += 1
+        loss.backward()
+        optimizer.step()
+
+    train_loss = np.mean(loss_list)
+    train_accuracy = 100 * (np.array(train_preds) == np.array(train_labels)).mean()
+    train_f1 = f1_score(train_labels, train_preds, average='macro')
+    train_precision = precision_score(train_labels, train_preds, average='macro')
+    
+    model.eval()
+    val_preds, val_labels = [], []
+    loss_list = []
+    with torch.no_grad():
+        for i, (x, l) in tqdm(enumerate(valid_loader), total=len(valid_loader)):
+            x = x.to(DEVICE)
+            l = l.to(DEVICE)
+
+            outputs = model(x, labels=l)
+
+            loss = outputs.loss
+            
+            _, predicted = torch.max(outputs.logits, 1)
+            val_preds.extend(predicted.cpu().numpy())
+            val_labels.extend(l.cpu().numpy())
+            
+            writer.add_scalar("Step/Validation Loss", loss.item(), step2)
+
+            loss_list.append(loss.item())
+            step2 += 1
+    val_loss = np.mean(loss_list)
+    val_accuracy = 100 * (np.array(val_preds) == np.array(val_labels)).mean()
+    val_f1 = f1_score(val_labels, val_preds, average='macro')
+    val_precision = precision_score(val_labels, val_preds, average='macro')
+    
+    val_labels_bin = label_binarize(val_labels, classes=[0, 1, 2])
+    val_preds_bin = label_binarize(val_preds, classes=[0, 1, 2])
+    val_auc = roc_auc_score(val_labels_bin, val_preds_bin, average='macro', multi_class='ovr')
+
+    # Save the model if the validation loss is the best we've seen so far.
+    if val_loss < best_val_loss:
+        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_loss_model.pt')
+        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_loss_model')
+        best_val_loss = val_loss
+
+    # Save the model if the validation F1 score is the best we've seen so far.
+    if val_f1 > best_val_f1:
+        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_f1_model.pt')
+        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_f1_model')
+        best_val_f1 = val_f1
+
+    lr = optimizer.param_groups[0]['lr']
+    writer.add_scalar("Epoch/Lr", lr, epoch)
+    writer.add_scalar("Epoch/Validation ROC AUC", val_auc, epoch)
+    writer.add_scalars("Epoch/Loss", {'Train Loss': train_loss, 'Validation Loss': val_loss}, epoch)
+    writer.add_scalars("Epoch/ACC", {'Train ACC': train_accuracy, 'Validation ACC': val_accuracy}, epoch)
+    writer.add_scalars("Epoch/Precision", {'Train Precision': train_precision, 'Validation Precision': val_precision}, epoch)
+    writer.add_scalars("Epoch/F1_Score", {'Train F1 Score': train_f1, 'Validation F1 Score': val_f1}, epoch)
+
+    print(f"Epoch {epoch}, Train Loss: {train_loss:.4f}, Train ACC: {train_accuracy:.4f}%, Train F1: {train_f1:.4f}, Train Precision: {train_precision:.4f}")
+    print(f"Epoch {epoch}, Validation Loss: {val_loss:.4f}, Validation ACC: {val_accuracy:.4f}%, Validation F1: {val_f1:.4f}, Validation Precision: {val_precision:.4f}, Validation ROC AUC: {val_auc:.4f}")
+
+    scheduler.step()
+
+# Test with best f1 model
+transform = transforms.Compose([
+    transforms.Resize((112, 112)), 
+    transforms.Grayscale(),
+    # transforms.RandomHorizontalFlip(p=0.5), 
+    # transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+dataset = ImageFolder(DATA_DIR, transform=transform)
+test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler)
+model_path = f'{folder_name}/{PORJECT_NAME}_best_f1_model.pt'
+model.load_state_dict(torch.load(model_path))
+model.to(DEVICE) 
+model.eval()
+test_preds, test_labels = [], []
+test_probs = []
+
+with torch.no_grad():
+    for x, l in tqdm(test_loader, total=len(test_loader)):
+        x = x.to(DEVICE)
+        l = l.to(DEVICE)
+
+        outputs = model(x)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        _, predicted = torch.max(probabilities, 1)
+
+        test_preds.extend(predicted.cpu().numpy())
+        test_labels.extend(l.cpu().numpy())
+        test_probs.extend(probabilities.cpu().numpy())
+
+test_probs = np.array(test_probs)
+
+df = pd.DataFrame({
+    'True Labels': test_labels,
+    'Predicted Labels': test_preds
+})
+
+for i in range(test_probs.shape[1]):
+    df[f'Prob_Class{i}'] = test_probs[:, i]
+
+df.to_csv(f'{PORJECT_NAME}.csv', index=False)
+print("Test labels, predictions, and probabilities have been saved")
+
+test_labels_binarized = label_binarize(test_labels, classes=[0, 1, 2])
+test_preds_binarized = label_binarize(test_preds, classes=[0, 1, 2])
+
+accuracy = accuracy_score(test_labels, test_preds)
+f1 = f1_score(test_labels, test_preds, average='macro')
+precision = precision_score(test_labels, test_preds, average='macro')
+recall = recall_score(test_labels, test_preds, average='macro')
+rocauc = roc_auc_score(test_labels_binarized, test_preds_binarized, average='macro')
+
+print(f'Accuracy: {accuracy:.4f}')
+print(f'F1 Score: {f1:.4f}')
+print(f'Precision: {precision:.4f}')
+print(f'Recall: {recall:.4f}')
+print(f'ROC AUC: {rocauc:.4f}')
+
+conf_matrix = confusion_matrix(test_labels, test_preds)
+print("Confusion Matrix:")
+print(conf_matrix)

codes/Fine-tuning/cell_type_classification/NuSPIRe_partial_fine-tuning.py

@@ -0,0 +1,310 @@
+import torch
+import torch.nn as nn
+import random
+import numpy as np
+import os
+from torch.utils.tensorboard import SummaryWriter
+import pandas as pd
+from torchvision import transforms
+from torch.utils.data import DataLoader, SubsetRandomSampler
+from tqdm import tqdm
+from transformers import ViTForImageClassification
+from torchvision.datasets import ImageFolder
+from sklearn.metrics import accuracy_score, f1_score, precision_score, roc_auc_score, recall_score, confusion_matrix
+from sklearn.preprocessing import label_binarize
+from torch.optim.lr_scheduler import LambdaLR
+import argparse
+
+
+def set_seeds(seed_value=42, cuda_deterministic=False):
+    """Set seeds for reproducibility."""
+    random.seed(seed_value)
+    os.environ['PYTHONHASHSEED'] = str(seed_value)
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed_value)
+        torch.cuda.manual_seed_all(seed_value)
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    if cuda_deterministic:  # slower, more reproducible
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:  # faster, less reproducible
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+        
+
+def warmup_lr_lambda(current_epoch: int, warmup_epochs: int):
+    if (current_epoch < warmup_epochs):
+        return float(current_epoch + 1) / float(max(1, warmup_epochs))
+    return 1.0
+
+# set up
+parser = argparse.ArgumentParser(description="Setup experiment parameters")
+parser.add_argument('--num', type=int, required=True, help='Number of samples per class')
+parser.add_argument('--device', type=int, default=0, help='CUDA device number (default: 0)')
+parser.add_argument('--rep', type=int, required=True, help='Number of replicate')
+args = parser.parse_args()
+num_samples_per_class = args.num
+device = args.device
+num_repeats = args.rep
+    
+SEED = 42
+DEVICE = torch.device(f"cuda:{device}")
+DATA_DIR  = '../lung5_rep1_cancer_nuclear_image_15micron/'
+BATCH_SIZE = 300
+NUM_EPOCHS = 30
+PORJECT_NAME = f'MLP_Frozen_{num_samples_per_class}_lung5_rep1_Classification'
+set_seeds(SEED)
+folder_name = f'./{PORJECT_NAME}_checkpoint'
+
+if not os.path.exists(folder_name):
+    os.makedirs(folder_name)
+    print(f"'{folder_name}'has been created.")
+else:
+    print(f"'{folder_name}' already exists.")
+
+# Dataset
+transform = transforms.Compose([
+    transforms.Resize((112, 112)),
+    transforms.Grayscale(),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+dataset = ImageFolder(DATA_DIR, transform=transform)
+labels = [dataset[i][1] for i in range(len(dataset))]
+
+# Define train and test sizes
+train_size = int(0.8 * len(dataset))
+valid_size = int(0.1 * len(dataset))
+test_size = len(dataset) - train_size - valid_size
+
+indices = np.arange(len(dataset))
+np.random.shuffle(indices)
+
+# Split
+train_indices = indices[:train_size]
+valid_indices = indices[train_size:train_size + valid_size]
+test_indices = indices[train_size + valid_size:]
+class_1_train_indices = [i for i in train_indices if labels[i] == 1]
+class_2_train_indices = [i for i in train_indices if labels[i] == 2]
+class_0_train_indices = [i for i in train_indices if labels[i] == 0]
+
+for repeat in range(num_repeats):
+    np.random.shuffle(class_1_train_indices)
+    np.random.shuffle(class_2_train_indices)
+    np.random.shuffle(class_0_train_indices)
+
+class_1_train_indices = class_1_train_indices[:num_samples_per_class]
+class_2_train_indices = class_2_train_indices[:num_samples_per_class]
+class_0_train_indices = class_0_train_indices[:num_samples_per_class]
+
+balanced_train_indices = (
+    class_1_train_indices +
+    class_2_train_indices +
+    class_0_train_indices
+)
+np.random.shuffle(balanced_train_indices)
+
+train_sampler = SubsetRandomSampler(balanced_train_indices)
+valid_sampler = SubsetRandomSampler(valid_indices)
+test_sampler = SubsetRandomSampler(test_indices)
+
+# print(balanced_train_indices)
+# print(valid_indices)
+# print(test_indices)
+
+train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers= 4)
+valid_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=valid_sampler, num_workers= 4)
+test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler, num_workers= 4)
+
+
+# Model
+model = ViTForImageClassification.from_pretrained("/mnt/Storage/home/huayuwei/container_workspace/spCS/2.result/0.pretrain_model/V5/epoch69",num_labels=3)
+for name, param in model.named_parameters():
+    if 'classifier' not in name:
+        param.requires_grad = False
+
+class MLP(nn.Module):
+    def __init__(self, input_dim, hidden_dim1, hidden_dim2, hidden_dim3, hidden_dim4, output_dim):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim1)
+        self.fc2 = nn.Linear(hidden_dim1, hidden_dim2)
+        self.fc3 = nn.Linear(hidden_dim2, hidden_dim3)
+        self.fc4 = nn.Linear(hidden_dim3, hidden_dim4)
+        self.fc5 = nn.Linear(hidden_dim4, output_dim)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.fc2(x))
+        x = self.relu(self.fc3(x))
+        x = self.relu(self.fc4(x))
+        x = self.fc5(x)
+        return x
+
+model.classifier = MLP(input_dim=768, hidden_dim1=512, hidden_dim2=256, hidden_dim3=128, hidden_dim4=64, output_dim=3)
+model.to(DEVICE)
+
+# Training
+optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)
+writer = SummaryWriter(f'./tensorboard/{PORJECT_NAME}')
+step1 = 0
+step2 = 0
+best_val_loss = float('inf')
+best_val_f1 = 0
+warmup_epochs = 5
+scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: warmup_lr_lambda(epoch, warmup_epochs))
+
+
+for epoch in range(NUM_EPOCHS):
+    print(f"Epoch: {epoch+1}/{NUM_EPOCHS}")
+    model.train()
+    train_preds, train_labels = [], []
+    loss_list = []
+    for i, (x, l) in tqdm(enumerate(train_loader), total=len(train_loader)):
+        x = x.to(DEVICE)
+        l = l.to(DEVICE)
+        
+        optimizer.zero_grad()
+        
+        outputs = model(x, labels=l)
+        
+        loss = outputs.loss
+        
+        _, predicted = torch.max(outputs.logits, 1)
+        train_preds.extend(predicted.cpu().numpy())
+        train_labels.extend(l.cpu().numpy())
+        
+        writer.add_scalar("Step/Train Loss", loss.item(), step1)
+        loss_list.append(loss.item())
+        
+        step1 += 1
+        loss.backward()
+        optimizer.step()
+
+    train_loss = np.mean(loss_list)
+    train_accuracy = 100 * (np.array(train_preds) == np.array(train_labels)).mean()
+    train_f1 = f1_score(train_labels, train_preds, average='macro')
+    train_precision = precision_score(train_labels, train_preds, average='macro')
+    
+    model.eval()
+    val_preds, val_labels = [], []
+    loss_list = []
+    with torch.no_grad():
+        for i, (x, l) in tqdm(enumerate(valid_loader), total=len(valid_loader)):
+            x = x.to(DEVICE)
+            l = l.to(DEVICE)
+
+            outputs = model(x, labels=l)
+
+            loss = outputs.loss
+            
+            _, predicted = torch.max(outputs.logits, 1)
+            val_preds.extend(predicted.cpu().numpy())
+            val_labels.extend(l.cpu().numpy())
+            
+            writer.add_scalar("Step/Validation Loss", loss.item(), step2)
+
+            loss_list.append(loss.item())
+            step2 += 1
+    val_loss = np.mean(loss_list)
+    val_accuracy = 100 * (np.array(val_preds) == np.array(val_labels)).mean()
+    val_f1 = f1_score(val_labels, val_preds, average='macro')
+    val_precision = precision_score(val_labels, val_preds, average='macro')
+    
+    val_labels_bin = label_binarize(val_labels, classes=[0, 1, 2])
+    val_preds_bin = label_binarize(val_preds, classes=[0, 1, 2])
+    val_auc = roc_auc_score(val_labels_bin, val_preds_bin, average='macro', multi_class='ovr')
+
+    # Save the model if the validation loss is the best we've seen so far.
+    if val_loss < best_val_loss:
+        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_loss_model.pt')
+        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_loss_model')
+        best_val_loss = val_loss
+
+    # Save the model if the validation F1 score is the best we've seen so far.
+    if val_f1 > best_val_f1:
+        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_f1_model.pt')
+        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_f1_model')
+        best_val_f1 = val_f1
+
+    lr = optimizer.param_groups[0]['lr']
+    writer.add_scalar("Epoch/Lr", lr, epoch)
+    writer.add_scalar("Epoch/Validation ROC AUC", val_auc, epoch)
+    writer.add_scalars("Epoch/Loss", {'Train Loss': train_loss, 'Validation Loss': val_loss}, epoch)
+    writer.add_scalars("Epoch/ACC", {'Train ACC': train_accuracy, 'Validation ACC': val_accuracy}, epoch)
+    writer.add_scalars("Epoch/Precision", {'Train Precision': train_precision, 'Validation Precision': val_precision}, epoch)
+    writer.add_scalars("Epoch/F1_Score", {'Train F1 Score': train_f1, 'Validation F1 Score': val_f1}, epoch)
+
+    print(f"Epoch {epoch}, Train Loss: {train_loss:.4f}, Train ACC: {train_accuracy:.4f}%, Train F1: {train_f1:.4f}, Train Precision: {train_precision:.4f}")
+    print(f"Epoch {epoch}, Validation Loss: {val_loss:.4f}, Validation ACC: {val_accuracy:.4f}%, Validation F1: {val_f1:.4f}, Validation Precision: {val_precision:.4f}, Validation ROC AUC: {val_auc:.4f}")
+
+    scheduler.step()
+
+# Test with best f1 model
+transform = transforms.Compose([
+    transforms.Resize((112, 112)),
+    transforms.Grayscale(),
+    # transforms.RandomHorizontalFlip(p=0.5),
+    # transforms.RandomVerticalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])
+])
+
+dataset = ImageFolder(DATA_DIR, transform=transform)
+test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler)
+model_path = f'{folder_name}/{PORJECT_NAME}_best_f1_model.pt'
+model.load_state_dict(torch.load(model_path))
+model.to(DEVICE)
+model.eval()
+test_preds, test_labels = [], []
+test_probs = []
+
+with torch.no_grad():
+    for x, l in tqdm(test_loader, total=len(test_loader)):
+        x = x.to(DEVICE)
+        l = l.to(DEVICE)
+
+        outputs = model(x)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        _, predicted = torch.max(probabilities, 1)
+
+        test_preds.extend(predicted.cpu().numpy())
+        test_labels.extend(l.cpu().numpy())
+        test_probs.extend(probabilities.cpu().numpy())
+
+test_probs = np.array(test_probs)
+
+df = pd.DataFrame({
+    'True Labels': test_labels,
+    'Predicted Labels': test_preds
+})
+
+for i in range(test_probs.shape[1]):
+    df[f'Prob_Class{i}'] = test_probs[:, i]
+
+df.to_csv(f'{PORJECT_NAME}.csv', index=False)
+print("Test labels, predictions, and probabilities have been saved")
+
+test_labels_binarized = label_binarize(test_labels, classes=[0, 1, 2])
+test_preds_binarized = label_binarize(test_preds, classes=[0, 1, 2])
+
+accuracy = accuracy_score(test_labels, test_preds)
+f1 = f1_score(test_labels, test_preds, average='macro')
+precision = precision_score(test_labels, test_preds, average='macro')
+recall = recall_score(test_labels, test_preds, average='macro')
+rocauc = roc_auc_score(test_labels_binarized, test_preds_binarized, average='macro')
+
+print(f'Accuracy: {accuracy:.4f}')
+print(f'F1 Score: {f1:.4f}')
+print(f'Precision: {precision:.4f}')
+print(f'Recall: {recall:.4f}')
+print(f'ROC AUC: {rocauc:.4f}')
+
+conf_matrix = confusion_matrix(test_labels, test_preds)
+print("Confusion Matrix:")
+print(conf_matrix)

codes/Fine-tuning/expression_prediction/NuSPIRe_full_fine-tuning.ipynb

@@ -0,0 +1,1105 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b0410ca4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import random\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import pandas as pd\n",
+    "from torchvision import transforms\n",
+    "from PIL import Image\n",
+    "from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler\n",
+    "from tqdm import tqdm\n",
+    "from transformers import ViTForImageClassification, ViTConfig\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2f99710",
+   "metadata": {},
+   "source": [
+    "# hyperparameter"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "b1a22094",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SEED = 42\n",
+    "DEVICE = torch.device(\"cuda:0\")\n",
+    "DATA_DIR  = '../train_nucleus_128_with_env_15dis_cell_scale/all/'\n",
+    "BATCH_SIZE = 300\n",
+    "NUM_EPOCHS = 30\n",
+    "PORJECT_NAME = f'Nuspire_mouse_brain_Regression'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "924045aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def set_seeds(seed_value=42, cuda_deterministic=False):\n",
+    "    \"\"\"Set seeds for reproducibility.\"\"\"\n",
+    "    random.seed(seed_value)\n",
+    "    os.environ['PYTHONHASHSEED'] = str(seed_value)\n",
+    "    np.random.seed(seed_value)\n",
+    "    torch.manual_seed(seed_value)\n",
+    "    if torch.cuda.is_available():\n",
+    "        torch.cuda.manual_seed(seed_value)\n",
+    "        torch.cuda.manual_seed_all(seed_value)\n",
+    "    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html\n",
+    "    if cuda_deterministic:  # slower, more reproducible\n",
+    "        torch.backends.cudnn.deterministic = True\n",
+    "        torch.backends.cudnn.benchmark = False\n",
+    "    else:  # faster, less reproducible\n",
+    "        torch.backends.cudnn.deterministic = False\n",
+    "        torch.backends.cudnn.benchmark = True\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9caab5e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "set_seeds(SEED)\n",
+    "timestamp = \"07\"\n",
+    "folder_name = f'./{PORJECT_NAME}_{timestamp}_checkpoint'\n",
+    "\n",
+    "if not os.path.exists(folder_name):\n",
+    "    os.makedirs(folder_name)\n",
+    "    # print(f\"'{folder_name}'has been created.\")\n",
+    "else:\n",
+    "    print(f\"'{folder_name}' already exists.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "4a354850",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ImageDataset(Dataset):\n",
+    "    def __init__(self, data_dir, transform=None):\n",
+    "        self.data_dir = data_dir\n",
+    "        self.transform = transform\n",
+    "        self.file_list = os.listdir(data_dir)\n",
+    "        self.cell_expression = pd.read_csv('../processed_data/cell_expression_filtered_size_allgene.csv', index_col=0)\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.file_list)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        img_name = os.path.join(self.data_dir, self.file_list[idx])\n",
+    "        img_index = img_name.split(\"/\")[-1].replace('image_', '').replace('.png', '')\n",
+    "        image = Image.open(img_name).convert('L')\n",
+    "        if self.transform:\n",
+    "            image = self.transform(image)\n",
+    "        \n",
+    "        if img_index in self.cell_expression.index:\n",
+    "            target = self.cell_expression.loc[img_index].values\n",
+    "        else:\n",
+    "            target = None\n",
+    "        return image, target"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "96d9003d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = transforms.Compose([\n",
+    "    transforms.Resize((112, 112)),\n",
+    "    transforms.RandomHorizontalFlip(p=0.5),\n",
+    "    transforms.RandomVerticalFlip(p=0.5),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "772d8f7b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = ImageDataset(DATA_DIR, transform=transform)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c52f7512",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "total_size = len(dataset)\n",
+    "train_size = int(total_size * 0.8)\n",
+    "remaining_size = total_size - train_size\n",
+    "\n",
+    "valid_size = int(remaining_size * 0.5)\n",
+    "test_size = remaining_size - valid_size\n",
+    "\n",
+    "indices = list(range(total_size))\n",
+    "np.random.shuffle(indices)\n",
+    "\n",
+    "train_indices = indices[:train_size]\n",
+    "remaining_indices = indices[train_size:]\n",
+    "valid_indices = remaining_indices[:valid_size]\n",
+    "test_indices = remaining_indices[valid_size:]\n",
+    "\n",
+    "train_sampler = SubsetRandomSampler(train_indices)\n",
+    "valid_sampler = SubsetRandomSampler(valid_indices)\n",
+    "test_sampler = SubsetRandomSampler(test_indices)\n",
+    "\n",
+    "train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers=4)\n",
+    "valid_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=valid_sampler, num_workers=4)\n",
+    "test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler, num_workers=4)\n",
+    "\n",
+    "# print(train_indices)\n",
+    "# print(valid_indices)\n",
+    "# print(test_indices)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c6e1da23",
+   "metadata": {},
+   "source": [
+    "# model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "224f2cab",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "You are using a model of type vit_mae to instantiate a model of type vit. This is not supported for all configurations of models and can yield errors.\n",
+      "Some weights of ViTForImageClassification were not initialized from the model checkpoint at /mnt/Storage/home/huayuwei/container_workspace/spCS/2.result/0.pretrain_model/V5/epoch69 and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
+      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = ViTConfig.from_pretrained(\"/mnt/Storage/home/huayuwei/container_workspace/spCS/2.result/0.pretrain_model/V5/epoch69\")\n",
+    "\n",
+    "config.hidden_dropout_prob = 0\n",
+    "config.attention_probs_dropout_prob = 0\n",
+    "config.num_labels = 347\n",
+    "\n",
+    "model = ViTForImageClassification.from_pretrained(\n",
+    "    \"/mnt/Storage/home/huayuwei/container_workspace/spCS/2.result/0.pretrain_model/V5/epoch69\",\n",
+    "    config=config\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "3704f75c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ViTForImageClassification(\n",
+       "  (vit): ViTModel(\n",
+       "    (embeddings): ViTEmbeddings(\n",
+       "      (patch_embeddings): ViTPatchEmbeddings(\n",
+       "        (projection): Conv2d(1, 768, kernel_size=(8, 8), stride=(8, 8))\n",
+       "      )\n",
+       "      (dropout): Dropout(p=0, inplace=False)\n",
+       "    )\n",
+       "    (encoder): ViTEncoder(\n",
+       "      (layer): ModuleList(\n",
+       "        (0-11): 12 x ViTLayer(\n",
+       "          (attention): ViTAttention(\n",
+       "            (attention): ViTSelfAttention(\n",
+       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (dropout): Dropout(p=0, inplace=False)\n",
+       "            )\n",
+       "            (output): ViTSelfOutput(\n",
+       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (dropout): Dropout(p=0, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "          (intermediate): ViTIntermediate(\n",
+       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
+       "            (intermediate_act_fn): GELUActivation()\n",
+       "          )\n",
+       "          (output): ViTOutput(\n",
+       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
+       "            (dropout): Dropout(p=0, inplace=False)\n",
+       "          )\n",
+       "          (layernorm_before): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "          (layernorm_after): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (layernorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "  )\n",
+       "  (classifier): Linear(in_features=768, out_features=347, bias=True)\n",
+       ")"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.to(DEVICE)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ea686b1b",
+   "metadata": {},
+   "source": [
+    "# Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "d9c8456e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)\n",
+    "criterion = nn.MSELoss()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "18c5aee0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch: 1/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [04:57<00:00,  2.08s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.37it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 0, Train Loss: 0.1923, Validation Loss: 0.1672\n",
+      "Epoch: 2/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [04:59<00:00,  2.10s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.38it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1, Train Loss: 0.1617, Validation Loss: 0.1588\n",
+      "Epoch: 3/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [04:59<00:00,  2.10s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.37it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2, Train Loss: 0.1528, Validation Loss: 0.1526\n",
+      "Epoch: 4/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:06<00:00,  2.14s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.33it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3, Train Loss: 0.1480, Validation Loss: 0.1482\n",
+      "Epoch: 5/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:07<00:00,  2.15s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.34it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4, Train Loss: 0.1445, Validation Loss: 0.1473\n",
+      "Epoch: 6/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:02<00:00,  2.12s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.37it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5, Train Loss: 0.1421, Validation Loss: 0.1455\n",
+      "Epoch: 7/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [07:31<00:00,  3.16s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.38it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 6, Train Loss: 0.1400, Validation Loss: 0.1428\n",
+      "Epoch: 8/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:02<00:00,  2.12s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.38it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 7, Train Loss: 0.1380, Validation Loss: 0.1426\n",
+      "Epoch: 9/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [04:55<00:00,  2.06s/it]\n",
+      "100%|██████████| 18/18 [00:12<00:00,  1.41it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 8, Train Loss: 0.1366, Validation Loss: 0.1429\n",
+      "Epoch: 10/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [04:53<00:00,  2.05s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.38it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 9, Train Loss: 0.1352, Validation Loss: 0.1408\n",
+      "Epoch: 11/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [04:57<00:00,  2.08s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.33it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 10, Train Loss: 0.1340, Validation Loss: 0.1419\n",
+      "Epoch: 12/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|█��████████| 143/143 [04:59<00:00,  2.09s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.38it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 11, Train Loss: 0.1325, Validation Loss: 0.1418\n",
+      "Epoch: 13/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [04:59<00:00,  2.09s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.38it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 12, Train Loss: 0.1314, Validation Loss: 0.1403\n",
+      "Epoch: 14/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:00<00:00,  2.10s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.37it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 13, Train Loss: 0.1304, Validation Loss: 0.1395\n",
+      "Epoch: 15/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:00<00:00,  2.10s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.37it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 14, Train Loss: 0.1289, Validation Loss: 0.1393\n",
+      "Epoch: 16/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [05:31<00:00,  2.32s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.17it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 15, Train Loss: 0.1279, Validation Loss: 0.1398\n",
+      "Epoch: 17/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:02<00:00,  2.53s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.16it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 16, Train Loss: 0.1265, Validation Loss: 0.1394\n",
+      "Epoch: 18/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:01<00:00,  2.53s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.15it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 17, Train Loss: 0.1254, Validation Loss: 0.1403\n",
+      "Epoch: 19/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:03<00:00,  2.54s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.17it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 18, Train Loss: 0.1241, Validation Loss: 0.1389\n",
+      "Epoch: 20/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 143/143 [06:07<00:00,  2.57s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.13it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 19, Train Loss: 0.1234, Validation Loss: 0.1399\n",
+      "Epoch: 21/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:13<00:00,  2.61s/it]\n",
+      "100%|██████████| 18/18 [00:16<00:00,  1.10it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 20, Train Loss: 0.1225, Validation Loss: 0.1407\n",
+      "Epoch: 22/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:13<00:00,  2.62s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.14it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 21, Train Loss: 0.1213, Validation Loss: 0.1395\n",
+      "Epoch: 23/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:15<00:00,  2.62s/it]\n",
+      "100%|██████████| 18/18 [00:16<00:00,  1.12it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 22, Train Loss: 0.1203, Validation Loss: 0.1399\n",
+      "Epoch: 24/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:14<00:00,  2.62s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.14it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 23, Train Loss: 0.1190, Validation Loss: 0.1400\n",
+      "Epoch: 25/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:13<00:00,  2.61s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.14it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 24, Train Loss: 0.1182, Validation Loss: 0.1413\n",
+      "Epoch: 26/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [05:39<00:00,  2.37s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.34it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 25, Train Loss: 0.1173, Validation Loss: 0.1407\n",
+      "Epoch: 27/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [05:08<00:00,  2.16s/it]\n",
+      "100%|██████████| 18/18 [00:13<00:00,  1.37it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 26, Train Loss: 0.1163, Validation Loss: 0.1411\n",
+      "Epoch: 28/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [10:22<00:00,  4.35s/it]\n",
+      "100%|██████████| 18/18 [00:32<00:00,  1.78s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 27, Train Loss: 0.1156, Validation Loss: 0.1412\n",
+      "Epoch: 29/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [09:45<00:00,  4.09s/it]\n",
+      "100%|██████████| 18/18 [00:16<00:00,  1.09it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 28, Train Loss: 0.1144, Validation Loss: 0.1416\n",
+      "Epoch: 30/30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 143/143 [06:34<00:00,  2.76s/it]\n",
+      "100%|██████████| 18/18 [00:15<00:00,  1.14it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 29, Train Loss: 0.1134, Validation Loss: 0.1412\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "writer = SummaryWriter(f'./tensorboard/{PORJECT_NAME}_{timestamp}')\n",
+    "step1 = 0\n",
+    "step2 = 0\n",
+    "best_val_loss = 1\n",
+    "\n",
+    "for epoch in range(NUM_EPOCHS):\n",
+    "    print(f\"Epoch: {epoch+1}/{NUM_EPOCHS}\")\n",
+    "    model.train()\n",
+    "    loss_list = []\n",
+    "    for i, (x,l) in tqdm(enumerate(train_loader), total=len(train_loader)):\n",
+    "        x = x.to(DEVICE)\n",
+    "        l = l.to(DEVICE)\n",
+    "        \n",
+    "        optimizer.zero_grad()\n",
+    "        \n",
+    "        outputs = model(x)\n",
+    "        \n",
+    "        loss = criterion(outputs.logits, l.float())\n",
+    "        \n",
+    "        writer.add_scalar(\"Step/Train Loss\", loss.item(),step1)\n",
+    "        loss_list.append(loss.item())\n",
+    "        \n",
+    "        step1+=1\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "    train_loss = np.mean(loss_list)\n",
+    "\n",
+    "    model.eval()\n",
+    "    loss_list = []\n",
+    "    with torch.no_grad():\n",
+    "         for i, (x,l) in tqdm(enumerate(valid_loader), total=len(valid_loader)):\n",
+    "            x = x.to(DEVICE)\n",
+    "            l = l.to(DEVICE)\n",
+    "\n",
+    "            optimizer.zero_grad()\n",
+    "\n",
+    "            outputs = model(x)\n",
+    "\n",
+    "            loss = criterion(outputs.logits, l.float())\n",
+    "            \n",
+    "            writer.add_scalar(\"Step/Validation Loss\", loss.item(),step2)\n",
+    "\n",
+    "            loss_list.append(loss.item())\n",
+    "            step2+=1\n",
+    "    val_loss = np.mean(loss_list)\n",
+    "    \n",
+    "    # Save the model if the validation loss is the best we've seen so far.\n",
+    "    if val_loss < best_val_loss:\n",
+    "        torch.save(model.state_dict(), f'{folder_name}/{PORJECT_NAME}_best_model.pt')\n",
+    "        model.save_pretrained(f'{folder_name}/{PORJECT_NAME}_best_model')\n",
+    "        best_epoch=epoch\n",
+    "        best_val_loss = val_loss\n",
+    "\n",
+    "    lr = optimizer.param_groups[0]['lr']\n",
+    "    writer.add_scalar(\"Epoch/Lr\", lr, epoch)\n",
+    "    writer.add_scalars(\"Epoch/Loss\",{'Train Loss':train_loss,'Validation Loss':val_loss},epoch)\n",
+    "    print(f\"Epoch {epoch}, Train Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0af358d9",
+   "metadata": {},
+   "source": [
+    "# Test"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "c7a6c5db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = transforms.Compose([\n",
+    "    transforms.Resize((112, 112)),\n",
+    "    # transforms.RandomHorizontalFlip(p=0.5),\n",
+    "    # transforms.RandomVerticalFlip(p=0.5),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(mean=[0.21869252622127533], std=[0.1809280514717102])\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "aa667594",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = ImageDataset(DATA_DIR, transform=transform)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "85b09b94",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=test_sampler, num_workers=4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "2e43c973",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ViTForImageClassification(\n",
+       "  (vit): ViTModel(\n",
+       "    (embeddings): ViTEmbeddings(\n",
+       "      (patch_embeddings): ViTPatchEmbeddings(\n",
+       "        (projection): Conv2d(1, 768, kernel_size=(8, 8), stride=(8, 8))\n",
+       "      )\n",
+       "      (dropout): Dropout(p=0, inplace=False)\n",
+       "    )\n",
+       "    (encoder): ViTEncoder(\n",
+       "      (layer): ModuleList(\n",
+       "        (0-11): 12 x ViTLayer(\n",
+       "          (attention): ViTAttention(\n",
+       "            (attention): ViTSelfAttention(\n",
+       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (dropout): Dropout(p=0, inplace=False)\n",
+       "            )\n",
+       "            (output): ViTSelfOutput(\n",
+       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (dropout): Dropout(p=0, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "          (intermediate): ViTIntermediate(\n",
+       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
+       "            (intermediate_act_fn): GELUActivation()\n",
+       "          )\n",
+       "          (output): ViTOutput(\n",
+       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
+       "            (dropout): Dropout(p=0, inplace=False)\n",
+       "          )\n",
+       "          (layernorm_before): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "          (layernorm_after): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (layernorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "  )\n",
+       "  (classifier): Linear(in_features=768, out_features=347, bias=True)\n",
+       ")"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_path = f'{folder_name}/{PORJECT_NAME}_best_model.pt'\n",
+    "model.load_state_dict(torch.load(model_path))\n",
+    "model.to(DEVICE) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "08a01b05",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 18/18 [00:15<00:00,  1.14it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.eval()\n",
+    "true_labels = []\n",
+    "predicted_outputs = []\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    for i, (x, l) in tqdm(enumerate(test_loader), total=len(test_loader)):\n",
+    "        x = x.to(DEVICE)\n",
+    "        l = l.to(DEVICE)\n",
+    "\n",
+    "        outputs = model(x)\n",
+    "\n",
+    "        # Collect true labels and predicted outputs\n",
+    "        true_labels.append(l.cpu())\n",
+    "        predicted_outputs.append(outputs.logits.cpu())\n",
+    "        \n",
+    "    true_labels = torch.cat(true_labels).numpy()\n",
+    "    predicted_outputs = torch.cat(predicted_outputs).numpy()  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "10fdc2f5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.save(f'{PORJECT_NAME}_{timestamp}_all_outputs.npy', predicted_outputs)\n",
+    "np.save(f'{PORJECT_NAME}_{timestamp}_all_targets.npy', true_labels)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "ee0c734a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MSE: [0.13185952 0.09007648 0.16135108 ... 0.12363786 0.18178999 0.25879715]\n",
+      "Pearson: [0.77104155 0.82377504 0.70091003 ... 0.77836889 0.64078275 0.49941442]\n",
+      "MSE - Mean: 0.1386, Std: 0.0514\n",
+      "Pearson - Mean: 0.7380, Std: 0.1021\n"
+     ]
+    }
+   ],
+   "source": [
+    "from scipy.stats import pearsonr\n",
+    "from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, mean_absolute_percentage_error, explained_variance_score\n",
+    "\n",
+    "n_samples, n_features = true_labels.shape\n",
+    "\n",
+    "results = {metric: [] for metric in ['MSE',\n",
+    "                                    #  'RMSE',\n",
+    "                                    #  'MAE', \n",
+    "                                    #  'MAPE', \n",
+    "                                    #  'R_squared', \n",
+    "                                    #  'Explained_Variance',\n",
+    "                                     'Pearson']}\n",
+    "\n",
+    "for i in range(n_samples):\n",
+    "    mse = mean_squared_error(true_labels[i, :], predicted_outputs[i, :])\n",
+    "    # rmse = np.sqrt(mse)\n",
+    "    # mae = mean_absolute_error(true_labels[i, :], predicted_outputs[i, :])\n",
+    "    # mape = mean_absolute_percentage_error(true_labels[i, :], predicted_outputs[i, :])\n",
+    "    # r2 = r2_score(true_labels[i, :], predicted_outputs[i, :])\n",
+    "    # explained_var = explained_variance_score(true_labels[i, :], predicted_outputs[i, :])\n",
+    "    pcc, _ = pearsonr(true_labels[i, :], predicted_outputs[i, :])\n",
+    "\n",
+    "    results['MSE'].append(mse)\n",
+    "    # results['RMSE'].append(rmse)\n",
+    "    # results['MAE'].append(mae)\n",
+    "    # results['MAPE'].append(mape)\n",
+    "    # results['R_squared'].append(r2)\n",
+    "    # results['Explained_Variance'].append(explained_var)\n",
+    "    results['Pearson'].append(pcc)\n",
+    "\n",
+    "for metric in results:\n",
+    "    results[metric] = np.array(results[metric])\n",
+    "\n",
+    "for metric in results:\n",
+    "    print(f\"{metric}: {results[metric]}\")\n",
+    "\n",
+    "for metric in results:\n",
+    "    print(f\"{metric} - Mean: {np.mean(results[metric]):.4f}, Std: {np.std(results[metric]):.4f}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}