204 lines
7.8 KiB
Python
204 lines
7.8 KiB
Python
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from torchvision import models, transforms
|
|
from PIL import Image
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import random
|
|
import time
|
|
import psutil
|
|
|
|
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(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
|
|
set_seed(42)
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
IMG_DIR = os.path.join(BASE_DIR, "..", "datos_para_ai")
|
|
MODEL_SAVE_PATH = os.path.join(BASE_DIR, "..", "models", "msi_encoder_trained.pth")
|
|
|
|
# VRAM OPTIMIZATION ENGINE
|
|
BATCH_SIZE = 8 # Physical batch size (keeps VRAM safely under 3.5 GB)
|
|
ACCUMULATION_STEPS = 4 # 8 x 4 = Effective Batch Size of 32!
|
|
EPOCHS = 60 # Slightly extended to give the larger batch size room to converge
|
|
LEARNING_RATE = 0.0003 # Tuned for an effective batch size of 32
|
|
TEMPERATURE = 0.5
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
class MSIAugmentations:
|
|
def __init__(self, size=256): # Enforce 256 dimension consistency
|
|
self.spatial_transform = transforms.Compose([
|
|
transforms.Resize((size, size)),
|
|
transforms.ToTensor()
|
|
])
|
|
|
|
def apply_msi_chemical_noise(self, tensor):
|
|
x = tensor.clone()
|
|
|
|
# 1. Intensity Scaling (Global Concentration Shifts)
|
|
scale_factor = random.uniform(0.8, 1.2)
|
|
x = x * scale_factor
|
|
|
|
# 2. Poisson Noise Simulation (Shot Noise)
|
|
# Scaled to avoid breaking the 0.0 - 1.0 float bounds
|
|
vals = len(torch.unique(x))
|
|
vals = 2 ** np.ceil(np.log2(vals)) if vals > 0 else 32
|
|
# Add soft, random Gaussian-Poisson approximation
|
|
noise = torch.randn_like(x) * (torch.sqrt(x + 1e-5) / float(vals))
|
|
x = x + noise
|
|
|
|
# 3. Random Missing Values (Standard Dropout - e.g., 2% pixels lost)
|
|
random_dropout_mask = (torch.rand_like(x) > 0.02).float()
|
|
x = x * random_dropout_mask
|
|
|
|
# 4. Intensity-Dependent Missing Values (Low-Signal Clipping)
|
|
# Pixels with very low values have a higher chance of dropping out completely
|
|
low_signal_mask = torch.rand_like(x)
|
|
# Drop threshold scales inversely with intensity: lower values drop more easily
|
|
drop_threshold = 0.15 * (1.0 - x)
|
|
x = torch.where(low_signal_mask > drop_threshold, x, torch.zeros_like(x))
|
|
|
|
return torch.clamp(x, 0.0, 1.0)
|
|
|
|
def __call__(self, x):
|
|
base_tensor = self.spatial_transform(x)
|
|
# Create two distinct chemical representations of the same underlying coordinate geometry
|
|
return self.apply_msi_chemical_noise(base_tensor), self.apply_msi_chemical_noise(base_tensor)
|
|
|
|
class SimCLRMSIDataset(Dataset):
|
|
def __init__(self, img_dir, transform=None):
|
|
self.img_dir = img_dir
|
|
self.img_names = [f for f in os.listdir(img_dir) if f.endswith('.png')]
|
|
self.transform = transform
|
|
|
|
def __len__(self):
|
|
return len(self.img_names)
|
|
|
|
def __getitem__(self, idx):
|
|
img_path = os.path.join(self.img_dir, self.img_names[idx])
|
|
image = Image.open(img_path).convert("L")
|
|
if self.transform:
|
|
xi, xj = self.transform(image)
|
|
return xi, xj
|
|
|
|
class MSI_Structural_Encoder(nn.Module):
|
|
def __init__(self, projection_dim=128):
|
|
super(MSI_Structural_Encoder, self).__init__()
|
|
|
|
# High-Fidelity Custom Encoder for Padded Data
|
|
self.features = nn.Sequential(
|
|
# Input: 256x256x1 (Resize in transform)
|
|
nn.Conv2d(1, 32, kernel_size=5, stride=2, padding=2, bias=False),
|
|
nn.GroupNorm(4, 32), # Replaces BatchNorm for perfect stability at batch size = 8
|
|
nn.ReLU(),
|
|
|
|
# 112x112 -> 56x56
|
|
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False),
|
|
nn.GroupNorm(8, 64),
|
|
nn.ReLU(),
|
|
|
|
# 56x56 -> 28x28 (Captures true raw pixel leaf structures)
|
|
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
|
|
nn.GroupNorm(16, 128),
|
|
nn.ReLU(),
|
|
|
|
# 28x28 -> 14x14
|
|
nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=False),
|
|
nn.GroupNorm(32, 256),
|
|
nn.ReLU(),
|
|
|
|
nn.AdaptiveAvgPool2d((1, 1)) # Yields a clean 256-dimensional vector
|
|
)
|
|
|
|
self.projector = nn.Sequential(
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, projection_dim)
|
|
)
|
|
|
|
def forward(self, x):
|
|
h = self.features(x).squeeze(-1).squeeze(-1)
|
|
h_norm = F.normalize(h, p=2, dim=1) # Projects onto the unit hypersphere
|
|
z = self.projector(h_norm)
|
|
return h, z
|
|
|
|
def nt_xent_loss(z1, z2, temperature=0.5):
|
|
batch_size = z1.shape[0]
|
|
z = torch.cat([z1, z2], dim=0)
|
|
z = F.normalize(z, p=2, dim=1)
|
|
|
|
sim_matrix = torch.mm(z, z.t()) / temperature
|
|
mask = torch.eye(2 * batch_size, device=z.device).bool()
|
|
sim_matrix = sim_matrix.masked_fill(mask, -9e15)
|
|
|
|
targets = torch.arange(2 * batch_size, device=z.device)
|
|
targets[:batch_size] += batch_size
|
|
targets[batch_size:] -= batch_size
|
|
|
|
return F.cross_entropy(sim_matrix, targets)
|
|
|
|
if __name__ == "__main__":
|
|
print(f"Launching VRAM-Optimized Contrastive Pipeline on: {DEVICE}")
|
|
print(f"Physical Batch Size: {BATCH_SIZE} | Accumulation Steps: {ACCUMULATION_STEPS} | Effective Batch Size: {BATCH_SIZE * ACCUMULATION_STEPS}")
|
|
os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
|
|
|
|
aug_pipeline = MSIAugmentations(size=256)
|
|
dataset = SimCLRMSIDataset(IMG_DIR, transform=aug_pipeline)
|
|
loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=2, pin_memory=True)
|
|
|
|
model = MSI_Structural_Encoder().to(DEVICE)
|
|
optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
|
|
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)
|
|
|
|
process = psutil.Process(os.getpid())
|
|
|
|
# Environment tweak to combat memory fragmentation
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
|
torch.cuda.empty_cache()
|
|
|
|
for epoch in range(EPOCHS):
|
|
model.train()
|
|
total_loss = 0
|
|
epoch_start = time.time()
|
|
|
|
optimizer.zero_grad() # Initialize gradients outside the accumulation loop
|
|
|
|
for batch_idx, (xi, xj) in enumerate(loader):
|
|
xi, xj = xi.to(DEVICE), xj.to(DEVICE)
|
|
|
|
_, z1 = model(xi)
|
|
_, z2 = model(xj)
|
|
|
|
# Scale loss by accumulation steps to ensure correct gradient weighting
|
|
loss = nt_xent_loss(z1, z2, temperature=TEMPERATURE) / ACCUMULATION_STEPS
|
|
loss.backward()
|
|
|
|
total_loss += loss.item() * ACCUMULATION_STEPS
|
|
|
|
# Step the optimizer only after gathering enough gradients
|
|
if (batch_idx + 1) % ACCUMULATION_STEPS == 0 or (batch_idx + 1) == len(loader):
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
|
|
scheduler.step()
|
|
|
|
# Periodic cache flushing keeps VRAM overhead low
|
|
if epoch % 5 == 0:
|
|
torch.cuda.empty_cache()
|
|
|
|
print(f"Epoch [{epoch+1}/{EPOCHS}], Loss: {total_loss / len(loader):.4f}, Time: {time.time()-epoch_start:.2f}s, RAM Used: {process.memory_info().rss / 1024 / 1024:.1f} MB")
|
|
|
|
print(f"--> Exporting optimized 32-BS encoder weights to: {MODEL_SAVE_PATH}")
|
|
torch.save(model.features.state_dict(), MODEL_SAVE_PATH)
|
|
print(f"Model saved to: {MODEL_SAVE_PATH}")
|