""" cloud_server/train_model_fhe_compatible.py SecureLens — FHE-Compatible Training Script CRITICAL DIFFERENCE FROM train_model.py: - NO ReLU between linear layers (ReLU not FHE-compatible) - BatchNorm folded into weights after training - Architecture matches HE inference EXACTLY This ensures zero accuracy loss between training and FHE inference. """ import os, sys, json import numpy as np import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler from torchvision import transforms, models from PIL import Image from tqdm import tqdm BASE_DIR = os.path.dirname(os.path.abspath(__file__)) MODELS_DIR = os.path.join(BASE_DIR, "models") DATA_DIR = os.path.join(BASE_DIR, "..", "data", "chest_xray") os.makedirs(MODELS_DIR, exist_ok=True) IMAGE_SIZE = 224 BATCH_SIZE = 32 EPOCHS = 20 LR_HEAD = 1e-3 LR_BACKBONE = 1e-5 NUM_CLASSES = 2 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"[Train] Device : {DEVICE}") print(f"[Train] Mode : FHE-COMPATIBLE (no ReLU in head)") # ── Dataset (same as original) ─────────────────────────────────────── class ChestXRayDataset(Dataset): CLASSES = {"NORMAL": 0, "PNEUMONIA": 1} def __init__(self, root_dir, split="train", transform=None): self.transform = transform self.samples = [] split_dir = os.path.join(root_dir, split) if not os.path.exists(split_dir): raise FileNotFoundError(f"Not found: {split_dir}") for cls, label in self.CLASSES.items(): d = os.path.join(split_dir, cls) if not os.path.exists(d): continue for f in os.listdir(d): if f.lower().endswith((".jpeg",".jpg",".png")): self.samples.append((os.path.join(d,f), label)) n = sum(1 for _,l in self.samples if l==0) p = sum(1 for _,l in self.samples if l==1) print(f" [{split:5s}] {len(self.samples):5d} images" f" NORMAL:{n} PNEUMONIA:{p}") def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label = self.samples[idx] try: img = Image.open(path).convert("RGB") except: img = Image.new("RGB",(IMAGE_SIZE,IMAGE_SIZE),128) if self.transform: img = self.transform(img) return img, label def get_transforms(): mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] train_tf = transforms.Compose([ transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(15), transforms.RandomAffine(degrees=0, translate=(0.1,0.1), scale=(0.9,1.1)), transforms.ColorJitter(brightness=0.3, contrast=0.3), transforms.RandomGrayscale(p=0.1), transforms.ToTensor(), transforms.Normalize(mean, std), transforms.RandomErasing(p=0.2), ]) val_tf = transforms.Compose([ transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)), transforms.ToTensor(), transforms.Normalize(mean, std), ]) return train_tf, val_tf # ── FHE-Compatible Model ────────────────────────────────────────────── class SecureLensNetFHE(nn.Module): """ FHE-Compatible Architecture: - ResNet-18 backbone (512-dim features) - Linear head WITHOUT ReLU (FHE cannot compute ReLU efficiently) - BatchNorm for training stability (will be folded into weights) - NO Dropout (not needed for inference) Architecture: 512 → [Linear+BN] → 256 → [Linear] → 2 After training, BatchNorm is folded into the Linear weights. """ def __init__(self, num_classes=2): super().__init__() # Pretrained ResNet-18 backbone backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) self.backbone = nn.Sequential(*list(backbone.children())[:-1]) # FHE-compatible head: Linear → BN → Linear (NO ReLU!) self.head = nn.Sequential( nn.Linear(512, 256), # [0] nn.BatchNorm1d(256), # [1] - for training only, will be folded nn.Linear(256, num_classes), # [2] ) def forward(self, x): x = self.backbone(x) # (B, 512, 1, 1) x = x.view(x.size(0), -1) # (B, 512) x = self.head(x) return x def get_backbone_features(self, x): """Extract 512-dim feature vector from image.""" with torch.no_grad(): f = self.backbone(x) return f.view(f.size(0), -1) def fold_batchnorm_into_linear(self): """ Folds BatchNorm parameters into the preceding Linear layer. After this, BN becomes identity and can be removed. Formula: y = gamma * (x - mean) / sqrt(var + eps) + beta = (gamma / sqrt(var + eps)) * x + (beta - gamma * mean / sqrt(var + eps)) So: W_folded = gamma / sqrt(var + eps) * W b_folded = gamma / sqrt(var + eps) * b + (beta - gamma * mean / sqrt(var + eps)) """ self.eval() # Use running stats linear1 = self.head[0] # First linear bn = self.head[1] # BatchNorm # Get BN parameters gamma = bn.weight.data beta = bn.bias.data mean = bn.running_mean var = bn.running_var eps = bn.eps # Compute scale factor scale = gamma / torch.sqrt(var + eps) # Fold into Linear1 linear1.weight.data = linear1.weight.data * scale.unsqueeze(1) linear1.bias.data = linear1.bias.data * scale + (beta - gamma * mean / torch.sqrt(var + eps)) # Reset BN to identity bn.weight.data.fill_(1.0) bn.bias.data.fill_(0.0) bn.running_mean.fill_(0.0) bn.running_var.fill_(1.0) print("[Fold] BatchNorm folded into Linear layer") def extract_feature_weights(self): """512 → 256 (with BN folded)""" l = self.head[0] return { "W": l.weight.detach().cpu().numpy().tolist(), "b": l.bias.detach().cpu().numpy().tolist(), } def extract_linear_weights(self): """256 → 2""" l = self.head[2] return { "W": l.weight.detach().cpu().numpy().tolist(), "b": l.bias.detach().cpu().numpy().tolist(), } # ── Training helpers (same as original) ─────────────────────────────── def make_sampler(dataset): labels = [s[1] for s in dataset.samples] counts = [labels.count(0), labels.count(1)] weights = [1.0/counts[l] for l in labels] return WeightedRandomSampler(weights, len(weights)) def train_epoch(model, loader, optimizer, criterion): model.train() loss_sum, correct, total = 0.0, 0, 0 for imgs, labels in tqdm(loader, desc=" Train", leave=False): imgs, labels = imgs.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() out = model(imgs) loss = criterion(out, labels) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() loss_sum += loss.item() * imgs.size(0) correct += (out.argmax(1)==labels).sum().item() total += imgs.size(0) return loss_sum/total, correct/total def evaluate(model, loader, criterion): model.eval() loss_sum, correct, total = 0.0, 0, 0 with torch.no_grad(): for imgs, labels in tqdm(loader, desc=" Eval ", leave=False): imgs, labels = imgs.to(DEVICE), labels.to(DEVICE) out = model(imgs) loss = criterion(out, labels) loss_sum += loss.item() * imgs.size(0) correct += (out.argmax(1)==labels).sum().item() total += imgs.size(0) return loss_sum/total, correct/total # ── Main Training Loop ──────────────────────────────────────────────── def main(): print("\n"+"="*60) print(" SecureLens — FHE-Compatible Training (No ReLU)") print("="*60) train_tf, val_tf = get_transforms() print("\n[Datasets]") train_ds = ChestXRayDataset(DATA_DIR, "train", train_tf) val_ds = ChestXRayDataset(DATA_DIR, "val", val_tf) test_ds = ChestXRayDataset(DATA_DIR, "test", val_tf) sampler = make_sampler(train_ds) train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, sampler=sampler, num_workers=0) val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) model = SecureLensNetFHE(NUM_CLASSES).to(DEVICE) print(f"\n[Model] Total params : " f"{sum(p.numel() for p in model.parameters()):,}") print(f"[Model] Architecture : 512 → 256 → 2 (NO ReLU)") criterion = nn.CrossEntropyLoss(label_smoothing=0.1) optimizer = optim.AdamW([ {"params": model.backbone.parameters(), "lr": LR_BACKBONE}, {"params": model.head.parameters(), "lr": LR_HEAD}, ], weight_decay=1e-3) scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=EPOCHS, eta_min=1e-7) best_val_acc = 0.0 patience = 6 no_improve = 0 history = {"train_loss":[],"train_acc":[], "val_loss":[],"val_acc":[]} print("\n[Training]\n") for epoch in range(1, EPOCHS+1): tr_loss, tr_acc = train_epoch(model, train_loader, optimizer, criterion) vl_loss, vl_acc = evaluate(model, val_loader, criterion) scheduler.step() history["train_loss"].append(round(tr_loss,4)) history["train_acc"].append(round(tr_acc,4)) history["val_loss"].append(round(vl_loss,4)) history["val_acc"].append(round(vl_acc,4)) gap = abs(vl_acc - tr_acc) flag = "" if vl_acc > best_val_acc: best_val_acc = vl_acc no_improve = 0 torch.save(model.state_dict(), os.path.join(MODELS_DIR,"best_model_fhe.pth")) flag = " ✅ saved" else: no_improve += 1 print(f" Epoch {epoch:02d}/{EPOCHS} " f"Train:{tr_acc:.2%}({tr_loss:.4f}) " f"Val:{vl_acc:.2%}({vl_loss:.4f}) " f"Gap:{gap:.2%}{flag}") if no_improve >= patience: print(f"\n Early stopping at epoch {epoch}.") break # Test print("\n[Test] Loading best model...") model.load_state_dict( torch.load(os.path.join(MODELS_DIR,"best_model_fhe.pth"), map_location=DEVICE)) ts_loss, ts_acc = evaluate(model, test_loader, criterion) print(f" Test Loss : {ts_loss:.4f}") print(f" Test Accuracy : {ts_acc:.2%}") # CRITICAL: Fold BatchNorm into weights print("\n[Export] Folding BatchNorm into Linear weights...") model.fold_batchnorm_into_linear() # Verify folding didn't break anything print("[Export] Verifying folded model...") ts_loss_fold, ts_acc_fold = evaluate(model, test_loader, criterion) print(f" After folding accuracy : {ts_acc_fold:.2%}") assert abs(ts_acc - ts_acc_fold) < 0.001, "Folding changed accuracy!" # Export weights feat_w = model.extract_feature_weights() linear_w = model.extract_linear_weights() exports = { "feature_weights.json": feat_w, "linear_weights.json": linear_w, } for fname, data in exports.items(): path = os.path.join(MODELS_DIR, fname) with open(path,"w") as f: json.dump(data, f) W = np.array(data["W"]) print(f" {fname:30s} shape: {W.shape}") # Save the folded model torch.save(model.state_dict(), os.path.join(MODELS_DIR,"best_model.pth")) # Overwrite original # Also save as FHE version torch.save(model.state_dict(), os.path.join(MODELS_DIR,"securelens_fhe.pth")) with open(os.path.join(MODELS_DIR,"training_history_fhe.json"),"w") as f: json.dump(history, f, indent=2) # Save metadata metadata = { "architecture": "ResNet18 + Linear (NO ReLU)", "fhe_compatible": True, "batchnorm_folded": True, "test_accuracy": float(ts_acc_fold), "val_accuracy": float(best_val_acc), "relu_used": False, "notes": "This model matches the HE inference architecture exactly." } with open(os.path.join(MODELS_DIR,"model_versions.json"),"w") as f: json.dump(metadata, f, indent=2) print(f"\n Best Val Accuracy : {best_val_acc:.2%}") print(f" Final Test Accuracy : {ts_acc_fold:.2%}") print(f"\n✅ FHE-compatible model trained successfully.") print(f"✅ BatchNorm folded - zero inference accuracy loss.") print(f"✅ Architecture matches HE inference exactly.") if __name__ == "__main__": main()